Skip to content

Commit 6262c66

Browse files
committed
feat(fork): fork chat from any assistant message
1 parent 94f5411 commit 6262c66

4 files changed

Lines changed: 215 additions & 1 deletion

File tree

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
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) {
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 title = `${parent.title ?? 'New task'} | Fork`
78+
const now = new Date()
79+
80+
const [newChat] = await db
81+
.insert(copilotChats)
82+
.values({
83+
id: newId,
84+
userId,
85+
workspaceId: parent.workspaceId,
86+
type: parent.type,
87+
title,
88+
model: parent.model,
89+
messages: forkedMessages,
90+
resources: parentResources,
91+
conversationId: null,
92+
updatedAt: now,
93+
lastSeenAt: now,
94+
})
95+
.returning({ id: copilotChats.id, workspaceId: copilotChats.workspaceId })
96+
97+
if (!newChat) {
98+
return createInternalServerErrorResponse('Failed to create forked chat')
99+
}
100+
101+
// Clone copilot-service conversation state (messages, active_messages, memory files).
102+
// Best-effort: if the copilot service doesn't have a row for the source chat yet, skip.
103+
try {
104+
const copilotHeaders: Record<string, string> = { 'Content-Type': 'application/json' }
105+
if (env.COPILOT_API_KEY) {
106+
copilotHeaders['x-api-key'] = env.COPILOT_API_KEY
107+
}
108+
const copilotRes = await fetchGo(`${SIM_AGENT_API_URL}/api/chats/fork`, {
109+
method: 'POST',
110+
headers: copilotHeaders,
111+
body: JSON.stringify({
112+
sourceChatId: chatId,
113+
newChatId: newId,
114+
upToMessageId,
115+
userId,
116+
}),
117+
spanName: 'sim → go /api/chats/fork',
118+
operation: 'fork_chat',
119+
})
120+
if (!copilotRes.ok) {
121+
const text = await copilotRes.text().catch(() => '')
122+
logger.warn('Copilot fork returned non-OK', { status: copilotRes.status, body: text })
123+
}
124+
} catch (err) {
125+
// The copilot service may not have a row for this chat if no messages
126+
// have been sent yet, or if it's unreachable. Log and continue.
127+
logger.warn('Failed to fork copilot-service conversation, skipping', { err })
128+
}
129+
130+
if (newChat.workspaceId) {
131+
taskPubSub?.publishStatusChanged({
132+
workspaceId: newChat.workspaceId,
133+
chatId: newId,
134+
type: 'created',
135+
})
136+
}
137+
138+
captureServerEvent(
139+
userId,
140+
'task_forked',
141+
{ workspace_id: parent.workspaceId ?? '', source_chat_id: chatId },
142+
{ groups: { workspace: parent.workspaceId ?? '' } }
143+
)
144+
145+
return NextResponse.json({ success: true, id: newId })
146+
} catch (error) {
147+
if (error instanceof z.ZodError) {
148+
return createBadRequestResponse('upToMessageId is required')
149+
}
150+
logger.error('Error forking chat:', error)
151+
return createInternalServerErrorResponse('Failed to fork chat')
152+
}
153+
}
154+
)

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

Lines changed: 36 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,
@@ -16,6 +18,7 @@ import {
1618
Tooltip,
1719
} from '@/components/emcn'
1820
import { useSubmitCopilotFeedback } from '@/hooks/queries/copilot-feedback'
21+
import { useForkTask } from '@/hooks/queries/tasks'
1922

2023
const SPECIAL_TAGS = 'thinking|options|usage_upgrade|credential|mothership-error|file'
2124

@@ -48,21 +51,26 @@ interface MessageActionsProps {
4851
chatId?: string
4952
userQuery?: string
5053
requestId?: string
54+
messageId?: string
5155
}
5256

5357
export const MessageActions = memo(function MessageActions({
5458
content,
5559
chatId,
5660
userQuery,
5761
requestId,
62+
messageId,
5863
}: MessageActionsProps) {
64+
const router = useRouter()
65+
const params = useParams<{ workspaceId: string }>()
5966
const [copied, setCopied] = useState(false)
6067
const [copiedRequestId, setCopiedRequestId] = useState(false)
6168
const [pendingFeedback, setPendingFeedback] = useState<'up' | 'down' | null>(null)
6269
const [feedbackText, setFeedbackText] = useState('')
6370
const resetTimeoutRef = useRef<number | null>(null)
6471
const requestIdTimeoutRef = useRef<number | null>(null)
6572
const submitFeedback = useSubmitCopilotFeedback()
73+
const forkTask = useForkTask()
6674

6775
useEffect(() => {
6876
return () => {
@@ -140,9 +148,20 @@ export const MessageActions = memo(function MessageActions({
140148
}
141149
}
142150

151+
const handleFork = async () => {
152+
if (!chatId || !messageId || forkTask.isPending) return
153+
try {
154+
const result = await forkTask.mutateAsync({ chatId, upToMessageId: messageId })
155+
router.push(`/workspace/${params.workspaceId}/task/${result.id}`)
156+
} catch {
157+
/* fork failed — button resets, user can retry */
158+
}
159+
}
160+
143161
const hasContent = Boolean(content)
144162
const canSubmitFeedback = Boolean(chatId && userQuery)
145-
if (!hasContent && !canSubmitFeedback) return null
163+
const canFork = Boolean(chatId && messageId)
164+
if (!hasContent && !canSubmitFeedback && !canFork) return null
146165

147166
return (
148167
<>
@@ -194,6 +213,22 @@ export const MessageActions = memo(function MessageActions({
194213
</Tooltip.Root>
195214
</>
196215
)}
216+
{canFork && (
217+
<Tooltip.Root>
218+
<Tooltip.Trigger asChild>
219+
<button
220+
type='button'
221+
aria-label='Fork from here'
222+
onClick={handleFork}
223+
disabled={forkTask.isPending}
224+
className={BUTTON_CLASS}
225+
>
226+
<GitBranch className={ICON_CLASS} />
227+
</button>
228+
</Tooltip.Trigger>
229+
<Tooltip.Content side='top'>Fork from here</Tooltip.Content>
230+
</Tooltip.Root>
231+
)}
197232
</div>
198233

199234
<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+
}

0 commit comments

Comments
 (0)