Skip to content

Commit 12d12f5

Browse files
waleedlatif1waleedlatif1
committed
fix(cleanup): purge a chat's inline images and only its own attachments
- Chat images are stored under the chat's id but were never deleted with the chat. Their keys are rebuilt from the assistant messages that published them, the same way a fork finds them to copy. - Message attachments were deleted as copilot storage whatever their key. Attachments in Chat are workspace-bucket files owned by their workspace_files row, which a fork can share with its source (a failed copy, a deleted file), and the copilot bucket falls back to the workspace bucket on GCS and can be configured to it on S3. Only copilot keys are deleted from messages now.
1 parent ae75942 commit 12d12f5

7 files changed

Lines changed: 201 additions & 57 deletions

File tree

‎apps/sim/lib/cleanup/chat-cleanup.test.ts‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
import { copilotChats, copilotMessages, workspaceFiles } from '@sim/db/schema'
2+
import { queueTableRows, resetDbChainMock } from '@sim/testing/mocks/database.mock'
3+
import { storageServiceMockFns } from '@sim/testing/mocks/storage-service.mock'
4+
import { uploadsMock, uploadsMockFns } from '@sim/testing/mocks/uploads.mock'
5+
import { beforeEach, describe, expect, it, vi } from 'vitest'
6+
7+
vi.mock('@/lib/uploads', () => uploadsMock)
8+
9+
import { prepareChatCleanup } from '@/lib/cleanup/chat-cleanup'
10+
import { inlineChatImageKey } from '@/lib/mothership/chat/inline-image-key'
11+
12+
const chatId = '3f0c2a52-8a43-4d4b-9b5f-0b3c7a1e2d10'
13+
14+
function attachment(key: string) {
15+
return { id: 'wf_x', key, filename: 'x.png', media_type: 'image/png', size: 1 }
16+
}
17+
18+
/** Purges one deleted chat whose message rows are `messages`; returns the deleted keys by context. */
19+
async function purge(messages: Record<string, unknown>[]) {
20+
queueTableRows(workspaceFiles, [])
21+
queueTableRows(
22+
copilotMessages,
23+
messages.map((content) => ({ chatId, content }))
24+
)
25+
const cleanup = await prepareChatCleanup([chatId], 'test')
26+
queueTableRows(copilotChats, [])
27+
await cleanup.execute()
28+
return Object.fromEntries(
29+
storageServiceMockFns.mockDeleteFiles.mock.calls.map(([keys, context]) => [context, keys])
30+
)
31+
}
32+
33+
describe('chat purge storage', () => {
34+
beforeEach(() => {
35+
resetDbChainMock()
36+
uploadsMockFns.mockIsUsingCloudStorage.mockReturnValue(true)
37+
storageServiceMockFns.mockDeleteFiles.mockReset()
38+
storageServiceMockFns.mockDeleteFiles.mockResolvedValue({ deleted: 1, failed: [] })
39+
})
40+
41+
it('deletes only copilot-storage attachment keys as copilot storage', async () => {
42+
// A fork carries its source's key when the file's copy failed or the file was deleted, and
43+
// the copilot bucket can be the workspace bucket, so a workspace key here is another chat's file.
44+
const deleted = await purge([
45+
{
46+
role: 'user',
47+
content: 'look',
48+
fileAttachments: [
49+
attachment('workspace/ws-1/1700-abc-shared.png'),
50+
attachment('assistant/org-1/user-1/u1/shared.png'),
51+
attachment('copilot/1234/legacy.png'),
52+
],
53+
},
54+
])
55+
expect(deleted).toEqual({ copilot: ['copilot/1234/legacy.png'] })
56+
})
57+
58+
it('deletes the inline images an assistant message published', async () => {
59+
const deleted = await purge([
60+
{
61+
role: 'assistant',
62+
requestId: 'req-1',
63+
content: 'Here: ![chart](files/chart.png) and `![code](files/no.png)`',
64+
contentBlocks: [{ type: 'text', content: '![other](/tmp/out.png)' }],
65+
},
66+
{ role: 'user', requestId: 'req-2', content: '![u](files/u.png)' },
67+
{ role: 'assistant', content: '![x](files/x.png)' },
68+
])
69+
expect(deleted).toEqual({
70+
mothership: [
71+
inlineChatImageKey(chatId, 'req-1', 'files/chart.png'),
72+
inlineChatImageKey(chatId, 'req-1', '/tmp/out.png'),
73+
],
74+
})
75+
})
76+
})

‎apps/sim/lib/cleanup/chat-cleanup.ts‎

Lines changed: 53 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,17 @@ import { dbFor } from '@sim/db'
22
import { copilotChats, copilotMessages, workspaceFiles } from '@sim/db/schema'
33
import { createLogger } from '@sim/logger'
44
import { chunkArray } from '@sim/utils/helpers'
5+
import { isRecordLike } from '@sim/utils/object'
56
import { and, inArray, isNull } from 'drizzle-orm'
67
import { env } from '@/lib/core/config/env'
8+
import {
9+
inlineChatImageKey,
10+
inlineChatImageReferences,
11+
} from '@/lib/mothership/chat/inline-image-key'
712
import { SIM_AGENT_API_URL } from '@/lib/mothership/constants'
813
import type { StorageContext } from '@/lib/uploads'
914
import { isUsingCloudStorage, StorageService } from '@/lib/uploads'
15+
import { tryInferContextFromKey } from '@/lib/uploads/utils/file-utils'
1016

1117
const logger = createLogger('ChatCleanup')
1218

@@ -33,9 +39,47 @@ interface FileRef {
3339
}
3440

3541
/**
36-
* Collect all file storage keys for the given chat IDs from two sources:
42+
* The chat images an assistant message published under its request id, keyed by chat id,
43+
* so they are purged with the chat. A row this cannot read as such a message has none.
44+
*/
45+
function inlineChatImageKeys(chatId: string, content: Record<string, unknown>): string[] {
46+
if (content.role !== 'assistant' || typeof content.requestId !== 'string') return []
47+
const published = inlineChatImageReferences({
48+
role: 'assistant',
49+
requestId: content.requestId,
50+
content: typeof content.content === 'string' ? content.content : '',
51+
contentBlocks: Array.isArray(content.contentBlocks)
52+
? content.contentBlocks.flatMap((block) =>
53+
isRecordLike(block) && block.type === 'text' && typeof block.content === 'string'
54+
? [{ type: 'text' as const, content: block.content }]
55+
: []
56+
)
57+
: undefined,
58+
})
59+
if (!published) return []
60+
const keys: string[] = []
61+
for (const reference of published.references) {
62+
try {
63+
keys.push(inlineChatImageKey(chatId, published.requestId, reference))
64+
} catch {
65+
// A chat or request id outside the key grammar never had an image stored under it.
66+
}
67+
}
68+
return keys
69+
}
70+
71+
/**
72+
* Collect all file storage keys for the given chat IDs from three sources:
3773
* 1. workspaceFiles rows with chatId FK (chat-scoped contexts only)
38-
* 2. fileAttachments[].key inside each copilot_messages.content
74+
* 2. fileAttachments[].key inside each copilot_messages.content, for copilot-storage keys only
75+
* 3. the chat-scoped inline images each assistant message published
76+
*
77+
* An attachment whose key belongs to another storage context is owned by its
78+
* `workspace_files` row (source 1, or the workspace file lifecycle), never by the message:
79+
* a fork carries the same attachment when its copy failed or its file was deleted, and the
80+
* copilot bucket falls back to the workspace bucket on GCS (and may be configured to it on
81+
* S3), so deleting such a key as copilot storage could delete a file another chat or the
82+
* workspace still uses.
3983
*/
4084
export async function collectChatFiles(chatIds: string[]): Promise<FileRef[]> {
4185
const files: FileRef[] = []
@@ -77,6 +121,12 @@ export async function collectChatFiles(chatIds: string[]): Promise<FileRef[]> {
77121
for (const row of messageRows) {
78122
const msg = row.content
79123
if (!msg || typeof msg !== 'object') continue
124+
for (const key of inlineChatImageKeys(row.chatId, msg as Record<string, unknown>)) {
125+
if (!seen.has(key)) {
126+
seen.add(key)
127+
files.push({ key, context: 'mothership', chatId: row.chatId })
128+
}
129+
}
80130
const attachments = (msg as Record<string, unknown>).fileAttachments
81131
if (!Array.isArray(attachments)) continue
82132
for (const attachment of attachments) {
@@ -86,6 +136,7 @@ export async function collectChatFiles(chatIds: string[]): Promise<FileRef[]> {
86136
(attachment as Record<string, unknown>).key
87137
) {
88138
const key = (attachment as Record<string, unknown>).key as string
139+
if (tryInferContextFromKey(key) !== 'copilot') continue
89140
if (!seen.has(key)) {
90141
seen.add(key)
91142
files.push({ key, context: 'copilot', chatId: row.chatId })

‎apps/sim/lib/mothership/chat/application/inline-images.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,8 +46,8 @@ import {
4646
materializeStreamImage,
4747
readInlineChatImage,
4848
} from '@/lib/mothership/chat/application/inline-images'
49+
import { inlineChatImageKey } from '@/lib/mothership/chat/inline-image-key'
4950
import { inlineChatImageUrl } from '@/lib/mothership/chat/inline-image-reference'
50-
import { inlineChatImageKey } from '@/lib/mothership/chat/inline-image-storage'
5151
import { GET } from '@/app/api/mothership/chats/[chatId]/images/[requestId]/route'
5252

5353
const mocks = {

‎apps/sim/lib/mothership/chat/fork-inline-images.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import { describe, expect, it, vi } from 'vitest'
44
vi.mock('@/lib/uploads/core/storage-service', () => storageServiceMock)
55

66
import { planForkInlineImages } from '@/lib/mothership/chat/fork-inline-images'
7-
import { inlineChatImageKey } from '@/lib/mothership/chat/inline-image-storage'
7+
import { inlineChatImageKey } from '@/lib/mothership/chat/inline-image-key'
88
import type { PersistedMessage } from '@/lib/mothership/chat/persisted-message'
99

1010
const message = (content: string): PersistedMessage => ({

‎apps/sim/lib/mothership/chat/fork-inline-images.ts‎

Lines changed: 22 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,10 @@
11
import { OrchestrationError } from '@/lib/core/orchestration/types'
22
import type { ChatBlobCopyTask } from '@/lib/mothership/chat/fork-chat-files'
33
import {
4-
inlineImageRequestIdSchema,
5-
inlineImageSourceSchema,
6-
} from '@/lib/mothership/chat/inline-image-reference'
7-
import {
8-
INLINE_CHAT_IMAGE_MAX_BYTES,
94
inlineChatImageKey,
10-
} from '@/lib/mothership/chat/inline-image-storage'
11-
import { collectMarkdownImageSources } from '@/lib/mothership/chat/markdown-images'
5+
inlineChatImageReferences,
6+
} from '@/lib/mothership/chat/inline-image-key'
7+
import { INLINE_CHAT_IMAGE_MAX_BYTES } from '@/lib/mothership/chat/inline-image-storage'
128
import type { PersistedMessage } from '@/lib/mothership/chat/persisted-message'
139

1410
/** Only retained assistant image references are copied; their Markdown and request IDs stay unchanged. */
@@ -19,38 +15,25 @@ export function planForkInlineImages(
1915
): ChatBlobCopyTask[] {
2016
const tasks = new Map<string, ChatBlobCopyTask>()
2117
for (const message of messages) {
22-
if (
23-
message.role !== 'assistant' ||
24-
!message.requestId ||
25-
!inlineImageRequestIdSchema.safeParse(message.requestId).success
26-
)
27-
continue
28-
const contents = [
29-
message.content,
30-
...(message.contentBlocks ?? [])
31-
.filter((block) => block.type === 'text')
32-
.map((block) => block.content ?? ''),
33-
]
34-
for (const content of contents) {
35-
for (const reference of collectMarkdownImageSources(content)) {
36-
if (!inlineImageSourceSchema.safeParse(reference).success) continue
37-
const sourceKey = inlineChatImageKey(sourceChatId, message.requestId, reference)
38-
tasks.set(sourceKey, {
39-
copyId: sourceKey,
40-
sourceKey,
41-
targetKey: inlineChatImageKey(newChatId, message.requestId, reference),
42-
context: 'mothership',
43-
fileName: 'chat-image.webp',
44-
contentType: 'image/webp',
45-
maxBytes: INLINE_CHAT_IMAGE_MAX_BYTES,
46-
persistMetadata: false,
47-
})
48-
if (tasks.size > 200)
49-
throw new OrchestrationError(
50-
'payload_too_large',
51-
'A chat fork can copy at most 200 inline images.'
52-
)
53-
}
18+
const published = inlineChatImageReferences(message)
19+
if (!published) continue
20+
for (const reference of published.references) {
21+
const sourceKey = inlineChatImageKey(sourceChatId, published.requestId, reference)
22+
tasks.set(sourceKey, {
23+
copyId: sourceKey,
24+
sourceKey,
25+
targetKey: inlineChatImageKey(newChatId, published.requestId, reference),
26+
context: 'mothership',
27+
fileName: 'chat-image.webp',
28+
contentType: 'image/webp',
29+
maxBytes: INLINE_CHAT_IMAGE_MAX_BYTES,
30+
persistMetadata: false,
31+
})
32+
if (tasks.size > 200)
33+
throw new OrchestrationError(
34+
'payload_too_large',
35+
'A chat fork can copy at most 200 inline images.'
36+
)
5437
}
5538
}
5639
return [...tasks.values()]

‎apps/sim/lib/mothership/chat/inline-image-key.ts‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
import { createHash } from 'node:crypto'
2+
import {
3+
INLINE_CHAT_IMAGE_PREFIX,
4+
inlineImageRequestIdSchema,
5+
inlineImageSourceSchema,
6+
normalizeInlineFileReference,
7+
} from '@/lib/mothership/chat/inline-image-reference'
8+
import { collectMarkdownImageSources } from '@/lib/mothership/chat/markdown-images'
9+
import type { PersistedMessage } from '@/lib/mothership/chat/persisted-message'
10+
11+
/** Stable per-message source identity makes replay independent of sandbox lifetime. */
12+
export function inlineChatImageKey(chatId: string, requestId: string, reference: string): string {
13+
inlineImageRequestIdSchema.parse(chatId)
14+
inlineImageRequestIdSchema.parse(requestId)
15+
const digest = createHash('sha256').update(normalizeInlineFileReference(reference)).digest('hex')
16+
return `${INLINE_CHAT_IMAGE_PREFIX}${chatId}/${requestId}/${digest}.webp`
17+
}
18+
19+
/**
20+
* The inline image references an assistant message published under its request id: every
21+
* first-party Markdown image source in its text and text blocks. Anything else (user text,
22+
* a message with no request receipt, remote URLs) never had a chat image stored.
23+
*/
24+
export function inlineChatImageReferences(
25+
message: Pick<PersistedMessage, 'role' | 'requestId' | 'content' | 'contentBlocks'>
26+
): { requestId: string; references: string[] } | null {
27+
if (
28+
message.role !== 'assistant' ||
29+
!message.requestId ||
30+
!inlineImageRequestIdSchema.safeParse(message.requestId).success
31+
)
32+
return null
33+
const contents = [
34+
message.content,
35+
...(message.contentBlocks ?? [])
36+
.filter((block) => block.type === 'text')
37+
.map((block) => block.content ?? ''),
38+
]
39+
const references = new Set<string>()
40+
for (const content of contents) {
41+
// Every Markdown image, inline or reference-style, opens with `![`; most messages have none.
42+
if (typeof content !== 'string' || !content.includes('![')) continue
43+
for (const reference of collectMarkdownImageSources(content))
44+
if (inlineImageSourceSchema.safeParse(reference).success) references.add(reference)
45+
}
46+
return { requestId: message.requestId, references: [...references] }
47+
}

‎apps/sim/lib/mothership/chat/inline-image-storage.ts‎

Lines changed: 1 addition & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,6 @@
1-
import { createHash } from 'node:crypto'
21
import sharp from 'sharp'
32
import { OrchestrationError } from '@/lib/core/orchestration/types'
4-
import {
5-
INLINE_CHAT_IMAGE_PREFIX,
6-
inlineImageRequestIdSchema,
7-
normalizeInlineFileReference,
8-
} from '@/lib/mothership/chat/inline-image-reference'
3+
import { inlineChatImageKey } from '@/lib/mothership/chat/inline-image-key'
94
import { isObjectNotFoundError } from '@/lib/uploads/core/errors'
105
import { downloadFile, uploadFile } from '@/lib/uploads/core/storage-service'
116

@@ -112,11 +107,3 @@ export async function loadInlineChatImage(
112107
throw error
113108
}
114109
}
115-
116-
/** Stable per-message source identity makes replay independent of sandbox lifetime. */
117-
export function inlineChatImageKey(chatId: string, requestId: string, reference: string): string {
118-
inlineImageRequestIdSchema.parse(chatId)
119-
inlineImageRequestIdSchema.parse(requestId)
120-
const digest = createHash('sha256').update(normalizeInlineFileReference(reference)).digest('hex')
121-
return `${INLINE_CHAT_IMAGE_PREFIX}${chatId}/${requestId}/${digest}.webp`
122-
}

0 commit comments

Comments
 (0)