Skip to content

Commit ae75942

Browse files
waleedlatif1waleedlatif1
committed
fix(mothership): point every reference a fork copies at what the fork holds
- A file whose blob copy failed is never published, yet the fork's messages and the worker's history were rewritten to its id and key. References to it now stay on the source file, and its resource tab is dropped as before. - In-app /workspace/<id>/files/<fileId> links were left on the source file; the fork stays in the same workspace, so only the file id moves. - Tool-call arguments and display titles kept the source file's id or key. - A Sources tab for a response past the cut was copied, pointing at a message the fork does not have.
1 parent d504f2e commit ae75942

3 files changed

Lines changed: 199 additions & 19 deletions

File tree

‎apps/sim/app/api/mothership/chats/[chatId]/fork/route.test.ts‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -355,6 +355,103 @@ describe('POST /api/mothership/chats/[chatId]/fork', () => {
355355
expect(mockPublishStatusChanged).not.toHaveBeenCalled()
356356
})
357357

358+
it('keeps references on the source file when its copy fails', async () => {
359+
const oldKey = 'workspace/ws-1/old-cat.png'
360+
mockListForkableChatFiles.mockResolvedValue([
361+
{ id: OLD_FILE_ID, key: oldKey, messageId: 'msg-1', workspaceId: 'ws-1' },
362+
])
363+
mockPlanChatFileCopies.mockReturnValue({
364+
idMap: new Map([[OLD_FILE_ID, NEW_FILE_ID]]),
365+
keyMap: new Map([[oldKey, 'workspace/ws-1/new-cat.png']]),
366+
blobTasks: [
367+
{
368+
copyId: NEW_FILE_ID,
369+
sourceKey: oldKey,
370+
targetKey: 'workspace/ws-1/new-cat.png',
371+
context: 'mothership',
372+
fileName: 'cat.png',
373+
contentType: 'image/png',
374+
},
375+
],
376+
})
377+
mockExecuteChatFileBlobCopies.mockResolvedValue({
378+
copied: 0,
379+
failed: 1,
380+
failedCopyIds: [NEW_FILE_ID],
381+
})
382+
const res = await POST(createRequest('chat-1'), createRouteContext({ chatId: 'chat-1' }))
383+
expect(res.status).toBe(200)
384+
// The copy is never published, so neither Sim's transcript nor the worker's may name it.
385+
expect(mockAppendCopilotChatMessages.mock.calls[0][1][0].content).toBe(
386+
`See ![cat](/api/files/view/${OLD_FILE_ID})`
387+
)
388+
const forkRequest = JSON.parse(mockFetchGo.mock.calls[0][1].body)
389+
expect(forkRequest.fileIds).toEqual({})
390+
expect(forkRequest.fileKeys).toEqual({})
391+
})
392+
393+
it('re-points in-app file links and tool-call arguments at the copied file', async () => {
394+
mockLoadCopilotChatMessages.mockResolvedValue([
395+
{ ...threeMessages[0], content: `Open /workspace/ws-1/files/${OLD_FILE_ID}` },
396+
{
397+
...threeMessages[1],
398+
contentBlocks: [
399+
{
400+
type: 'tool',
401+
toolCall: {
402+
id: 'call-1',
403+
name: 'sim_cli',
404+
state: 'success',
405+
params: { fileId: OLD_FILE_ID, args: ['files', 'read', OLD_FILE_ID], n: 2 },
406+
display: { title: `Read /api/files/view/${OLD_FILE_ID}` },
407+
},
408+
},
409+
],
410+
},
411+
])
412+
mockPlanChatFileCopies.mockReturnValue({
413+
idMap: new Map([[OLD_FILE_ID, NEW_FILE_ID]]),
414+
keyMap: new Map(),
415+
blobTasks: [],
416+
})
417+
const res = await POST(createRequest('chat-1'), createRouteContext({ chatId: 'chat-1' }))
418+
expect(res.status).toBe(200)
419+
const [user, assistant] = mockAppendCopilotChatMessages.mock.calls[0][1]
420+
expect(user.content).toBe(`Open /workspace/ws-1/files/${NEW_FILE_ID}`)
421+
expect(assistant.contentBlocks[0].toolCall).toMatchObject({
422+
params: { fileId: NEW_FILE_ID, args: ['files', 'read', NEW_FILE_ID], n: 2 },
423+
display: { title: `Read /api/files/view/${NEW_FILE_ID}` },
424+
})
425+
})
426+
427+
it('drops a Sources tab whose response is past the cut', async () => {
428+
const kept = { type: 'sources', id: 'cited-sources', title: 'Sources' }
429+
dbChainMockFns.limit.mockResolvedValue([
430+
{ ...parentRow, resources: [{ ...kept, sources: { messageId: 'msg-3' } }] },
431+
])
432+
await POST(createRequest('chat-1'), createRouteContext({ chatId: 'chat-1' }))
433+
expect(dbChainMockFns.values).toHaveBeenCalledWith(expect.objectContaining({ resources: [] }))
434+
435+
dbChainMockFns.values.mockClear()
436+
dbChainMockFns.limit.mockResolvedValue([
437+
{
438+
...parentRow,
439+
resources: [{ ...kept, sources: { messageId: 'live-id', requestId: 'req-2' } }],
440+
},
441+
])
442+
mockLoadCopilotChatMessages.mockResolvedValue([
443+
threeMessages[0],
444+
{ ...threeMessages[1], requestId: 'req-2' },
445+
threeMessages[2],
446+
])
447+
await POST(createRequest('chat-1'), createRouteContext({ chatId: 'chat-1' }))
448+
expect(dbChainMockFns.values).toHaveBeenCalledWith(
449+
expect.objectContaining({
450+
resources: [{ ...kept, sources: { messageId: 'live-id', requestId: 'req-2' } }],
451+
})
452+
)
453+
})
454+
358455
it('copies pre-cut uploads and drops only post-cut ghosts', async () => {
359456
// The source chat owns two more uploads (apple pre-cut, banana post-cut)
360457
// beside the kept one, plus one shared workspace-file resource. The fork

‎apps/sim/lib/mothership/chat/application/fork.ts‎

Lines changed: 20 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import {
2626
import { loadCopilotChatMessages } from '@/lib/mothership/chat/lifecycle'
2727
import { appendCopilotChatMessages } from '@/lib/mothership/chat/messages-store'
2828
import {
29+
publishedFileRefMaps,
2930
rewriteMessageFileRefs,
3031
rewriteResourceFileRefs,
3132
} from '@/lib/mothership/chat/rewrite-file-references'
@@ -100,18 +101,24 @@ export const forkChat = defineAuthorizedChatUseCase({
100101
throw new OrchestrationError('validation', 'Message not found in chat')
101102
}
102103
const forkedMessages = messages.slice(0, forkIdx + 1)
104+
const keptMessageIds = new Set(forkedMessages.map((m) => m.id))
105+
const keptRequestIds = new Set(
106+
forkedMessages.flatMap((m) => (m.requestId ? [m.requestId] : []))
107+
)
108+
/** The Sources panel reads its response by message or request id; a response past the cut is not in the fork. */
109+
const addressesKeptMessage = ({ sources }: MothershipResource) =>
110+
!!sources &&
111+
(keptMessageIds.has(sources.messageId) ||
112+
(!!sources.requestId && keptRequestIds.has(sources.requestId)))
103113

104114
/** Single workspace_files read per fork: every chat-owned upload. The copied set is timeline-cut to the kept message slice in memory (files born after the fork point stay behind). */
105115
const chatOwnedFiles = context.workspaceId ? await listForkableChatFiles(db, chatId) : []
106-
const sourceFiles = filterForkableChatFiles(
107-
chatOwnedFiles,
108-
new Set(forkedMessages.map((m) => m.id))
109-
)
116+
const sourceFiles = filterForkableChatFiles(chatOwnedFiles, keptMessageIds)
110117

111118
/** Resources are stored as a jsonb array on the chat row. They carry no timestamps, so they can't be timeline-cut like messages — instead, file resources whose chat-owned file is NOT copied (uploads born after the cut) are dropped in the rewrite below; everything else is copied. */
112119
const parentResources = sanitizeChatResources(
113120
Array.isArray(parent.resources) ? (parent.resources as MothershipResource[]) : []
114-
)
121+
).filter((resource) => resource.type !== 'sources' || addressesKeptMessage(resource))
115122

116123
/** The source chat's chat-owned file ids (no cut) — the "is this resource a ghost?" test set for the rewrite. */
117124
const chatOwnedFileIds = new Set(chatOwnedFiles.map((row) => row.id))
@@ -125,12 +132,12 @@ export const forkChat = defineAuthorizedChatUseCase({
125132
preparedBlobs = [...plan.blobTasks, ...planForkInlineImages(forkedMessages, chatId, newId)]
126133
const { failed, failedCopyIds } = await executeChatFileBlobCopies(preparedBlobs)
127134
const failedIds = new Set(failedCopyIds)
128-
const maps = { fileIds: plan.idMap, fileKeys: plan.keyMap }
129-
const newChatResources = rewriteResourceFileRefs(
130-
parentResources,
131-
maps,
132-
chatOwnedFileIds
133-
).filter((resource) => resource.type !== 'file' || !failedIds.has(resource.id))
135+
const maps = {
136+
...publishedFileRefMaps(plan, failedIds),
137+
workspaceId: context.workspaceId,
138+
}
139+
/** A chat-owned file whose copy failed is a ghost here too: no published copy stands in for it. */
140+
const newChatResources = rewriteResourceFileRefs(parentResources, maps, chatOwnedFileIds)
134141
const cutUser = [...forkedMessages].reverse().find((message) => message.role === 'user')
135142
if (!cutUser) throw new Error('The fork has no user message')
136143
workerCopyRequested = true
@@ -143,8 +150,8 @@ export const forkChat = defineAuthorizedChatUseCase({
143150
userId,
144151
upToMessageId: cutUser.id,
145152
includeResponse: forkedMessages.at(-1)?.role === 'assistant',
146-
fileIds: Object.fromEntries(plan.idMap),
147-
fileKeys: Object.fromEntries(plan.keyMap),
153+
fileIds: Object.fromEntries(maps.fileIds),
154+
fileKeys: Object.fromEntries(maps.fileKeys),
148155
})
149156

150157
/** Publish only after both the file bytes and the worker conversation are prepared. */

‎apps/sim/lib/mothership/chat/rewrite-file-references.ts‎

Lines changed: 82 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import { isPlainRecord } from '@sim/utils/object'
2+
import type { PersistedContentBlock } from '@/lib/api/contracts/copilot-messages'
13
import type { PersistedMessage } from '@/lib/mothership/chat/persisted-message'
24
import type { MothershipResource } from '@/lib/mothership/resources/types'
35
import { rewriteForkContentRefs } from '@/ee/workspace-forking/lib/remap/remap-content-refs'
@@ -10,22 +12,98 @@ import { rewriteForkContentRefs } from '@/ee/workspace-forking/lib/remap/remap-c
1012
export interface ChatFileRefMaps {
1113
fileIds: ReadonlyMap<string, string>
1214
fileKeys: ReadonlyMap<string, string>
15+
/**
16+
* The workspace both chats live in. A fork stays in its source's workspace, so in-app
17+
* `/workspace/<id>/files/<fileId>` links keep their workspace and only the file id moves.
18+
*/
19+
workspaceId?: string | null
20+
}
21+
22+
/**
23+
* The copy plan's id/key maps restricted to copies whose bytes were prepared. A failed copy
24+
* is never published, so the fork's messages, resources and worker history keep naming the
25+
* source file (alive while the source chat is) instead of an id or key that never exists.
26+
*/
27+
export function publishedFileRefMaps(
28+
plan: {
29+
idMap: ReadonlyMap<string, string>
30+
keyMap: ReadonlyMap<string, string>
31+
blobTasks: readonly { copyId: string; targetKey: string }[]
32+
},
33+
failedCopyIds: ReadonlySet<string>
34+
): { fileIds: Map<string, string>; fileKeys: Map<string, string> } {
35+
const failedKeys = new Set(
36+
plan.blobTasks.filter((task) => failedCopyIds.has(task.copyId)).map((task) => task.targetKey)
37+
)
38+
return {
39+
fileIds: new Map([...plan.idMap].filter(([, copyId]) => !failedCopyIds.has(copyId))),
40+
fileKeys: new Map([...plan.keyMap].filter(([, copyKey]) => !failedKeys.has(copyKey))),
41+
}
1342
}
1443

1544
function hasMappings(maps: ChatFileRefMaps): boolean {
1645
return maps.fileIds.size > 0 || maps.fileKeys.size > 0
1746
}
1847

1948
function rewriteText(text: string, maps: ChatFileRefMaps): string {
20-
return rewriteForkContentRefs(text, { fileIds: maps.fileIds, fileKeys: maps.fileKeys })
49+
return rewriteForkContentRefs(text, {
50+
fileIds: maps.fileIds,
51+
fileKeys: maps.fileKeys,
52+
...(maps.workspaceId ? { workspaceId: { from: maps.workspaceId, to: maps.workspaceId } } : {}),
53+
})
54+
}
55+
56+
/**
57+
* Tool arguments name a file by its bare id or storage key (`{ fileId }`, `["files", "read", id]`),
58+
* so a string that IS a mapped id or key is replaced whole; any other string gets the URL grammar.
59+
*/
60+
function rewriteToolValue(value: unknown, maps: ChatFileRefMaps): unknown {
61+
if (typeof value === 'string')
62+
return maps.fileIds.get(value) ?? maps.fileKeys.get(value) ?? rewriteText(value, maps)
63+
if (Array.isArray(value)) return value.map((entry) => rewriteToolValue(entry, maps))
64+
if (isPlainRecord(value))
65+
return Object.fromEntries(
66+
Object.entries(value).map(([key, entry]) => [key, rewriteToolValue(entry, maps)])
67+
)
68+
return value
69+
}
70+
71+
function rewriteBlock(block: PersistedContentBlock, maps: ChatFileRefMaps): PersistedContentBlock {
72+
const toolCall = block.toolCall
73+
return {
74+
...block,
75+
...(block.content ? { content: rewriteText(block.content, maps) } : {}),
76+
...(toolCall
77+
? {
78+
toolCall: {
79+
...toolCall,
80+
...(toolCall.params
81+
? { params: rewriteToolValue(toolCall.params, maps) as Record<string, unknown> }
82+
: {}),
83+
...(toolCall.activityDescription
84+
? { activityDescription: rewriteText(toolCall.activityDescription, maps) }
85+
: {}),
86+
...(toolCall.display?.title
87+
? {
88+
display: {
89+
...toolCall.display,
90+
title: rewriteText(toolCall.display.title, maps),
91+
},
92+
}
93+
: {}),
94+
},
95+
}
96+
: {}),
97+
}
2198
}
2299

23100
/**
24101
* Re-point every file reference in a copied transcript at the copied files, so
25102
* the fork is self-contained (it survives the original chat's deletion).
26103
* Rewrites: free-text URLs in `content` and text content blocks (serve/view/
27-
* in-app/`sim:file` forms, via the shared fork grammar), attachment chip
28-
* ids+keys, and `@`-mention context chip file ids. References to anything not
104+
* in-app/`sim:file` forms, via the shared fork grammar), tool-call arguments
105+
* and display text, attachment chip ids+keys, and `@`-mention context chip
106+
* file ids. References to anything not
29107
* in the maps (shared workspace files, workflows, other chats) pass through
30108
* unchanged. Pure; returns the input array untouched when there is nothing to
31109
* rewrite.
@@ -50,9 +128,7 @@ export function rewriteMessageFileRefs(
50128
content: rewriteText(message.content, maps),
51129
}
52130
if (message.contentBlocks?.length) {
53-
rewritten.contentBlocks = message.contentBlocks.map((block) =>
54-
block.content ? { ...block, content: rewriteText(block.content, maps) } : block
55-
)
131+
rewritten.contentBlocks = message.contentBlocks.map((block) => rewriteBlock(block, maps))
56132
}
57133
if (message.fileAttachments?.length) {
58134
rewritten.fileAttachments = message.fileAttachments.map((att) => ({

0 commit comments

Comments
 (0)