Skip to content

Commit 1a321c5

Browse files
authored
feat(fork): fork chat from any assistant message (#4343)
* feat(fork): fork chat from any assistant message * fix(fork): toast on failure, disabled opacity, copy previewYaml/planArtifact/config * fix(fork): type guard for mothership-only, prevent title accumulation on re-fork * fix(fork): register task_forked PostHog event type
1 parent 7d8ec24 commit 1a321c5

5 files changed

Lines changed: 226 additions & 1 deletion

File tree

Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,158 @@
1+
import { db } from '@sim/db'
2+
import { copilotChats } from '@sim/db/schema'
3+
import { createLogger } from '@sim/logger'
4+
import { generateId } from '@sim/utils/id'
5+
import { eq } from 'drizzle-orm'
6+
import { type NextRequest, NextResponse } from 'next/server'
7+
import { z } from 'zod'
8+
import type { PersistedMessage } from '@/lib/copilot/chat/persisted-message'
9+
import { SIM_AGENT_API_URL } from '@/lib/copilot/constants'
10+
import { fetchGo } from '@/lib/copilot/request/go/fetch'
11+
import {
12+
authenticateCopilotRequestSessionOnly,
13+
createBadRequestResponse,
14+
createInternalServerErrorResponse,
15+
createNotFoundResponse,
16+
createUnauthorizedResponse,
17+
} from '@/lib/copilot/request/http'
18+
import type { MothershipResource } from '@/lib/copilot/resources/types'
19+
import { taskPubSub } from '@/lib/copilot/tasks'
20+
import { env } from '@/lib/core/config/env'
21+
import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
22+
import { captureServerEvent } from '@/lib/posthog/server'
23+
import { assertActiveWorkspaceAccess } from '@/lib/workspaces/permissions/utils'
24+
25+
const logger = createLogger('ForkChatAPI')
26+
27+
const ForkChatSchema = z.object({
28+
upToMessageId: z.string().min(1),
29+
})
30+
31+
/**
32+
* POST /api/mothership/chats/[chatId]/fork
33+
* Creates a new chat branched from the given chat, keeping messages up to and
34+
* including the specified message. Resources and copilot-side state are copied.
35+
*/
36+
export const POST = withRouteHandler(
37+
async (request: NextRequest, { params }: { params: Promise<{ chatId: string }> }) => {
38+
try {
39+
const { userId, isAuthenticated } = await authenticateCopilotRequestSessionOnly()
40+
if (!isAuthenticated || !userId) {
41+
return createUnauthorizedResponse()
42+
}
43+
44+
const { chatId } = await params
45+
const body = await request.json()
46+
const { upToMessageId } = ForkChatSchema.parse(body)
47+
48+
// Load parent chat and verify ownership.
49+
const [parent] = await db
50+
.select()
51+
.from(copilotChats)
52+
.where(eq(copilotChats.id, chatId))
53+
.limit(1)
54+
55+
if (!parent || parent.userId !== userId || parent.type !== 'mothership') {
56+
return createNotFoundResponse('Chat not found')
57+
}
58+
59+
if (parent.workspaceId) {
60+
await assertActiveWorkspaceAccess(parent.workspaceId, userId)
61+
}
62+
63+
// Find the fork point in the Sim-side messages array.
64+
const messages = Array.isArray(parent.messages) ? (parent.messages as PersistedMessage[]) : []
65+
const forkIdx = messages.findIndex((m) => m.id === upToMessageId)
66+
if (forkIdx < 0) {
67+
return createBadRequestResponse('Message not found in chat')
68+
}
69+
const forkedMessages = messages.slice(0, forkIdx + 1)
70+
71+
// Resources are stored as a jsonb array on the chat row — copy them directly.
72+
const parentResources = Array.isArray(parent.resources)
73+
? (parent.resources as MothershipResource[])
74+
: []
75+
76+
const newId = generateId()
77+
const baseTitle = (parent.title ?? 'New task').replace(/ \| Fork$/, '')
78+
const title = `${baseTitle} | Fork`
79+
const now = new Date()
80+
81+
const [newChat] = await db
82+
.insert(copilotChats)
83+
.values({
84+
id: newId,
85+
userId,
86+
workspaceId: parent.workspaceId,
87+
type: parent.type,
88+
title,
89+
model: parent.model,
90+
messages: forkedMessages,
91+
resources: parentResources,
92+
previewYaml: parent.previewYaml,
93+
planArtifact: parent.planArtifact,
94+
config: parent.config,
95+
conversationId: null,
96+
updatedAt: now,
97+
lastSeenAt: now,
98+
})
99+
.returning({ id: copilotChats.id, workspaceId: copilotChats.workspaceId })
100+
101+
if (!newChat) {
102+
return createInternalServerErrorResponse('Failed to create forked chat')
103+
}
104+
105+
// Clone copilot-service conversation state (messages, active_messages, memory files).
106+
// Best-effort: if the copilot service doesn't have a row for the source chat yet, skip.
107+
try {
108+
const copilotHeaders: Record<string, string> = { 'Content-Type': 'application/json' }
109+
if (env.COPILOT_API_KEY) {
110+
copilotHeaders['x-api-key'] = env.COPILOT_API_KEY
111+
}
112+
const copilotRes = await fetchGo(`${SIM_AGENT_API_URL}/api/chats/fork`, {
113+
method: 'POST',
114+
headers: copilotHeaders,
115+
body: JSON.stringify({
116+
sourceChatId: chatId,
117+
newChatId: newId,
118+
upToMessageId,
119+
userId,
120+
}),
121+
spanName: 'sim → go /api/chats/fork',
122+
operation: 'fork_chat',
123+
})
124+
if (!copilotRes.ok) {
125+
const text = await copilotRes.text().catch(() => '')
126+
logger.warn('Copilot fork returned non-OK', { status: copilotRes.status, body: text })
127+
}
128+
} catch (err) {
129+
// The copilot service may not have a row for this chat if no messages
130+
// have been sent yet, or if it's unreachable. Log and continue.
131+
logger.warn('Failed to fork copilot-service conversation, skipping', { err })
132+
}
133+
134+
if (newChat.workspaceId) {
135+
taskPubSub?.publishStatusChanged({
136+
workspaceId: newChat.workspaceId,
137+
chatId: newId,
138+
type: 'created',
139+
})
140+
}
141+
142+
captureServerEvent(
143+
userId,
144+
'task_forked',
145+
{ workspace_id: parent.workspaceId ?? '', source_chat_id: chatId },
146+
{ groups: { workspace: parent.workspaceId ?? '' } }
147+
)
148+
149+
return NextResponse.json({ success: true, id: newId })
150+
} catch (error) {
151+
if (error instanceof z.ZodError) {
152+
return createBadRequestResponse('upToMessageId is required')
153+
}
154+
logger.error('Error forking chat:', error)
155+
return createInternalServerErrorResponse('Failed to fork chat')
156+
}
157+
}
158+
)

apps/sim/app/workspace/[workspaceId]/components/message-actions/message-actions.tsx

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
'use client'
22

33
import { memo, useEffect, useRef, useState } from 'react'
4+
import { GitBranch } from 'lucide-react'
5+
import { useParams, useRouter } from 'next/navigation'
46
import {
57
Button,
68
Check,
@@ -14,8 +16,11 @@ import {
1416
ThumbsDown,
1517
ThumbsUp,
1618
Tooltip,
19+
toast,
1720
} from '@/components/emcn'
21+
import { cn } from '@/lib/core/utils/cn'
1822
import { useSubmitCopilotFeedback } from '@/hooks/queries/copilot-feedback'
23+
import { useForkTask } from '@/hooks/queries/tasks'
1924

2025
const SPECIAL_TAGS = 'thinking|options|usage_upgrade|credential|mothership-error|file'
2126

@@ -48,21 +53,26 @@ interface MessageActionsProps {
4853
chatId?: string
4954
userQuery?: string
5055
requestId?: string
56+
messageId?: string
5157
}
5258

5359
export const MessageActions = memo(function MessageActions({
5460
content,
5561
chatId,
5662
userQuery,
5763
requestId,
64+
messageId,
5865
}: MessageActionsProps) {
66+
const router = useRouter()
67+
const params = useParams<{ workspaceId: string }>()
5968
const [copied, setCopied] = useState(false)
6069
const [copiedRequestId, setCopiedRequestId] = useState(false)
6170
const [pendingFeedback, setPendingFeedback] = useState<'up' | 'down' | null>(null)
6271
const [feedbackText, setFeedbackText] = useState('')
6372
const resetTimeoutRef = useRef<number | null>(null)
6473
const requestIdTimeoutRef = useRef<number | null>(null)
6574
const submitFeedback = useSubmitCopilotFeedback()
75+
const forkTask = useForkTask()
6676

6777
useEffect(() => {
6878
return () => {
@@ -140,9 +150,20 @@ export const MessageActions = memo(function MessageActions({
140150
}
141151
}
142152

153+
const handleFork = async () => {
154+
if (!chatId || !messageId || forkTask.isPending) return
155+
try {
156+
const result = await forkTask.mutateAsync({ chatId, upToMessageId: messageId })
157+
router.push(`/workspace/${params.workspaceId}/task/${result.id}`)
158+
} catch {
159+
toast.error('Failed to fork chat')
160+
}
161+
}
162+
143163
const hasContent = Boolean(content)
144164
const canSubmitFeedback = Boolean(chatId && userQuery)
145-
if (!hasContent && !canSubmitFeedback) return null
165+
const canFork = Boolean(chatId && messageId)
166+
if (!hasContent && !canSubmitFeedback && !canFork) return null
146167

147168
return (
148169
<>
@@ -194,6 +215,22 @@ export const MessageActions = memo(function MessageActions({
194215
</Tooltip.Root>
195216
</>
196217
)}
218+
{canFork && (
219+
<Tooltip.Root>
220+
<Tooltip.Trigger asChild>
221+
<button
222+
type='button'
223+
aria-label='Fork from here'
224+
onClick={handleFork}
225+
disabled={forkTask.isPending}
226+
className={cn(BUTTON_CLASS, forkTask.isPending && 'cursor-not-allowed opacity-50')}
227+
>
228+
<GitBranch className={ICON_CLASS} />
229+
</button>
230+
</Tooltip.Trigger>
231+
<Tooltip.Content side='top'>Fork from here</Tooltip.Content>
232+
</Tooltip.Root>
233+
)}
197234
</div>
198235

199236
<Modal open={pendingFeedback !== null} onOpenChange={handleModalClose}>

apps/sim/app/workspace/[workspaceId]/home/components/mothership-chat/mothership-chat.tsx

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,7 @@ export function MothershipChat({
217217
chatId={chatId}
218218
userQuery={precedingUserContent}
219219
requestId={msg.requestId}
220+
messageId={msg.id}
220221
/>
221222
</div>
222223
)}

apps/sim/hooks/queries/tasks.ts

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -590,3 +590,27 @@ export function useMarkTaskUnread(workspaceId?: string) {
590590
},
591591
})
592592
}
593+
594+
async function forkChat(params: {
595+
chatId: string
596+
upToMessageId: string
597+
}): Promise<{ id: string }> {
598+
const response = await fetch(`/api/mothership/chats/${params.chatId}/fork`, {
599+
method: 'POST',
600+
headers: { 'Content-Type': 'application/json' },
601+
body: JSON.stringify({ upToMessageId: params.upToMessageId }),
602+
})
603+
if (!response.ok) throw new Error('Failed to fork chat')
604+
const data = await response.json()
605+
return { id: data.id }
606+
}
607+
608+
export function useForkTask() {
609+
const queryClient = useQueryClient()
610+
return useMutation({
611+
mutationFn: forkChat,
612+
onSettled: () => {
613+
queryClient.invalidateQueries({ queryKey: taskKeys.lists() })
614+
},
615+
})
616+
}

apps/sim/lib/posthog/events.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -406,6 +406,11 @@ export interface PostHogEventMap {
406406
workspace_id: string
407407
}
408408

409+
task_forked: {
410+
workspace_id: string
411+
source_chat_id: string
412+
}
413+
409414
task_marked_unread: {
410415
workspace_id: string
411416
}

0 commit comments

Comments
 (0)