From 43abe940fd05d77ba050910ab3fc463f6e5cb20d Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Fri, 11 Sep 2026 11:18:02 -0700 Subject: [PATCH 1/2] feat(assistant): support image attachments --- apps/sim/app/api/files/authorization.test.ts | 8 + apps/sim/app/api/files/authorization.ts | 2 + .../api/files/serve/[...path]/route.test.ts | 50 +++- .../app/api/files/serve/[...path]/route.ts | 17 ++ apps/sim/app/api/files/uploads/finalizers.ts | 4 + apps/sim/app/api/files/uploads/purposes.ts | 1 + apps/sim/app/api/files/uploads/route.test.ts | 43 ++++ .../components/composer/composer.test.tsx | 133 ++++++++++- .../home/components/composer/composer.tsx | 124 +++++++--- .../home/organization-home.test.tsx | 94 ++++++++ .../home/organization-home.tsx | 45 +++- .../components/drop-overlay/drop-overlay.tsx | 13 +- .../home/hooks/use-chat.mount-send.test.tsx | 31 +++ .../[workspaceId]/home/hooks/use-chat.ts | 4 +- .../hooks/use-file-attachments.test.tsx | 51 +++- .../user-input/hooks/use-file-attachments.ts | 32 ++- apps/sim/lib/api/contracts/upload-sessions.ts | 30 ++- .../lib/copilot/chat/assistant-images.test.ts | 104 ++++++++ apps/sim/lib/copilot/chat/assistant-images.ts | 68 ++++++ apps/sim/lib/copilot/chat/payload.test.ts | 27 +++ apps/sim/lib/copilot/chat/payload.ts | 5 + apps/sim/lib/copilot/chat/post.test.ts | 138 +++++++++++ apps/sim/lib/copilot/chat/post.ts | 152 +++++++----- apps/sim/lib/core/utils/browser-storage.ts | 30 ++- apps/sim/lib/mothership/events.ts | 2 +- apps/sim/lib/uploads/client/admission.ts | 9 +- apps/sim/lib/uploads/client/session-upload.ts | 12 +- .../application.test.ts | 226 ++++++++++++++++++ .../organization-assistant/application.ts | 181 ++++++++++++++ .../organization-assistant/binding.ts | 56 +++++ .../lib/uploads/shared/assistant-images.ts | 15 ++ .../upload-session/application.test.ts | 75 +++++- .../lib/uploads/upload-session/application.ts | 20 ++ .../uploads/upload-session/service.test.ts | 58 +++++ .../sim/lib/uploads/upload-session/service.ts | 64 ++++- apps/sim/lib/uploads/utils/file-utils.ts | 1 + 36 files changed, 1779 insertions(+), 146 deletions(-) create mode 100644 apps/sim/lib/copilot/chat/assistant-images.test.ts create mode 100644 apps/sim/lib/copilot/chat/assistant-images.ts create mode 100644 apps/sim/lib/uploads/contexts/organization-assistant/application.test.ts create mode 100644 apps/sim/lib/uploads/contexts/organization-assistant/application.ts create mode 100644 apps/sim/lib/uploads/contexts/organization-assistant/binding.ts create mode 100644 apps/sim/lib/uploads/shared/assistant-images.ts diff --git a/apps/sim/app/api/files/authorization.test.ts b/apps/sim/app/api/files/authorization.test.ts index 22525582e75..a3b17bfab40 100644 --- a/apps/sim/app/api/files/authorization.test.ts +++ b/apps/sim/app/api/files/authorization.test.ts @@ -63,6 +63,14 @@ function grantAccess(cloudKey: string) { } describe('verifyKBFileAccess (binding-only)', () => { + it.each(['mothership', 'profile-pictures', 'general'] as const)( + 'refuses organization image keys through legacy %s authorization', + async (context) => { + await expect( + verifyFileAccess('assistant/org-1/user-1/upload-1/image.png', USER_ID, undefined, context) + ).resolves.toBe(false) + } + ) beforeEach(() => { vi.clearAllMocks() // Default liveness query result: one active document references the exact storage key. diff --git a/apps/sim/app/api/files/authorization.ts b/apps/sim/app/api/files/authorization.ts index 618b4e9711e..56f5d8e2cd8 100644 --- a/apps/sim/app/api/files/authorization.ts +++ b/apps/sim/app/api/files/authorization.ts @@ -150,6 +150,8 @@ export async function verifyFileAccess( isLocal?: boolean, options?: { requireWrite?: boolean; knowledgeAccess?: KnowledgeFileAccess } ): Promise { + /** Organization images require the Principal-aware Assistant application resolver. */ + if (cloudKey.startsWith('assistant/')) return false const requireWrite = options?.requireWrite ?? false try { const keyContext = inferContextFromKey(cloudKey) diff --git a/apps/sim/app/api/files/serve/[...path]/route.test.ts b/apps/sim/app/api/files/serve/[...path]/route.test.ts index 755861df05f..33b08a07e0c 100644 --- a/apps/sim/app/api/files/serve/[...path]/route.test.ts +++ b/apps/sim/app/api/files/serve/[...path]/route.test.ts @@ -3,7 +3,12 @@ * * @vitest-environment node */ -import { hybridAuthMockFns, storageServiceMock, storageServiceMockFns } from '@sim/testing' +import { + authMockFns, + hybridAuthMockFns, + storageServiceMock, + storageServiceMockFns, +} from '@sim/testing' import { NextRequest } from 'next/server' import { beforeEach, describe, expect, it, vi } from 'vitest' import { MAX_BUFFERED_TRANSFER_BYTES } from '@/lib/uploads/shared/types' @@ -34,6 +39,7 @@ const { mockCreateErrorResponse, FileNotFoundError, serveLogger, + mockReadOrganizationAssistantImage, } = vi.hoisted(() => { class FileNotFoundErrorClass extends Error { constructor(message: string) { @@ -43,6 +49,7 @@ const { } return { serveLogger: { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn() }, + mockReadOrganizationAssistantImage: vi.fn(), mockVerifyFileAccess: vi.fn(), mockReadFile: vi.fn(), mockIsUsingCloudStorage: vi.fn(), @@ -62,6 +69,10 @@ const { } }) +vi.mock('@/lib/uploads/contexts/organization-assistant/application', () => ({ + readOrganizationAssistantImage: mockReadOrganizationAssistantImage, +})) + vi.mock('fs/promises', () => ({ readFile: mockReadFile, access: vi.fn().mockResolvedValue(undefined), @@ -204,6 +215,43 @@ describe('File Serve API Route', () => { }) }) + it('serves private Assistant images through session authorization and disables caching', async () => { + authMockFns.mockGetSession.mockResolvedValue({ + user: { id: 'user-1' }, + session: { id: 'session-1' }, + }) + const key = 'assistant/org-1/user-1/upload-1/image.png' + mockReadOrganizationAssistantImage.mockResolvedValue({ + name: 'image.png', + contentType: 'image/webp', + buffer: Buffer.from('decoded-image'), + }) + const response = await GET(new NextRequest(`http://localhost/api/files/serve/${key}`), { + params: Promise.resolve({ path: key.split('/') }), + }) + expect(response.status).toBe(200) + expect(mockReadOrganizationAssistantImage).toHaveBeenCalledWith({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + key, + signal: expect.any(AbortSignal), + }) + expect(mockCreateFileResponse).toHaveBeenCalledWith( + expect.objectContaining({ cacheControl: 'private, no-store', contentType: 'image/webp' }) + ) + expect(mockVerifyFileAccess).not.toHaveBeenCalled() + expect(hybridAuthMockFns.mockCheckSessionOrInternalAuth).not.toHaveBeenCalled() + }) + + it('requires a real session for private Assistant images even when legacy auth succeeds', async () => { + authMockFns.mockGetSession.mockResolvedValue(null) + const key = 'assistant/org-1/user-1/upload-1/image.png' + const response = await GET(new NextRequest(`http://localhost/api/files/serve/${key}`), { + params: Promise.resolve({ path: key.split('/') }), + }) + expect(response.status).toBe(401) + expect(mockReadOrganizationAssistantImage).not.toHaveBeenCalled() + }) + it('bounds the local read rather than trusting the stored size', async () => { await GET(new NextRequest('http://localhost:3000/api/files/serve/workspace/ws/test-file.txt'), { params: Promise.resolve({ path: ['workspace', 'ws', 'test-file.txt'] }), diff --git a/apps/sim/app/api/files/serve/[...path]/route.ts b/apps/sim/app/api/files/serve/[...path]/route.ts index ffb8845aeff..6ccd6e7eb6e 100644 --- a/apps/sim/app/api/files/serve/[...path]/route.ts +++ b/apps/sim/app/api/files/serve/[...path]/route.ts @@ -7,6 +7,7 @@ import { fileServeParamsSchema, fileServeQuerySchema } from '@/lib/api/contracts import { concealCrossTenantResourceError, InternalUnauthenticatedError, + internalSessionAuth, } from '@/lib/api/server/routes' import { AuthType, checkSessionOrInternalAuth } from '@/lib/auth/hybrid' import { resolveServableDocBytes } from '@/lib/copilot/tools/server/files/doc-compile' @@ -16,6 +17,7 @@ import { assertKnownSizeWithinLimit, isPayloadSizeLimitError } from '@/lib/core/ import { withRouteHandler } from '@/lib/core/utils/with-route-handler' import { CopilotFiles, isUsingCloudStorage } from '@/lib/uploads' import type { StorageContext } from '@/lib/uploads/config' +import { readOrganizationAssistantImage } from '@/lib/uploads/contexts/organization-assistant/application' import { parseWorkspaceFileKey } from '@/lib/uploads/contexts/workspace/workspace-file-manager' import { downloadFile } from '@/lib/uploads/core/storage-service' import { resolveServableImageBytes } from '@/lib/uploads/server/image-derivative' @@ -218,6 +220,21 @@ export const GET = withRouteHandler( const isCloudPath = isS3Path || isBlobPath || isGcsPath const cloudKey = isCloudPath ? path.slice(1).join('/') : fullPath + if (cloudKey.startsWith('assistant/')) { + const principal = await internalSessionAuth.authenticate() + const image = await readOrganizationAssistantImage({ + principal, + key: cloudKey, + signal: request.signal, + }) + return createFileResponse({ + buffer: image.buffer, + filename: image.name, + contentType: image.contentType, + cacheControl: 'private, no-store', + }) + } + const isPublicByKeyPrefix = cloudKey.startsWith('profile-pictures/') || cloudKey.startsWith('og-images/') || diff --git a/apps/sim/app/api/files/uploads/finalizers.ts b/apps/sim/app/api/files/uploads/finalizers.ts index bc51d3a034a..3eef20d27ff 100644 --- a/apps/sim/app/api/files/uploads/finalizers.ts +++ b/apps/sim/app/api/files/uploads/finalizers.ts @@ -9,6 +9,7 @@ import type { OrchestrationRequestContext } from '@/lib/core/orchestration/types import { captureServerEvent } from '@/lib/posthog/server' import { notifyWorkspaceFilesChanged } from '@/lib/realtime/notify' import { getServeStoragePrefix } from '@/lib/uploads/config' +import { finalizeOrganizationAssistantAttachment } from '@/lib/uploads/contexts/organization-assistant/application' import { getWorkspaceFile, registerUploadedWorkspaceFile, @@ -107,6 +108,9 @@ export async function finalizeUploadPurpose({ case 'workspace_logo': return finalizeWorkspaceLogo(session, actor, request) case 'mothership_attachment': + if (session.workspaceId === null) { + return { value: await finalizeOrganizationAssistantAttachment(principal, session) } + } return finalizeMothershipAttachment(session) case 'execution_attachment': return finalizeExecutionAttachment(session) diff --git a/apps/sim/app/api/files/uploads/purposes.ts b/apps/sim/app/api/files/uploads/purposes.ts index dddddf82663..41e50569af9 100644 --- a/apps/sim/app/api/files/uploads/purposes.ts +++ b/apps/sim/app/api/files/uploads/purposes.ts @@ -78,6 +78,7 @@ export async function createPurposeUploadSession( localOrigin, }) case 'mothership_attachment': + if (!body.workspaceId) throw new UploadSessionError('validation', 'workspaceId is required') await requireWorkspacePermission(userId, body.workspaceId, 'write') return createUploadSession({ purpose: body.purpose, diff --git a/apps/sim/app/api/files/uploads/route.test.ts b/apps/sim/app/api/files/uploads/route.test.ts index f36e29862d1..4dc5aab2175 100644 --- a/apps/sim/app/api/files/uploads/route.test.ts +++ b/apps/sim/app/api/files/uploads/route.test.ts @@ -234,6 +234,49 @@ describe('/api/files/uploads', () => { ) }) + it.each([ + { organizationId: 'org-1', contentType: 'application/pdf', size: 100 }, + { organizationId: 'org-1', contentType: 'image/svg+xml', size: 100 }, + { organizationId: 'org-1', contentType: 'image/png', size: 5 * 1024 * 1024 + 1 }, + { organizationId: 'org-1', workspaceId: 'ws-1', contentType: 'image/png', size: 100 }, + ])('rejects unsupported organization attachments before application loading', async (body) => { + const response = await createUpload( + new NextRequest('http://localhost/api/files/uploads', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ purpose: 'mothership_attachment', name: 'image.png', ...body }), + }) + ) + expect(response.status).toBe(400) + expect(mockCreateInternalPurposeUploadSession).not.toHaveBeenCalled() + }) + + it('creates organization image attachments through the same upload lifecycle', async () => { + mockCreateInternalPurposeUploadSession.mockResolvedValue({ + ...session({ purpose: 'mothership_attachment', storageContext: 'mothership' }), + transfer: { method: 'put', url: 'https://storage.example/upload', headers: {} }, + }) + const response = await createUpload( + new NextRequest('http://localhost/api/files/uploads', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + purpose: 'mothership_attachment', + organizationId: 'org-1', + name: 'image.png', + contentType: 'image/png', + size: 100, + }), + }) + ) + expect(response.status).toBe(201) + expect(mockCreateInternalPurposeUploadSession).toHaveBeenCalledWith( + expect.objectContaining({ kind: 'session', userId: 'user-1' }), + expect.objectContaining({ purpose: 'mothership_attachment', organizationId: 'org-1' }), + expect.anything() + ) + }) + it('rejects mothership attachments above the 5 GiB direct-to-storage limit', async () => { const request = new NextRequest('http://localhost/api/files/uploads', { method: 'POST', diff --git a/apps/sim/app/o/[organizationId]/home/components/composer/composer.test.tsx b/apps/sim/app/o/[organizationId]/home/components/composer/composer.test.tsx index b8e873529b2..65ee1653466 100644 --- a/apps/sim/app/o/[organizationId]/home/components/composer/composer.test.tsx +++ b/apps/sim/app/o/[organizationId]/home/components/composer/composer.test.tsx @@ -9,9 +9,11 @@ const mocks = vi.hoisted(() => ({ toggleListening: vi.fn(), resetTranscript: vi.fn(), submit: vi.fn(), + upload: vi.fn(), })) vi.mock('@/hooks/use-speech-to-text', () => ({ useSpeechToText: mocks.speech })) +vi.mock('@/lib/uploads/client/session-upload', () => ({ uploadInternalFileSession: mocks.upload })) vi.mock('@/hooks/use-animated-placeholder', () => ({ useAnimatedPlaceholder: () => 'Ask Sim to' })) vi.mock('@/hooks/use-chat-input-focus', () => ({ useChatInputFocus: vi.fn() })) vi.mock('@/app/o/[organizationId]/providers/organization-provider', () => ({ @@ -19,6 +21,7 @@ vi.mock('@/app/o/[organizationId]/providers/organization-provider', () => ({ })) import { Composer } from '@/app/o/[organizationId]/home/components/composer/composer' +import { useFileAttachments } from '@/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/copilot/components/user-input/hooks/use-file-attachments' let root: Root let container: HTMLDivElement @@ -26,6 +29,17 @@ let container: HTMLDivElement beforeEach(() => { vi.clearAllMocks() vi.stubGlobal('IS_REACT_ACT_ENVIRONMENT', true) + vi.stubGlobal( + 'URL', + class extends URL { + static createObjectURL = vi.fn(() => 'blob:image-preview') + static revokeObjectURL = vi.fn() + } + ) + mocks.upload.mockResolvedValue({ + key: 'assistant/organization-a/user-a/image-a/screenshot.png', + path: '/api/files/serve/image-a?context=mothership', + }) vi.stubGlobal( 'matchMedia', vi.fn(() => ({ @@ -50,21 +64,25 @@ afterEach(async () => { await act(async () => root.unmount()) container.remove() vi.unstubAllGlobals() + vi.restoreAllMocks() }) -async function render(isInitialView: boolean) { +async function render(isInitialView: boolean, initialValue = 'Summarize') { function Harness() { - const [value, setValue] = useState('Summarize') + const [value, setValue] = useState(initialValue) + const files = useFileAttachments({ userId: 'user-a', organizationId: 'organization-a' }) return ( { - mocks.submit(value) + mocks.submit(value, files.attachedFiles) setValue('') + files.clearAttachedFiles() }} /> ) @@ -90,7 +108,7 @@ describe('organization voice composer', () => { await act(async () => { container.querySelector('button[aria-label="Send"]')!.click() }) - expect(mocks.submit).toHaveBeenCalledWith('Summarize the release') + expect(mocks.submit).toHaveBeenCalledWith('Summarize the release', []) expect(mocks.resetTranscript).toHaveBeenCalledOnce() await act(async () => mocks.speech.mock.calls.at(-1)![0].onTranscript('Next question')) expect(container.querySelector('textarea')!.value).toBe('Next question') @@ -109,3 +127,110 @@ describe('organization voice composer', () => { expect(container.querySelector('button[aria-label="Voice input"]')).toBeNull() }) }) + +function fileList(files: File[]): FileList { + return Object.assign(files, { item: (index: number) => files[index] ?? null }) +} + +async function paste(files: File[]) { + const event = new Event('paste', { bubbles: true, cancelable: true }) + Object.defineProperty(event, 'clipboardData', { value: { files: fileList(files) } }) + await act(async () => container.querySelector('textarea')!.dispatchEvent(event)) + return event +} + +describe('organization image composer', () => { + it.each([true, false])( + 'pastes and submits an image without text (initial: %s)', + async (initial) => { + await render(initial, '') + const image = new File(['image'], 'screenshot.png', { type: 'image/png' }) + const event = await paste([image]) + expect(event.defaultPrevented).toBe(true) + expect(mocks.upload).toHaveBeenCalledWith( + expect.objectContaining({ + purpose: 'mothership_attachment', + organizationId: 'organization-a', + file: image, + }) + ) + expect(container.querySelector('img')?.getAttribute('alt')).toBe('screenshot.png') + await act(async () => + container.querySelector('button[aria-label="Send"]')!.click() + ) + expect(mocks.submit).toHaveBeenCalledWith('', [ + expect.objectContaining({ + key: 'assistant/organization-a/user-a/image-a/screenshot.png', + uploading: false, + }), + ]) + expect(container.querySelector('img')).toBeNull() + } + ) + + it('leaves ordinary text paste to the textarea', async () => { + await render(true) + expect((await paste([])).defaultPrevented).toBe(false) + expect(mocks.upload).not.toHaveBeenCalled() + }) + + it('accepts dropped images through the same upload flow', async () => { + await render(false) + const image = new File(['image'], 'dropped.png', { type: 'image/png' }) + const drop = new Event('drop', { bubbles: true, cancelable: true }) + Object.defineProperty(drop, 'dataTransfer', { value: { files: fileList([image]) } }) + await act(async () => container.querySelector('textarea')!.dispatchEvent(drop)) + expect(drop.defaultPrevented).toBe(true) + expect(mocks.upload).toHaveBeenCalledWith(expect.objectContaining({ file: image })) + expect(container.querySelector('img')?.getAttribute('alt')).toBe('dropped.png') + }) + + it('blocks Send and Enter until an image upload finishes', async () => { + let finish!: (value: { key: string; path: string }) => void + mocks.upload.mockImplementation( + () => + new Promise((resolve) => { + finish = resolve + }) + ) + await render(true) + await paste([new File(['image'], 'screenshot.png', { type: 'image/png' })]) + expect(container.querySelector('button[aria-label="Send"]')!.disabled).toBe( + true + ) + await act(async () => + container + .querySelector('textarea')! + .dispatchEvent(new KeyboardEvent('keydown', { key: 'Enter', bubbles: true })) + ) + expect(mocks.submit).not.toHaveBeenCalled() + await act(async () => finish({ key: 'image-key', path: '/image-path' })) + expect(container.querySelector('button[aria-label="Send"]')!.disabled).toBe( + false + ) + }) + + it('uses the picker and lets an attachment be removed before sending', async () => { + await render(true, '') + const input = container.querySelector('input[type="file"]')! + const click = vi.spyOn(input, 'click') + await act(async () => + container.querySelector('button[aria-label="Attach images"]')!.click() + ) + expect(click).toHaveBeenCalledOnce() + expect(input.accept).toContain('image/png') + Object.defineProperty(input, 'files', { + value: fileList([new File(['image'], 'screenshot.png', { type: 'image/png' })]), + }) + await act(async () => input.dispatchEvent(new Event('change', { bubbles: true }))) + await act(async () => + container + .querySelector('button[aria-label="Remove screenshot.png"]')! + .click() + ) + expect(container.querySelector('img')).toBeNull() + expect(container.querySelector('button[aria-label="Send"]')!.disabled).toBe( + true + ) + }) +}) diff --git a/apps/sim/app/o/[organizationId]/home/components/composer/composer.tsx b/apps/sim/app/o/[organizationId]/home/components/composer/composer.tsx index 1c897d4d532..891af13473a 100644 --- a/apps/sim/app/o/[organizationId]/home/components/composer/composer.tsx +++ b/apps/sim/app/o/[organizationId]/home/components/composer/composer.tsx @@ -1,11 +1,15 @@ 'use client' import { useRef } from 'react' -import { Button, cn } from '@sim/emcn' -import { ArrowUp } from '@sim/emcn/icons' +import { Button, Chip, cn, Tooltip } from '@sim/emcn' +import { ArrowUp, Plus } from '@sim/emcn/icons' +import { ASSISTANT_IMAGE_ACCEPT_ATTRIBUTE } from '@/lib/uploads/shared/assistant-images' import { useOrganizationContext } from '@/app/o/[organizationId]/providers/organization-provider' +import { AttachedFilesList } from '@/app/workspace/[workspaceId]/home/components/user-input/components/attached-files-list/attached-files-list' +import { DropOverlay } from '@/app/workspace/[workspaceId]/home/components/user-input/components/drop-overlay/drop-overlay' import { MicButton } from '@/app/workspace/[workspaceId]/home/components/user-input/components/mic-button/mic-button' import { MicrophonePermissionHelp } from '@/app/workspace/[workspaceId]/home/components/user-input/components/microphone-permission-help/microphone-permission-help' +import type { useFileAttachments } from '@/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/copilot/components/user-input/hooks/use-file-attachments' import { useAnimatedPlaceholder } from '@/hooks/use-animated-placeholder' import { useChatInputFocus } from '@/hooks/use-chat-input-focus' import { useVoiceInput } from '@/hooks/use-voice-input' @@ -17,6 +21,7 @@ const SEND_BUTTON_DISABLED = 'bg-[#808080] dark:bg-[#808080]' interface ComposerProps { value: string + files: ReturnType /** On the empty home the placeholder types itself and the field is taller; in a chat it is the plain footer input. */ isInitialView: boolean isSending: boolean @@ -32,6 +37,7 @@ interface ComposerProps { */ export function Composer({ value, + files, isInitialView, isSending, onChange, @@ -46,7 +52,9 @@ export function Composer({ getValue: () => value, onChange, }) - const canSubmit = value.trim().length > 0 + const canSubmit = + !files.attachedFiles.some((file) => file.uploading) && + (value.trim().length > 0 || files.attachedFiles.some((file) => file.key)) const animatedPlaceholder = useAnimatedPlaceholder(isInitialView) const placeholder = isInitialView ? animatedPlaceholder : 'Send message to Sim' @@ -58,11 +66,20 @@ export function Composer({ return (
+
onChange(event.target.value)} + onPaste={(event) => { + const pasted = event.clipboardData.files + if (!pasted.length) return + event.preventDefault() + void files.processFiles(pasted) + }} onKeyDown={(event) => { if (event.key === 'Enter' && !event.shiftKey && !event.nativeEvent.isComposing) { event.preventDefault() @@ -86,43 +109,68 @@ export function Composer({ />
-
- {voice.isSupported && ( - - )} - {isSending ? ( - - ) : ( - - )} + + + + + ) : ( + + )} +
+ + {files.isDragging && } ({ apiKeys: vi.fn(), authorizedApps: vi.fn(), fetchNextPage: vi.fn(), + upload: vi.fn(), })) vi.mock('@/lib/auth/auth-client', () => ({ useSession: () => ({ data: { user: { id: 'reader' } } }), })) +vi.mock('@/lib/uploads/client/session-upload', () => ({ uploadInternalFileSession: mocks.upload })) vi.mock('@/lib/core/utils/browser-storage', () => ({ MothershipHandoffStorage: { consume: mocks.consume }, })) @@ -46,6 +48,14 @@ let container: HTMLDivElement beforeEach(() => { vi.clearAllMocks() vi.stubGlobal('IS_REACT_ACT_ENVIRONMENT', true) + vi.stubGlobal( + 'URL', + class extends URL { + static createObjectURL = vi.fn(() => 'blob:image-preview') + static revokeObjectURL = vi.fn() + } + ) + mocks.upload.mockResolvedValue({ key: 'image-key', path: '/image-path' }) mocks.context.mockReturnValue({ organization: { id: 'organization-a' }, searchAccess: { memberScoped: true }, @@ -290,6 +300,90 @@ describe('organization home', () => { await act(async () => composerProps().onSubmit()) expect(mocks.send).not.toHaveBeenCalled() }) + it('sends image-only turns with canonical attachment properties and clears the draft', async () => { + await act(async () => root.render()) + const files = [new File(['image'], 'screenshot.png', { type: 'image/png' })] + await act(async () => + composerProps().files.processFiles( + Object.assign(files, { item: (index: number) => files[index] ?? null }) + ) + ) + await act(async () => composerProps().onSubmit()) + expect(mocks.send).toHaveBeenCalledWith( + '', + [ + expect.objectContaining({ + id: expect.any(String), + key: 'image-key', + filename: 'screenshot.png', + media_type: 'image/png', + size: 5, + path: '/image-path', + }), + ], + undefined, + { requestMode: 'assistant' } + ) + expect(composerProps().files.attachedFiles).toEqual([]) + }) + + it('restores queued images when editing and includes them in the replacement turn', async () => { + mocks.renderer.mockImplementation(({ composer }: { composer: ReactNode }) => composer) + const attachments = [ + { + id: 'image-a', + key: 'image-key', + filename: 'screenshot.png', + media_type: 'image/png', + size: 5, + }, + ] + mocks.chat.mockReturnValue({ + messages: [], + sendMessage: mocks.send, + editQueuedMessage: () => ({ + id: 'queued-a', + content: 'Explain this', + fileAttachments: attachments, + }), + }) + await act(async () => root.render()) + await act(async () => mocks.renderer.mock.lastCall![0].onEditQueuedMessage('queued-a')) + expect(composerProps().files.attachedFiles[0]).toEqual( + expect.objectContaining({ + name: 'screenshot.png', + key: 'image-key', + uploading: false, + path: '/api/files/serve/image-key?context=mothership&preview=1', + }) + ) + await act(async () => composerProps().onSubmit()) + expect(mocks.send).toHaveBeenCalledWith( + 'Explain this', + [{ ...attachments[0], path: '/api/files/serve/image-key?context=mothership&preview=1' }], + undefined, + { + requestMode: 'assistant', + } + ) + }) + + it('resumes image-only handoffs without dropping their attachments', async () => { + const attachments = [ + { + id: 'image-a', + key: 'image-key', + filename: 'screenshot.png', + media_type: 'image/png', + size: 5, + }, + ] + mocks.consume.mockReturnValueOnce({ message: '', fileAttachments: attachments }) + await act(async () => root.render()) + expect(mocks.send).toHaveBeenCalledWith('', attachments, undefined, { + requestMode: 'assistant', + }) + }) it('resumes a scoped handoff with the original search filters', async () => { const assistantSearch = { documentIds: ['document-a'] } mocks.consume.mockReturnValueOnce({ message: 'Summarize', assistantSearch }) diff --git a/apps/sim/app/o/[organizationId]/home/organization-home.tsx b/apps/sim/app/o/[organizationId]/home/organization-home.tsx index 59d754b9117..34ec83488ae 100644 --- a/apps/sim/app/o/[organizationId]/home/organization-home.tsx +++ b/apps/sim/app/o/[organizationId]/home/organization-home.tsx @@ -2,6 +2,7 @@ import { useEffect, useState } from 'react' import { useSession } from '@/lib/auth/auth-client' +import { getMothershipAttachmentPreviewUrl } from '@/lib/copilot/chat/attachment-preview' import { MothershipHandoffStorage } from '@/lib/core/utils/browser-storage' import { Composer } from '@/app/o/[organizationId]/home/components/composer' import { GetStarted } from '@/app/o/[organizationId]/home/components/get-started' @@ -9,6 +10,8 @@ import { useOrganizationContext } from '@/app/o/[organizationId]/providers/organ import { SearchIntegrationConnection } from '@/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/search-integration-connection' import { MothershipChat } from '@/app/workspace/[workspaceId]/home/components/mothership-chat' import { useChat } from '@/app/workspace/[workspaceId]/home/hooks/use-chat' +import type { FileAttachmentForApi } from '@/app/workspace/[workspaceId]/home/types' +import { useFileAttachments } from '@/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/copilot/components/user-input/hooks/use-file-attachments' import { useMarkMothershipChatRead } from '@/hooks/queries/mothership-chats' interface OrganizationHomeProps { @@ -28,6 +31,7 @@ function OrganizationHomeContent({ userName, chatId }: OrganizationHomeProps) { const { data: session } = useSession() const [draft, setDraft] = useState('') const chat = useChat({ organizationId: organization.id }, chatId) + const files = useFileAttachments({ userId: session?.user?.id, organizationId: organization.id }) const { sendMessage } = chat const { mutate: markRead } = useMarkMothershipChatRead({ organizationId: organization.id }) const firstName = userName?.split(' ')[0] ?? '' @@ -40,8 +44,8 @@ function OrganizationHomeContent({ userName, chatId }: OrganizationHomeProps) { useEffect(() => { if (chatId) return const handoff = MothershipHandoffStorage.consume({ organizationId: organization.id }) - if (handoff?.message) { - void sendMessage(handoff.message, undefined, undefined, { + if (handoff && (handoff.message || handoff.fileAttachments?.length)) { + void sendMessage(handoff.message ?? '', handoff.fileAttachments, undefined, { requestMode: 'assistant', ...(handoff.resumeUserMessageId ? { resumeUserMessageId: handoff.resumeUserMessageId } @@ -51,21 +55,34 @@ function OrganizationHomeContent({ userName, chatId }: OrganizationHomeProps) { } }, [chatId, organization.id, sendMessage]) - const send = (message: string) => { - void sendMessage(message, undefined, undefined, { requestMode: 'assistant' }) + const send = (message: string, fileAttachments?: FileAttachmentForApi[]) => { + void sendMessage(message, fileAttachments, undefined, { requestMode: 'assistant' }) } const submit = () => { const message = draft.trim() - if (!message) return + if (files.attachedFiles.some((file) => file.uploading)) return + const attachments: FileAttachmentForApi[] = files.attachedFiles + .filter((file) => file.key) + .map((file) => ({ + id: file.id, + key: file.key!, + filename: file.name, + media_type: file.type, + size: file.size, + path: file.path, + })) + if (!message && !attachments.length) return setDraft('') - send(message) + send(message, attachments.length ? attachments : undefined) + files.clearAttachedFiles() } const hasChat = Boolean(chatId || chat.messages.length) const composer = ( { const queued = chat.editQueuedMessage(id) - if (queued) setDraft(queued.content) + if (queued) { + setDraft(queued.content) + files.restoreAttachedFiles( + (queued.fileAttachments ?? []).map((file) => ({ + id: file.id, + key: file.key, + name: file.filename, + type: file.media_type, + size: file.size, + path: file.path || getMothershipAttachmentPreviewUrl(file) || '', + previewUrl: getMothershipAttachmentPreviewUrl(file), + uploading: false, + })) + ) + } return queued }} onCancelQueueEdit={chat.cancelQueueEdit} diff --git a/apps/sim/app/workspace/[workspaceId]/home/components/user-input/components/drop-overlay/drop-overlay.tsx b/apps/sim/app/workspace/[workspaceId]/home/components/user-input/components/drop-overlay/drop-overlay.tsx index 03993db036e..08bd1e957d5 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/components/user-input/components/drop-overlay/drop-overlay.tsx +++ b/apps/sim/app/workspace/[workspaceId]/home/components/user-input/components/drop-overlay/drop-overlay.tsx @@ -1,6 +1,7 @@ 'use client' import { memo } from 'react' +import { ImageUp } from '@sim/emcn/icons' import { AudioIcon, CsvIcon, @@ -25,13 +26,19 @@ const DROP_OVERLAY_ICONS = [ VideoIcon, ] as const -export const DropOverlay = memo(function DropOverlay() { +interface DropOverlayProps { + imagesOnly?: boolean +} + +export const DropOverlay = memo(function DropOverlay({ imagesOnly = false }: DropOverlayProps) { return (
- Drop files + + {imagesOnly ? 'Drop images' : 'Drop files'} +
- {DROP_OVERLAY_ICONS.map((Icon, i) => ( + {(imagesOnly ? [ImageUp] : DROP_OVERLAY_ICONS).map((Icon, i) => ( ))}
diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.mount-send.test.tsx b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.mount-send.test.tsx index 7ec24fb8476..875dc709b53 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.mount-send.test.tsx +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.mount-send.test.tsx @@ -301,6 +301,37 @@ async function waitFor(predicate: () => boolean, budgetMs = 2000): Promise } describe('useChat remount send recovery', () => { + it('sends and recovers an image-only organization turn', async () => { + navigationMocks.usePathname.mockReturnValue('/o/org-1/home') + const { getResult, unmount } = renderUseChat({ organizationId: 'org-1' }) + const attachments = [ + { + id: 'image-a', + key: 'image-key', + filename: 'screenshot.png', + media_type: 'image/png', + size: 5, + }, + ] + await act(async () => { + void getResult().sendMessage('', attachments) + }) + await waitFor(() => state.postBodies.length === 1) + expect(state.postBodies[0]).toMatchObject({ + organizationId: 'org-1', + mode: 'assistant', + message: '', + fileAttachments: attachments, + }) + expect(state.postBodies[0]).not.toHaveProperty('workspaceId') + unmount() + await waitFor(() => window.localStorage.getItem('sim_mothership_handoff') !== null) + expect(MothershipHandoffStorage.consume({ organizationId: 'org-1' })).toMatchObject({ + message: '', + fileAttachments: attachments, + }) + }) + it('sends and recovers an organization turn without adding workspace scope', async () => { navigationMocks.usePathname.mockReturnValue('/o/org-1/home') const { getResult, unmount } = renderUseChat({ organizationId: 'org-1' }) diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts index 6752c4cd6e3..0ac26689ac1 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts @@ -3811,7 +3811,7 @@ export function useChat( contexts?: ChatContext[], options?: StartSendMessageOptions ): Promise => { - if (!message.trim() || !scopeKey) return false + if ((!message.trim() && !fileAttachments?.length) || !scopeKey) return false const { onOptimisticSendApplied, queuedSendHandoff } = options ?? {} const pendingStop = options?.pendingStop ?? pendingStopPromiseRef.current const pendingStopStreamId = pendingStop @@ -4339,7 +4339,7 @@ export function useChat( contexts?: ChatContext[], options?: SendMessageOptions ) => { - if (!message.trim() || !scopeKey) return + if ((!message.trim() && !fileAttachments?.length) || !scopeKey) return const queueStore = useMothershipQueueStore.getState() const activeChatKey = chatKeyRef.current diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/copilot/components/user-input/hooks/use-file-attachments.test.tsx b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/copilot/components/user-input/hooks/use-file-attachments.test.tsx index 9d1a3d26d63..520c7fd8c8b 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/copilot/components/user-input/hooks/use-file-attachments.test.tsx +++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/copilot/components/user-input/hooks/use-file-attachments.test.tsx @@ -16,6 +16,10 @@ vi.mock('@/lib/uploads/client/session-upload', () => ({ uploadInternalFileSession: mockUploadInternalFileSession, })) +import { + ASSISTANT_IMAGE_MAX_BYTES, + ASSISTANT_IMAGE_MAX_COUNT, +} from '@/lib/uploads/shared/assistant-images' import { MAX_WORKSPACE_FILE_SIZE } from '@/lib/uploads/shared/types' import { useFileAttachments } from '@/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/copilot/components/user-input/hooks/use-file-attachments' @@ -24,13 +28,15 @@ interface HookHarness { unmount: () => void } -function renderFileAttachmentsHook(): HookHarness { +function renderFileAttachmentsHook( + owner: { workspaceId: string } | { organizationId: string } = { workspaceId: 'workspace-1' } +): HookHarness { ;(globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true const root: Root = createRoot(document.createElement('div')) let latest: ReturnType function Probe() { - latest = useFileAttachments({ userId: 'user-1', workspaceId: 'workspace-1' }) + latest = useFileAttachments({ userId: 'user-1', ...owner }) return null } @@ -115,4 +121,45 @@ describe('useFileAttachments admission', () => { unmount() }) + + it.each(['unsupported', 'oversized', 'too many'] as const)( + 'rejects %s organization images before allocating previews or sessions', + async (kind) => { + const { result, unmount } = renderFileAttachmentsHook({ organizationId: 'organization-1' }) + const files = + kind === 'unsupported' + ? [new File(['pdf'], 'document.pdf', { type: 'application/pdf' })] + : kind === 'oversized' + ? [sizedFile('large.png', ASSISTANT_IMAGE_MAX_BYTES + 1)] + : Array.from({ length: ASSISTANT_IMAGE_MAX_COUNT + 1 }, (_, index) => + sizedFile(`image-${index}.png`, 10) + ) + await act(async () => result().processFiles(asFileList(files))) + expect(mockToastError).toHaveBeenCalledOnce() + expect(createObjectUrl).not.toHaveBeenCalled() + expect(mockUploadInternalFileSession).not.toHaveBeenCalled() + expect(result().attachedFiles).toEqual([]) + unmount() + } + ) + + it('uses organization scope for images and removes a failed upload', async () => { + mockUploadInternalFileSession.mockRejectedValueOnce(new Error('Upload failed')) + const { result, unmount } = renderFileAttachmentsHook({ organizationId: 'organization-1' }) + const file = sizedFile('screenshot.png', 10) + await act(async () => result().processFiles(asFileList([file]))) + expect(mockUploadInternalFileSession).toHaveBeenCalledWith( + expect.objectContaining({ + purpose: 'mothership_attachment', + organizationId: 'organization-1', + file, + }) + ) + expect(mockUploadInternalFileSession.mock.calls[0][0]).not.toHaveProperty('workspaceId') + expect(result().attachedFiles).toEqual([]) + expect(mockToastError).toHaveBeenCalledWith('Couldn\'t upload "screenshot.png"', { + description: 'Upload failed', + }) + unmount() + }) }) diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/copilot/components/user-input/hooks/use-file-attachments.ts b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/copilot/components/user-input/hooks/use-file-attachments.ts index 13a5be86e4f..edf9c7aa24d 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/copilot/components/user-input/hooks/use-file-attachments.ts +++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/copilot/components/user-input/hooks/use-file-attachments.ts @@ -9,6 +9,12 @@ import { getMothershipAttachmentPreviewUrl } from '@/lib/copilot/chat/attachment import { assertMultiFileUploadAdmission } from '@/lib/uploads/client/admission' import { runWithConcurrency, WHOLE_FILE_PARALLEL_UPLOADS } from '@/lib/uploads/client/concurrency' import { uploadInternalFileSession } from '@/lib/uploads/client/session-upload' +import { + ASSISTANT_IMAGE_MAX_BYTES, + ASSISTANT_IMAGE_MAX_COUNT, + ASSISTANT_IMAGE_MAX_TOTAL_BYTES, + isAssistantImageType, +} from '@/lib/uploads/shared/assistant-images' import { MAX_WORKSPACE_FILE_SIZE } from '@/lib/uploads/shared/types' import { resolveFileType } from '@/lib/uploads/utils/file-utils' @@ -78,6 +84,7 @@ export interface MessageFileAttachment { interface UseFileAttachmentsProps { userId?: string workspaceId?: string + organizationId?: string disabled?: boolean isLoading?: boolean } @@ -90,7 +97,7 @@ interface UseFileAttachmentsProps { * @returns File attachment state and operations */ export function useFileAttachments(props: UseFileAttachmentsProps) { - const { userId, workspaceId, disabled, isLoading } = props + const { userId, workspaceId, organizationId, disabled, isLoading } = props const [attachedFiles, setAttachedFiles] = useState([]) const [dragCounter, setDragCounter] = useState(0) @@ -152,16 +159,29 @@ export function useFileAttachments(props: UseFileAttachmentsProps) { logger.error('User ID not available for file upload') return } - if (!workspaceId) { - logger.error('workspaceId required for mothership uploads') + if (!workspaceId && !organizationId) { + logger.error('Workspace or organization context required for attachments') return } if (fileList.length === 0) return try { + if ( + organizationId && + Array.from(fileList).some((file) => !isAssistantImageType(resolveFileType(file))) + ) { + toast.error('Attach PNG, JPEG, GIF, or WebP images.') + return + } assertMultiFileUploadAdmission(fileList, { existingFiles: attachedFilesRef.current, - maxFileBytes: MAX_WORKSPACE_FILE_SIZE, + maxFileBytes: organizationId ? ASSISTANT_IMAGE_MAX_BYTES : MAX_WORKSPACE_FILE_SIZE, + ...(organizationId + ? { + maxFiles: ASSISTANT_IMAGE_MAX_COUNT, + maxTotalBytes: ASSISTANT_IMAGE_MAX_TOTAL_BYTES, + } + : {}), }) } catch (error) { toast.error("Couldn't add files", { description: toError(error).message }) @@ -198,7 +218,7 @@ export function useFileAttachments(props: UseFileAttachmentsProps) { const result = await uploadInternalFileSession({ purpose: 'mothership_attachment', file, - workspaceId, + ...(organizationId ? { organizationId } : { workspaceId: workspaceId! }), signal: controller.signal, }) @@ -236,7 +256,7 @@ export function useFileAttachments(props: UseFileAttachmentsProps) { } }) }, - [userId, workspaceId, updateAttachedFiles] + [userId, workspaceId, organizationId, updateAttachedFiles] ) /** diff --git a/apps/sim/lib/api/contracts/upload-sessions.ts b/apps/sim/lib/api/contracts/upload-sessions.ts index a60749ac704..dfc5938b859 100644 --- a/apps/sim/lib/api/contracts/upload-sessions.ts +++ b/apps/sim/lib/api/contracts/upload-sessions.ts @@ -2,6 +2,7 @@ import { z } from 'zod' import { folderIdSchema, noInputSchema, + organizationIdSchema, workflowIdSchema, workspaceIdSchema, } from '@/lib/api/contracts/primitives' @@ -16,6 +17,10 @@ import { v2UploadTransferSchema, } from '@/lib/api/contracts/v2/uploads' import { executionIdSchema } from '@/lib/api/contracts/workflows' +import { + ASSISTANT_IMAGE_CONTENT_TYPES, + ASSISTANT_IMAGE_MAX_BYTES, +} from '@/lib/uploads/shared/assistant-images' import { MAX_WORKSPACE_FILE_SIZE, MAX_WORKSPACE_FORMDATA_FILE_SIZE, @@ -62,9 +67,30 @@ export const createInternalFileUploadBodySchema = z.discriminatedUnion('purpose' purpose: z.literal('mothership_attachment'), ...internalFileUploadBaseShape, size: z.number().int().min(1).max(MAX_WORKSPACE_FILE_SIZE), - workspaceId: workspaceIdSchema, + workspaceId: workspaceIdSchema.optional(), + organizationId: organizationIdSchema.optional(), }) - .strict(), + .strict() + .superRefine((body, ctx) => { + if (Boolean(body.workspaceId) === Boolean(body.organizationId)) { + ctx.addIssue({ + code: 'custom', + path: ['workspaceId'], + message: 'Provide exactly one workspaceId or organizationId', + }) + } + if ( + body.organizationId && + (body.size > ASSISTANT_IMAGE_MAX_BYTES || + !ASSISTANT_IMAGE_CONTENT_TYPES.some((type) => type === body.contentType)) + ) { + ctx.addIssue({ + code: 'custom', + path: ['contentType'], + message: 'Assistant attachments must be PNG, JPEG, GIF, or WebP images up to 5 MB', + }) + } + }), z .object({ purpose: z.literal('execution_attachment'), diff --git a/apps/sim/lib/copilot/chat/assistant-images.test.ts b/apps/sim/lib/copilot/chat/assistant-images.test.ts new file mode 100644 index 00000000000..5a11b4905de --- /dev/null +++ b/apps/sim/lib/copilot/chat/assistant-images.test.ts @@ -0,0 +1,104 @@ +/** @vitest-environment node */ +import type { SessionPrincipal } from '@sim/auth/principal' +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { OrchestrationError } from '@/lib/core/orchestration/types' +import type { readOrganizationAssistantImage } from '@/lib/uploads/contexts/organization-assistant/application' +import { ASSISTANT_IMAGE_MAX_COUNT } from '@/lib/uploads/shared/assistant-images' + +const { readImage } = vi.hoisted(() => ({ + readImage: vi.fn(), +})) +vi.mock('@/lib/uploads/contexts/organization-assistant/application', () => ({ + readOrganizationAssistantImage: readImage, +})) + +import { prepareAssistantImages } from '@/lib/copilot/chat/assistant-images' +import { getMothershipAttachmentPreviewUrl } from '@/lib/copilot/chat/attachment-preview' + +const principal: SessionPrincipal = { kind: 'session', userId: 'user-1', sessionId: 'session-1' } +const key = 'assistant/org-1/user-1/upload-1/image.png' +const image = { + id: 'upload-1', + key, + name: 'image.png', + contentType: 'image/png', + size: 5, + buffer: Buffer.from('image'), +} + +describe('Assistant image preparation', () => { + beforeEach(() => { + vi.clearAllMocks() + readImage.mockResolvedValue(image) + }) + + it('uses canonical metadata and model content from the authorized reader', async () => { + const signal = new AbortController().signal + const result = await prepareAssistantImages({ + principal, + organizationId: 'org-1', + attachments: [{ key }], + signal, + }) + expect(readImage).toHaveBeenCalledWith({ principal, organizationId: 'org-1', key, signal }) + expect(result).toEqual({ + attachments: [ + { id: 'upload-1', key, filename: 'image.png', media_type: 'image/png', size: 5 }, + ], + content: [ + { + type: 'image', + filename: 'image.png', + source: { type: 'base64', media_type: 'image/png', data: 'aW1hZ2U=' }, + }, + ], + }) + expect(getMothershipAttachmentPreviewUrl(result.attachments[0])).toBe( + `/api/files/serve/${encodeURIComponent(key)}?context=mothership&preview=1` + ) + }) + + it('rejects an oversized batch before reading any image', async () => { + await expect( + prepareAssistantImages({ + principal, + organizationId: 'org-1', + attachments: Array.from({ length: ASSISTANT_IMAGE_MAX_COUNT + 1 }, () => ({ key })), + }) + ).rejects.toMatchObject({ code: 'validation' }) + expect(readImage).not.toHaveBeenCalled() + }) + + it('fails the entire turn if any image is inaccessible', async () => { + readImage.mockRejectedValueOnce(new OrchestrationError('not_found', 'Image not found')) + await expect( + prepareAssistantImages({ + principal, + organizationId: 'org-1', + attachments: [{ key }, { key: 'another-image' }], + }) + ).rejects.toMatchObject({ code: 'not_found' }) + expect(readImage).toHaveBeenCalledOnce() + }) + + it('rejects non-image content even if an upstream reader returns it', async () => { + readImage.mockResolvedValueOnce({ ...image, contentType: 'application/pdf' }) + await expect( + prepareAssistantImages({ principal, organizationId: 'org-1', attachments: [{ key }] }) + ).rejects.toMatchObject({ code: 'validation' }) + }) + + it('does not read images after the request is aborted', async () => { + const controller = new AbortController() + controller.abort() + await expect( + prepareAssistantImages({ + principal, + organizationId: 'org-1', + attachments: [{ key }], + signal: controller.signal, + }) + ).rejects.toMatchObject({ name: 'AbortError' }) + expect(readImage).not.toHaveBeenCalled() + }) +}) diff --git a/apps/sim/lib/copilot/chat/assistant-images.ts b/apps/sim/lib/copilot/chat/assistant-images.ts new file mode 100644 index 00000000000..c2f6f643743 --- /dev/null +++ b/apps/sim/lib/copilot/chat/assistant-images.ts @@ -0,0 +1,68 @@ +import type { SessionPrincipal } from '@sim/auth/principal' +import type { PersistedFileAttachment } from '@/lib/copilot/chat/persisted-message' +import { OrchestrationError } from '@/lib/core/orchestration/types' +import { readOrganizationAssistantImage } from '@/lib/uploads/contexts/organization-assistant/application' +import { + ASSISTANT_IMAGE_MAX_COUNT, + ASSISTANT_IMAGE_MAX_TOTAL_BYTES, +} from '@/lib/uploads/shared/assistant-images' +import { createFileContent, type MessageContent } from '@/lib/uploads/utils/file-utils' + +export interface AssistantImageContent extends MessageContent { + type: 'image' + filename: string +} + +interface PreparedAssistantImages { + attachments: PersistedFileAttachment[] + content: AssistantImageContent[] +} + +/** Resolves private uploads before any attachment metadata or bytes enter a chat turn. */ +export async function prepareAssistantImages({ + principal, + organizationId, + attachments, + signal, +}: { + principal: SessionPrincipal + organizationId: string + attachments: readonly { key: string }[] + signal?: AbortSignal +}): Promise { + if (attachments.length > ASSISTANT_IMAGE_MAX_COUNT) { + throw new OrchestrationError( + 'validation', + `Attach up to ${ASSISTANT_IMAGE_MAX_COUNT} images per message` + ) + } + + const prepared: PreparedAssistantImages = { attachments: [], content: [] } + let totalBytes = 0 + for (const attachment of attachments) { + signal?.throwIfAborted() + const image = await readOrganizationAssistantImage({ + principal, + organizationId, + key: attachment.key, + signal, + }) + totalBytes += image.buffer.length + if (totalBytes > ASSISTANT_IMAGE_MAX_TOTAL_BYTES) { + throw new OrchestrationError('payload_too_large', 'Attached images are too large') + } + const content = createFileContent(image.buffer, image.contentType) + if (content?.type !== 'image') { + throw new OrchestrationError('validation', 'Assistant attachments must be supported images') + } + prepared.attachments.push({ + id: image.id, + key: image.key, + filename: image.name, + media_type: image.contentType, + size: image.size, + }) + prepared.content.push({ ...content, type: 'image', filename: image.name }) + } + return prepared +} diff --git a/apps/sim/lib/copilot/chat/payload.test.ts b/apps/sim/lib/copilot/chat/payload.test.ts index 5c504effebe..93a86f0b825 100644 --- a/apps/sim/lib/copilot/chat/payload.test.ts +++ b/apps/sim/lib/copilot/chat/payload.test.ts @@ -647,6 +647,33 @@ describe('Assistant payload', () => { mockIsIntegrationDeploymentAvailable.mockReturnValue(true) mockCreateUserToolSchema.mockReturnValue({ type: 'object', properties: {} }) }) + it('sends prepared organization images as model-readable attachments without workspace tracking', async () => { + mockTrackChatUpload.mockClear() + const image = { + type: 'image' as const, + filename: 'image.png', + source: { type: 'base64' as const, media_type: 'image/png', data: 'aW1hZ2U=' }, + } + const payload = await buildCopilotRequestPayload( + { + message: '', + userId: 'user-1', + userMessageId: 'message-1', + organizationId: 'org-1', + mode: 'assistant', + model: '', + assistantImages: [image], + fileAttachments: [{ id: 'image', key: 'private-upload', size: 5 }], + }, + { selectedModel: '' } + ) + expect(payload.message).toBe('') + expect(payload.fileAttachments).toEqual([image]) + expect(payload).not.toHaveProperty('context') + expect(payload).not.toHaveProperty('workspaceId') + expect(mockTrackChatUpload).not.toHaveBeenCalled() + }) + it('forwards organization scope without workspace, integration, or desktop authority', async () => { const payload = await buildCopilotRequestPayload( { diff --git a/apps/sim/lib/copilot/chat/payload.ts b/apps/sim/lib/copilot/chat/payload.ts index ff76023488c..35a48a938a7 100644 --- a/apps/sim/lib/copilot/chat/payload.ts +++ b/apps/sim/lib/copilot/chat/payload.ts @@ -10,6 +10,7 @@ import { isAssistantIntegrationTool, } from '@/lib/copilot/assistant/tool-policy' import { getBlockVisibilityForCopilot, visibilitySignature } from '@/lib/copilot/block-visibility' +import type { AssistantImageContent } from '@/lib/copilot/chat/assistant-images' import type { VfsSnapshotV1 } from '@/lib/copilot/generated/vfs-snapshot-v1' import { type IntegrationGateConfig, @@ -52,6 +53,7 @@ interface BuildPayloadParams { */ mcpServerIds?: string[] fileAttachments?: Array<{ id: string; key: string; size: number; [key: string]: unknown }> + assistantImages?: AssistantImageContent[] commands?: string[] chatId?: string prefetch?: boolean @@ -412,6 +414,9 @@ export async function buildCopilotRequestPayload( ...(provider ? { provider } : {}), mode: transportMode, ...(isAssistant && params.assistantSearch ? { assistantSearch: params.assistantSearch } : {}), + ...(isAssistant && params.organizationId && params.assistantImages?.length + ? { fileAttachments: params.assistantImages } + : {}), messageId: userMessageId, ...(allContexts.length > 0 ? { context: allContexts } : {}), ...(chatId ? { chatId } : {}), diff --git a/apps/sim/lib/copilot/chat/post.test.ts b/apps/sim/lib/copilot/chat/post.test.ts index 6a0d6914736..a06b91590d8 100644 --- a/apps/sim/lib/copilot/chat/post.test.ts +++ b/apps/sim/lib/copilot/chat/post.test.ts @@ -40,6 +40,7 @@ const { resolveBillingAttribution, resolveOrganizationBillingAttribution, authorizeOrganizationChat, + readOrganizationAssistantImage, finalizeAssistantTurn, appendCopilotChatMessages, persistChatResources, @@ -62,6 +63,7 @@ const { resolveBillingAttribution: vi.fn(), resolveOrganizationBillingAttribution: vi.fn(), authorizeOrganizationChat: vi.fn(), + readOrganizationAssistantImage: vi.fn(), finalizeAssistantTurn: vi.fn(), appendCopilotChatMessages: vi.fn(), persistChatResources: vi.fn(), @@ -135,6 +137,10 @@ vi.mock('@/lib/copilot/chat/organization-chats', () => ({ authorizeOrganizationChat: { execute: authorizeOrganizationChat }, })) +vi.mock('@/lib/uploads/contexts/organization-assistant/application', () => ({ + readOrganizationAssistantImage, +})) + vi.mock('@/lib/credentials/application/personal-credentials', () => ({ listPersonalCredentials: { execute: listPersonal }, })) @@ -237,6 +243,14 @@ describe('handleUnifiedChatPost', () => { userId: 'user-1', role: 'member', }) + readOrganizationAssistantImage.mockResolvedValue({ + id: 'upload-1', + key: 'assistant/org-1/user-1/upload-1/image.png', + name: 'image.png', + contentType: 'image/png', + size: 5, + buffer: Buffer.from('image'), + }) getEffectiveEnvironmentSnapshot.mockResolvedValue({ personalEncrypted: { API_KEY: 'encrypted-secret' }, workspaceEncrypted: {}, @@ -341,6 +355,130 @@ describe('handleUnifiedChatPost', () => { ) }) + it.each(['Describe this image', ''])( + 'prepares organization image bytes and persists canonical metadata (message: %s)', + async (message) => { + getSession.mockResolvedValue({ user: { id: 'user-1' }, session: { id: 'session-1' } }) + dbChainMockFns.returning.mockResolvedValueOnce([{ model: null }]) + const key = 'assistant/org-1/user-1/upload-1/image.png' + const response = await handleUnifiedChatPost( + new NextRequest('http://localhost/api/mothership/chat', { + method: 'POST', + body: JSON.stringify({ + message, + organizationId: 'org-1', + mode: 'assistant', + fileAttachments: [ + { id: 'forged-id', key, filename: 'forged.txt', media_type: 'text/plain', size: 0 }, + ], + }), + }) + ) + expect(response.status).toBe(200) + expect(readOrganizationAssistantImage).toHaveBeenCalledWith({ + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + organizationId: 'org-1', + key, + signal: expect.any(AbortSignal), + }) + expect(buildCopilotRequestPayload).toHaveBeenCalledWith( + expect.objectContaining({ + message, + assistantImages: [ + { + type: 'image', + filename: 'image.png', + source: { type: 'base64', media_type: 'image/png', data: 'aW1hZ2U=' }, + }, + ], + }), + expect.anything() + ) + expect(appendCopilotChatMessages).toHaveBeenCalledWith( + 'chat-1', + [ + expect.objectContaining({ + content: message, + fileAttachments: [ + { id: 'upload-1', key, filename: 'image.png', media_type: 'image/png', size: 5 }, + ], + }), + ], + expect.anything(), + expect.anything() + ) + expect(getUserEntityPermissions).not.toHaveBeenCalled() + expect(generateWorkspaceSnapshot).not.toHaveBeenCalled() + } + ) + + it('rejects inaccessible images before creating or persisting a conversation', async () => { + getSession.mockResolvedValue({ user: { id: 'user-1' }, session: { id: 'session-1' } }) + readOrganizationAssistantImage.mockRejectedValueOnce( + new OrchestrationError('not_found', 'Image not found') + ) + const response = await handleUnifiedChatPost( + new NextRequest('http://localhost/api/mothership/chat', { + method: 'POST', + body: JSON.stringify({ + message: '', + organizationId: 'org-1', + mode: 'assistant', + fileAttachments: [ + { + id: 'image', + key: 'other-user-image', + filename: 'image.png', + media_type: 'image/png', + size: 5, + }, + ], + }), + }) + ) + expect(response.status).toBe(403) + expect(resolveOrCreateChat).not.toHaveBeenCalled() + expect(appendCopilotChatMessages).not.toHaveBeenCalled() + expect(createSSEStream).not.toHaveBeenCalled() + }) + + it('continues rejecting empty messages without organization images', async () => { + const response = await handleUnifiedChatPost( + new NextRequest('http://localhost/api/mothership/chat', { + method: 'POST', + body: JSON.stringify({ message: '', organizationId: 'org-1', mode: 'assistant' }), + }) + ) + expect(response.status).toBe(400) + expect(readOrganizationAssistantImage).not.toHaveBeenCalled() + expect(resolveOrCreateChat).not.toHaveBeenCalled() + }) + + it('keeps workspace files unavailable in workspace Assistant mode', async () => { + const response = await handleUnifiedChatPost( + new NextRequest('http://localhost/api/mothership/chat', { + method: 'POST', + body: JSON.stringify({ + message: 'Read this file', + workspaceId: 'ws-1', + mode: 'assistant', + fileAttachments: [ + { + id: 'file-1', + key: 'workspace/file.png', + filename: 'file.png', + media_type: 'image/png', + size: 5, + }, + ], + }), + }) + ) + expect(response.status).toBe(400) + expect(readOrganizationAssistantImage).not.toHaveBeenCalled() + expect(resolveOrCreateChat).not.toHaveBeenCalled() + }) + it.each([{ workspaceId: 'ws-1' }, { workflowId: 'wf-1' }, { mode: 'agent' }])( 'rejects mixed organization scope before persistence: %j', async (extra) => { diff --git a/apps/sim/lib/copilot/chat/post.ts b/apps/sim/lib/copilot/chat/post.ts index ba2f04f7afe..77b1570cb8d 100644 --- a/apps/sim/lib/copilot/chat/post.ts +++ b/apps/sim/lib/copilot/chat/post.ts @@ -19,6 +19,10 @@ import { resolveOrganizationBillingAttribution, } from '@/lib/billing/core/billing-attribution' import { chatOperations } from '@/lib/copilot/application/operations' +import { + type AssistantImageContent, + prepareAssistantImages, +} from '@/lib/copilot/chat/assistant-images' import { DESKTOP_TERMINAL_HINT_ID_MAX_LENGTH, DESKTOP_TERMINAL_HINT_TEXT_MAX_LENGTH, @@ -277,63 +281,70 @@ const ChatContextSchema = z } }) -const ChatMessageSchema = z.object({ - message: z.string().min(1, 'Message is required'), - /* Bounded because it becomes part of a Postgres key in `chatSendIdempotency`; +const ChatMessageSchema = z + .object({ + message: z.string(), + /* Bounded because it becomes part of a Postgres key in `chatSendIdempotency`; a client-supplied id longer than the btree entry limit would throw there. A generated id is 36 chars. */ - userMessageId: z.string().max(128).optional(), - chatId: z.string().optional(), - workflowId: z.string().optional(), - workspaceId: z.string().optional(), - organizationId: z.string().min(1).max(200).optional(), - workflowName: z.string().optional(), - model: z.string().optional().default(DEFAULT_MODEL), - mode: z.enum(COPILOT_REQUEST_MODES).optional().default('agent'), - assistantSearch: workspaceSearchFiltersSchema.optional(), - prefetch: z.boolean().optional(), - createNewChat: z.boolean().optional().default(false), - implicitFeedback: z.string().optional(), - fileAttachments: z.array(FileAttachmentSchema).optional(), - resourceAttachments: z - .preprocess(dropUnaddressableAttachments, z.array(ResourceAttachmentSchema)) - .optional(), - provider: z.string().optional(), - contexts: z.array(ChatContextSchema).optional(), - commands: z.array(z.string()).optional(), - userTimezone: z.string().optional(), - desktopCapabilities: z - .object({ - localFilesystem: z.boolean().optional(), - browser: z.boolean().optional(), - terminal: z.boolean().optional(), - terminals: z - .array( - z.object({ - id: z.string().max(DESKTOP_TERMINAL_HINT_ID_MAX_LENGTH), - cwd: z.string().max(DESKTOP_TERMINAL_HINT_TEXT_MAX_LENGTH).optional(), - running: z.string().max(DESKTOP_TERMINAL_HINT_TEXT_MAX_LENGTH).optional(), - interactive: z.boolean().optional(), - active: z.boolean().optional(), - }) - ) - .optional(), - browserSessions: z - .array( - z.object({ - hostname: z - .string() - .max(253) - .regex(/^[a-z0-9.-]+$/), - evidence: z.enum(['sign-in-completed', 'cookies']), - lastObservedAt: z.string().datetime(), - }) - ) - .max(20) - .optional(), - }) - .optional(), -}) + userMessageId: z.string().max(128).optional(), + chatId: z.string().optional(), + workflowId: z.string().optional(), + workspaceId: z.string().optional(), + organizationId: z.string().min(1).max(200).optional(), + workflowName: z.string().optional(), + model: z.string().optional().default(DEFAULT_MODEL), + mode: z.enum(COPILOT_REQUEST_MODES).optional().default('agent'), + assistantSearch: workspaceSearchFiltersSchema.optional(), + prefetch: z.boolean().optional(), + createNewChat: z.boolean().optional().default(false), + implicitFeedback: z.string().optional(), + fileAttachments: z.array(FileAttachmentSchema).optional(), + resourceAttachments: z + .preprocess(dropUnaddressableAttachments, z.array(ResourceAttachmentSchema)) + .optional(), + provider: z.string().optional(), + contexts: z.array(ChatContextSchema).optional(), + commands: z.array(z.string()).optional(), + userTimezone: z.string().optional(), + desktopCapabilities: z + .object({ + localFilesystem: z.boolean().optional(), + browser: z.boolean().optional(), + terminal: z.boolean().optional(), + terminals: z + .array( + z.object({ + id: z.string().max(DESKTOP_TERMINAL_HINT_ID_MAX_LENGTH), + cwd: z.string().max(DESKTOP_TERMINAL_HINT_TEXT_MAX_LENGTH).optional(), + running: z.string().max(DESKTOP_TERMINAL_HINT_TEXT_MAX_LENGTH).optional(), + interactive: z.boolean().optional(), + active: z.boolean().optional(), + }) + ) + .optional(), + browserSessions: z + .array( + z.object({ + hostname: z + .string() + .max(253) + .regex(/^[a-z0-9.-]+$/), + evidence: z.enum(['sign-in-completed', 'cookies']), + lastObservedAt: z.string().datetime(), + }) + ) + .max(20) + .optional(), + }) + .optional(), + }) + .refine( + (body) => + body.message.length > 0 || + (body.mode === 'assistant' && !!body.organizationId && !!body.fileAttachments?.length), + { message: 'Message is required', path: ['message'] } + ) type UnifiedChatRequest = z.infer type BrowserSessions = NonNullable['browserSessions'] @@ -406,6 +417,7 @@ type UnifiedChatBranch = contexts: Array<{ type: string; content: string; tag?: string; path?: string }> mcpServerIds?: string[] fileAttachments?: UnifiedChatRequest['fileAttachments'] + assistantImages?: AssistantImageContent[] userPermission?: string entitlements?: string[] userTimezone?: string @@ -1138,7 +1150,7 @@ export async function handleUnifiedChatPost(req: NextRequest) { body.mode === 'assistant' && (body.workflowId || body.workflowName || - body.fileAttachments?.length || + (body.fileAttachments?.length && !body.organizationId) || body.contexts?.length) ) { return createBadRequestResponse( @@ -1251,6 +1263,21 @@ export async function handleUnifiedChatPost(req: NextRequest) { return capabilityRefusalResponse(chatCapability) } + const assistantImages = + branch.kind === 'organization' && body.fileAttachments?.length + ? await prepareAssistantImages({ + principal: { + kind: 'session', + userId: authenticatedUserId, + sessionId: session.session.id, + }, + organizationId: branch.organizationId, + attachments: body.fileAttachments, + signal: req.signal, + }) + : undefined + const fileAttachments = assistantImages?.attachments ?? body.fileAttachments + /* Prompt content is captured only once the turn is going to run. Both calls are internally gated on OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, but the gate is on @@ -1478,7 +1505,7 @@ export async function handleUnifiedChatPost(req: NextRequest) { chatId: actualChatId, userMessageId, message: body.message, - fileAttachments: body.fileAttachments, + fileAttachments, contexts: normalizedContexts, workspaceId, notifyWorkspaceStatus: branch.notifyWorkspaceStatus, @@ -1543,7 +1570,7 @@ export async function handleUnifiedChatPost(req: NextRequest) { contexts: turnContexts, assistantSearch: body.mode === 'assistant' ? body.assistantSearch : undefined, mcpServerIds, - fileAttachments: body.fileAttachments, + fileAttachments, userPermission: userPermission ?? undefined, entitlements, userTimezone: body.userTimezone, @@ -1572,7 +1599,8 @@ export async function handleUnifiedChatPost(req: NextRequest) { contexts: turnContexts, assistantSearch: body.mode === 'assistant' ? body.assistantSearch : undefined, mcpServerIds, - fileAttachments: body.fileAttachments, + fileAttachments, + assistantImages: assistantImages?.content, userPermission: userPermission ?? undefined, entitlements, userTimezone: body.userTimezone, @@ -1702,6 +1730,12 @@ export async function handleUnifiedChatPost(req: NextRequest) { if (applicationError?.code === 'forbidden' || applicationError?.code === 'not_found') { return NextResponse.json({ error: 'Conversation access denied' }, { status: 403 }) } + if (applicationError?.code === 'validation' || applicationError?.code === 'payload_too_large') { + return NextResponse.json( + { error: applicationError.message }, + { status: applicationError.code === 'validation' ? 400 : 413 } + ) + } if (isWorkspaceAccessDeniedError(error)) { return NextResponse.json({ error: 'Workspace access denied' }, { status: 403 }) } diff --git a/apps/sim/lib/core/utils/browser-storage.ts b/apps/sim/lib/core/utils/browser-storage.ts index 32109dcfb58..7103bfa421a 100644 --- a/apps/sim/lib/core/utils/browser-storage.ts +++ b/apps/sim/lib/core/utils/browser-storage.ts @@ -362,22 +362,27 @@ export class MothershipHandoffStorage { * accumulate — "Add to chat" can fire twice before the route swap completes, * and the second write must not drop the first. * @returns True if stored, false when the workspace is empty or the handoff - * carries neither a message nor a context. + * carries no message, context, or attachment. */ static store(handoff: MothershipHandoff, owner: MothershipHandoffOwner): boolean { const workspaceId = typeof owner === 'string' ? owner : undefined const organizationId = typeof owner === 'string' ? undefined : owner.organizationId const message = handoff.message?.trim() + const hasAttachments = Boolean(handoff.fileAttachments?.length) const contexts = handoff.contexts ?? [] - if (!(workspaceId || organizationId) || (!message && contexts.length === 0)) { + if ( + !(workspaceId || organizationId) || + (!message && !hasAttachments && contexts.length === 0) + ) { return false } return BrowserStorage.setItem(MothershipHandoffStorage.KEY, { - ...(message ? { message } : {}), - contexts: message - ? contexts - : [...MothershipHandoffStorage.pendingContexts(owner), ...contexts], + ...(message || hasAttachments ? { message: message ?? '' } : {}), + contexts: + message || hasAttachments + ? contexts + : [...MothershipHandoffStorage.pendingContexts(owner), ...contexts], ...(handoff.fileAttachments?.length ? { fileAttachments: handoff.fileAttachments } : {}), ...(handoff.resumeUserMessageId ? { resumeUserMessageId: handoff.resumeUserMessageId } : {}), ...(handoff.requestMode ? { requestMode: handoff.requestMode } : {}), @@ -398,7 +403,13 @@ export class MothershipHandoffStorage { */ private static pendingContexts(owner: MothershipHandoffOwner): ChatContext[] { const data = BrowserStorage.getItem(MothershipHandoffStorage.KEY, null) - if (!data || data.message || !MothershipHandoffStorage.belongsTo(data, owner)) return [] + if ( + !data || + data.message || + data.fileAttachments?.length || + !MothershipHandoffStorage.belongsTo(data, owner) + ) + return [] if (!data.timestamp || Date.now() - data.timestamp > MothershipHandoffStorage.MAX_AGE_MS) { return [] } @@ -433,10 +444,11 @@ export class MothershipHandoffStorage { MothershipHandoffStorage.clear() const contexts = Array.isArray(data.contexts) ? data.contexts : [] + const hasAttachments = Array.isArray(data.fileAttachments) && data.fileAttachments.length > 0 if ( !(data.workspaceId || data.organizationId) || Boolean(data.workspaceId && data.organizationId) || - (!data.message && contexts.length === 0) || + (!data.message && !hasAttachments && contexts.length === 0) || !data.timestamp || Date.now() - data.timestamp > maxAge ) { @@ -447,7 +459,7 @@ export class MothershipHandoffStorage { if (!assistantSearch.success) return null return { - ...(data.message ? { message: data.message } : {}), + ...(data.message || hasAttachments ? { message: data.message ?? '' } : {}), contexts, ...(data.requestMode === 'assistant' ? { requestMode: 'assistant' as const } : {}), ...(data.assistantSearch ? { assistantSearch: assistantSearch.data } : {}), diff --git a/apps/sim/lib/mothership/events.ts b/apps/sim/lib/mothership/events.ts index 1ee0b53150c..0d3bdd13d65 100644 --- a/apps/sim/lib/mothership/events.ts +++ b/apps/sim/lib/mothership/events.ts @@ -61,7 +61,7 @@ export function sendMothershipMessage( assistantSearch?: WorkspaceSearchFilters ): boolean { const trimmed = message.trim() - if (!trimmed) { + if (!trimmed && !fileAttachments?.length) { logger.warn('sendMothershipMessage called with empty message') return false } diff --git a/apps/sim/lib/uploads/client/admission.ts b/apps/sim/lib/uploads/client/admission.ts index 670a570a2fd..489323c1c2a 100644 --- a/apps/sim/lib/uploads/client/admission.ts +++ b/apps/sim/lib/uploads/client/admission.ts @@ -34,6 +34,7 @@ interface UploadAdmissionFile { interface MultiFileUploadAdmissionOptions { existingFiles?: ArrayLike + maxFiles?: number maxFileBytes?: number maxTotalBytes?: number } @@ -48,9 +49,13 @@ export function assertMultiFileUploadAdmission( options: MultiFileUploadAdmissionOptions = {} ): void { const existingFiles = options.existingFiles + const maxFiles = options.maxFiles ?? MULTI_FILE_UPLOAD_MAX_FILES const maxFileBytes = options.maxFileBytes ?? MULTI_FILE_UPLOAD_MAX_FILE_BYTES const maxTotalBytes = options.maxTotalBytes ?? MULTI_FILE_UPLOAD_MAX_TOTAL_FILE_EQUIVALENTS * maxFileBytes + if (!Number.isSafeInteger(maxFiles) || maxFiles < 1) { + throw new Error('Invalid upload file count limit') + } if (!Number.isSafeInteger(maxFileBytes) || maxFileBytes < 1) { throw new Error('Invalid per-file upload limit') } @@ -59,9 +64,9 @@ export function assertMultiFileUploadAdmission( } const existingCount = existingFiles?.length ?? 0 const totalCount = existingCount + files.length - if (totalCount > MULTI_FILE_UPLOAD_MAX_FILES) { + if (totalCount > maxFiles) { throw new MultiFileUploadAdmissionError( - `Select up to ${MULTI_FILE_UPLOAD_MAX_FILES} files at a time.`, + `Select up to ${maxFiles} files at a time.`, 'UPLOAD_FILE_COUNT_EXCEEDED' ) } diff --git a/apps/sim/lib/uploads/client/session-upload.ts b/apps/sim/lib/uploads/client/session-upload.ts index 55d1a7466f6..1411a7bec50 100644 --- a/apps/sim/lib/uploads/client/session-upload.ts +++ b/apps/sim/lib/uploads/client/session-upload.ts @@ -39,7 +39,8 @@ type InternalUploadContext = | { purpose: 'workspace_file'; workspaceId: string; folderId?: string | null } | { purpose: 'profile_picture' } | { purpose: 'workspace_logo'; workspaceId: string } - | { purpose: 'mothership_attachment'; workspaceId: string } + | { purpose: 'mothership_attachment'; workspaceId: string; organizationId?: never } + | { purpose: 'mothership_attachment'; organizationId: string; workspaceId?: never } | { purpose: 'execution_attachment' workspaceId: string @@ -162,12 +163,19 @@ function internalUploadBody(params: UploadInternalFileSessionParams): CreateInte case 'profile_picture': return { purpose: params.purpose, ...fileFields } case 'workspace_logo': - case 'mothership_attachment': return { purpose: params.purpose, workspaceId: params.workspaceId, ...fileFields, } + case 'mothership_attachment': + return { + purpose: params.purpose, + ...(params.organizationId + ? { organizationId: params.organizationId } + : { workspaceId: params.workspaceId }), + ...fileFields, + } case 'execution_attachment': return { purpose: params.purpose, diff --git a/apps/sim/lib/uploads/contexts/organization-assistant/application.test.ts b/apps/sim/lib/uploads/contexts/organization-assistant/application.test.ts new file mode 100644 index 00000000000..dca45cbf243 --- /dev/null +++ b/apps/sim/lib/uploads/contexts/organization-assistant/application.test.ts @@ -0,0 +1,226 @@ +/** @vitest-environment node */ +import { dbChainMockFns, resetDbChainMock } from '@sim/testing' +import sharp from 'sharp' +import { beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' + +const mocks = vi.hoisted(() => ({ download: vi.fn(), config: vi.fn(), create: vi.fn() })) +vi.mock('@/lib/uploads/core/storage-service', () => ({ downloadFile: mocks.download })) +vi.mock('@/lib/uploads/upload-session/service', () => ({ createUploadSession: mocks.create })) +vi.mock('@/lib/permission-groups/resolve.server', () => ({ + getUserPermissionConfigForOrganization: mocks.config, +})) +vi.mock('@/lib/uploads/config', () => ({ getServeStoragePrefix: () => 's3' })) + +import { + authorizeOrganizationAttachmentControl, + createOrganizationAssistantAttachment, + finalizeOrganizationAssistantAttachment, + readOrganizationAssistantImage, +} from '@/lib/uploads/contexts/organization-assistant/application' +import { ASSISTANT_IMAGE_MAX_BYTES } from '@/lib/uploads/shared/assistant-images' +import type { UploadSessionRecord } from '@/lib/uploads/upload-session/service' + +const principal = { kind: 'session', userId: 'user-1', sessionId: 'session-1' } as const +const key = 'assistant/org-1/user-1/upload-1/image.png' +const session: UploadSessionRecord = { + id: 'upload-1', + purpose: 'mothership_attachment', + workspaceId: null, + userId: 'user-1', + metadata: { + organizationAttachment: { organizationId: 'org-1', userId: 'user-1', sessionId: 'session-1' }, + }, + finalKey: key, + storageKey: key, + fileName: 'image.png', + contentType: 'image/png', + fileSize: 100, + storageContext: 'mothership', + storageProvider: 's3', + status: 'completed', + method: 'put', + knowledgeBaseId: null, + workflowId: null, + executionId: null, + providerUploadId: null, + providerObjectVersion: 'v1', + partSize: null, + partCount: null, + uploadToken: '', + createdAt: new Date(), + expiresAt: new Date(), + completedFileId: null, + error: null, + completedAt: new Date(), + updatedAt: new Date(), +} +let png: Buffer +beforeAll(async () => { + png = await sharp({ create: { width: 4, height: 4, channels: 3, background: '#ff0000' } }) + .png() + .toBuffer() +}) +beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() + mocks.config.mockResolvedValue(null) + mocks.download.mockResolvedValue(png) + dbChainMockFns.limit.mockResolvedValue([{ role: 'member' }]) +}) + +function read(overrides: Partial[0]> = {}) { + return readOrganizationAssistantImage({ principal, organizationId: 'org-1', key, ...overrides }) +} + +describe('private organization Assistant images', () => { + it('creates uploads as the actual current member with no workspace fallback', async () => { + await createOrganizationAssistantAttachment(principal, { + organizationId: 'org-1', + name: 'image.png', + contentType: 'image/png', + size: 100, + localOrigin: 'http://localhost', + }) + expect(mocks.create).toHaveBeenCalledWith( + expect.objectContaining({ + principal, + userId: 'user-1', + organizationId: 'org-1', + purpose: 'mothership_attachment', + }) + ) + expect(mocks.create.mock.calls[0][0]).not.toHaveProperty('workspaceId') + }) + + it('rejects removed members before creating an upload', async () => { + dbChainMockFns.limit.mockResolvedValue([]) + await expect( + createOrganizationAssistantAttachment(principal, { + organizationId: 'org-1', + name: 'image.png', + contentType: 'image/png', + size: 100, + localOrigin: 'http://localhost', + }) + ).rejects.toMatchObject({ code: 'not_found' }) + expect(mocks.create).not.toHaveBeenCalled() + }) + + it('reads canonical completed uploads after a new login and emits bounded decoded bytes', async () => { + dbChainMockFns.limit.mockResolvedValueOnce([session]) + const signal = new AbortController().signal + const image = await read({ principal: { ...principal, sessionId: 'new-session' }, signal }) + expect(image).toMatchObject({ + id: 'upload-1', + key, + name: 'image.png', + contentType: 'image/webp', + }) + expect((await sharp(image.buffer).metadata()).format).toBe('webp') + expect(mocks.download).toHaveBeenCalledWith({ + key, + context: 'mothership', + maxBytes: ASSISTANT_IMAGE_MAX_BYTES, + signal, + }) + }) + + it.each([ + { key: 'https://example.com/image.png' }, + { key: 'assistant/org-1/user-1/../image.png' }, + { principal: { kind: 'personal_api_key', userId: 'user-1', keyId: 'key-1' } as const }, + ])('rejects invalid references or non-session callers before loading', async (input) => { + await expect(read(input)).rejects.toMatchObject({ code: 'not_found' }) + expect(dbChainMockFns.limit).not.toHaveBeenCalled() + expect(mocks.download).not.toHaveBeenCalled() + }) + + it.each([ + { organizationId: 'other-org' }, + { principal: { ...principal, userId: 'other-user' } }, + { key: 'assistant/other-org/user-1/upload-1/image.png' }, + ])('rejects a mismatched asserted owner', async (input) => { + dbChainMockFns.limit.mockResolvedValueOnce([session]) + await expect(read(input)).rejects.toMatchObject({ code: 'not_found' }) + expect(mocks.download).not.toHaveBeenCalled() + }) + + it('refuses absent, incomplete, or purged records', async () => { + dbChainMockFns.limit.mockResolvedValueOnce([]) + await expect(read()).rejects.toMatchObject({ code: 'not_found' }) + expect(mocks.download).not.toHaveBeenCalled() + }) + + it('rechecks membership for every preview/model read', async () => { + dbChainMockFns.limit.mockResolvedValueOnce([session]).mockResolvedValueOnce([]) + await expect(read()).rejects.toMatchObject({ code: 'not_found' }) + expect(mocks.download).not.toHaveBeenCalled() + }) + + it('rejects missing immutable scope metadata', async () => { + dbChainMockFns.limit.mockResolvedValueOnce([{ ...session, metadata: {} }]) + await expect(read()).rejects.toMatchObject({ code: 'not_found' }) + expect(mocks.download).not.toHaveBeenCalled() + }) + + it('refuses metadata above the byte cap without downloading', async () => { + dbChainMockFns.limit.mockResolvedValueOnce([ + { ...session, fileSize: ASSISTANT_IMAGE_MAX_BYTES + 1 }, + ]) + await expect(read()).rejects.toMatchObject({ code: 'payload_too_large' }) + expect(mocks.download).not.toHaveBeenCalled() + }) + + it.each([ + '', + '', + ])('rejects active content with a forged image MIME', async (content) => { + dbChainMockFns.limit.mockResolvedValueOnce([session]) + mocks.download.mockResolvedValue(Buffer.from(content)) + await expect(read()).rejects.toMatchObject({ code: 'validation' }) + }) + + it('propagates storage infrastructure failures unchanged', async () => { + const error = new Error('storage unavailable') + dbChainMockFns.limit.mockResolvedValueOnce([session]) + mocks.download.mockRejectedValue(error) + await expect(read()).rejects.toBe(error) + }) + + it('rejects compressed images above the 25 megapixel decode budget', async () => { + const largePng = await sharp({ + create: { width: 5001, height: 5000, channels: 3, background: '#000' }, + }) + .png() + .toBuffer() + expect(largePng.length).toBeLessThan(ASSISTANT_IMAGE_MAX_BYTES) + dbChainMockFns.limit.mockResolvedValueOnce([session]) + mocks.download.mockResolvedValue(largePng) + await expect(read()).rejects.toMatchObject({ code: 'validation', cause: expect.any(Error) }) + }) + + it('binds upload controls to the exact creating session', async () => { + await expect( + authorizeOrganizationAttachmentControl({ ...principal, sessionId: 'other-session' }, session) + ).rejects.toMatchObject({ code: 'not_found' }) + expect(dbChainMockFns.limit).not.toHaveBeenCalled() + }) + + it('reauthorizes finalization after decoding before returning an attachment', async () => { + dbChainMockFns.limit.mockResolvedValueOnce([{ role: 'member' }]).mockResolvedValueOnce([]) + await expect(finalizeOrganizationAssistantAttachment(principal, session)).rejects.toMatchObject( + { code: 'not_found' } + ) + expect(mocks.download).toHaveBeenCalledTimes(1) + }) + + it('returns the same durable metadata on completion replay', async () => { + const first = await finalizeOrganizationAssistantAttachment(principal, session) + const second = await finalizeOrganizationAssistantAttachment(principal, session) + expect(second).toEqual(first) + expect(first).toMatchObject({ + key, + path: `/api/files/serve/s3/${encodeURIComponent(key)}?context=mothership`, + }) + }) +}) diff --git a/apps/sim/lib/uploads/contexts/organization-assistant/application.ts b/apps/sim/lib/uploads/contexts/organization-assistant/application.ts new file mode 100644 index 00000000000..ea894123634 --- /dev/null +++ b/apps/sim/lib/uploads/contexts/organization-assistant/application.ts @@ -0,0 +1,181 @@ +import type { Principal } from '@sim/auth/principal' +import { db } from '@sim/db' +import { uploadSession } from '@sim/db/schema' +import { and, eq, isNull } from 'drizzle-orm' +import sharp from 'sharp' +import { authorizeOrganizationOperation } from '@/lib/core/application/organization-authorization' +import { defineOrganizationOperation } from '@/lib/core/application/organization-operation' +import { OrchestrationError } from '@/lib/core/orchestration/types' +import { getServeStoragePrefix } from '@/lib/uploads/config' +import { + assertOrganizationAttachmentControlBinding, + organizationAttachmentBinding, +} from '@/lib/uploads/contexts/organization-assistant/binding' +import { downloadFile } from '@/lib/uploads/core/storage-service' +import { + ASSISTANT_IMAGE_MAX_BYTES, + isAssistantImageType, +} from '@/lib/uploads/shared/assistant-images' +import { createUploadSession, type UploadSessionRecord } from '@/lib/uploads/upload-session/service' + +const organizationAttachmentOperation = defineOrganizationOperation({ + id: 'organization.assistant.attachments.use', + minimumRole: 'member', + principalKinds: ['session'], + capability: 'copilot.use', +}) + +const MAX_ASSISTANT_IMAGE_PIXELS = 25_000_000 + +export interface CreateOrganizationAssistantAttachmentInput { + organizationId: string + name: string + contentType: string + size: number + localOrigin: string +} + +export async function createOrganizationAssistantAttachment( + principal: Principal, + input: CreateOrganizationAssistantAttachmentInput +) { + const context = await authorizeOrganizationOperation( + principal, + organizationAttachmentOperation, + input + ) + return createUploadSession({ + purpose: 'mothership_attachment', + principal, + organizationId: context.organizationId, + userId: context.userId, + fileName: input.name, + contentType: input.contentType, + fileSize: input.size, + localOrigin: input.localOrigin, + }) +} + +export async function authorizeOrganizationAttachmentControl( + principal: Principal, + session: UploadSessionRecord +): Promise { + const binding = assertOrganizationAttachmentControlBinding(session, principal) + await authorizeOrganizationOperation(principal, organizationAttachmentOperation, binding) +} + +/** A bounded decode removes active content, metadata, and animation before preview or model use. */ +async function readImageBytes(key: string, contentType: string, signal?: AbortSignal) { + if (!isAssistantImageType(contentType)) + throw new OrchestrationError('validation', 'Unsupported image type') + const buffer = await downloadFile({ + key, + context: 'mothership', + maxBytes: ASSISTANT_IMAGE_MAX_BYTES, + signal, + }) + try { + const image = sharp(buffer, { limitInputPixels: MAX_ASSISTANT_IMAGE_PIXELS, pages: 1 }) + const metadata = await image.metadata() + if (!metadata.format || !['jpeg', 'png', 'gif', 'webp'].includes(metadata.format)) { + throw new OrchestrationError( + 'validation', + 'Attachment must contain a PNG, JPEG, GIF, or WebP image' + ) + } + const normalized = await image + .rotate() + .resize(1568, 1568, { fit: 'inside', withoutEnlargement: true }) + .webp({ quality: 85 }) + .toBuffer() + if (normalized.length > ASSISTANT_IMAGE_MAX_BYTES) + throw new OrchestrationError('payload_too_large', 'Image exceeds the 5 MB limit') + return normalized + } catch (cause) { + if (cause instanceof OrchestrationError) throw cause + const error = new OrchestrationError('validation', 'Attachment is not a valid supported image') + error.cause = cause + throw error + } +} + +export async function finalizeOrganizationAssistantAttachment( + principal: Principal, + session: UploadSessionRecord +) { + await authorizeOrganizationAttachmentControl(principal, session) + await readImageBytes(session.finalKey, session.contentType) + await authorizeOrganizationAttachmentControl(principal, session) + return { + path: `/api/files/serve/${getServeStoragePrefix()}/${encodeURIComponent(session.finalKey)}?context=mothership`, + key: session.finalKey, + name: session.fileName, + size: session.fileSize, + type: session.contentType, + } +} + +/** Resolves only completed images owned by the current user and their current organization. */ +export async function readOrganizationAssistantImage(input: { + principal: Principal + organizationId?: string + key: string + signal?: AbortSignal +}) { + if (input.principal.kind !== 'session') + throw new OrchestrationError('not_found', 'Attachment not found') + const keyParts = input.key.split('/') + if ( + keyParts.length !== 5 || + keyParts[0] !== 'assistant' || + keyParts.some((part) => !part || part === '.' || part === '..') + ) { + throw new OrchestrationError('not_found', 'Attachment not found') + } + const [session] = await db + .select({ + id: uploadSession.id, + purpose: uploadSession.purpose, + workspaceId: uploadSession.workspaceId, + userId: uploadSession.userId, + metadata: uploadSession.metadata, + fileName: uploadSession.fileName, + contentType: uploadSession.contentType, + fileSize: uploadSession.fileSize, + finalKey: uploadSession.finalKey, + }) + .from(uploadSession) + .where( + and( + eq(uploadSession.id, keyParts[3]), + eq(uploadSession.finalKey, input.key), + eq(uploadSession.userId, input.principal.userId), + eq(uploadSession.purpose, 'mothership_attachment'), + eq(uploadSession.status, 'completed'), + isNull(uploadSession.workspaceId) + ) + ) + .limit(1) + if (!session) throw new OrchestrationError('not_found', 'Attachment not found') + const binding = organizationAttachmentBinding(session) + if ( + session.userId !== input.principal.userId || + keyParts[1] !== binding.organizationId || + keyParts[2] !== binding.userId || + (input.organizationId && input.organizationId !== binding.organizationId) + ) { + throw new OrchestrationError('not_found', 'Attachment not found') + } + await authorizeOrganizationOperation(input.principal, organizationAttachmentOperation, binding) + if (session.fileSize > ASSISTANT_IMAGE_MAX_BYTES) + throw new OrchestrationError('payload_too_large', 'Image exceeds the 5 MB limit') + const buffer = await readImageBytes(session.finalKey, session.contentType, input.signal) + return { + id: session.id, + key: session.finalKey, + name: session.fileName, + size: buffer.length, + contentType: 'image/webp', + buffer, + } +} diff --git a/apps/sim/lib/uploads/contexts/organization-assistant/binding.ts b/apps/sim/lib/uploads/contexts/organization-assistant/binding.ts new file mode 100644 index 00000000000..df2743c533c --- /dev/null +++ b/apps/sim/lib/uploads/contexts/organization-assistant/binding.ts @@ -0,0 +1,56 @@ +import type { Principal } from '@sim/auth/principal' +import { isRecordLike } from '@sim/utils/object' +import { OrchestrationError } from '@/lib/core/orchestration/types' + +export interface OrganizationAttachmentBinding { + organizationId: string + userId: string + sessionId: string +} + +interface OrganizationAttachmentSession { + purpose: string + workspaceId: string | null + userId: string + metadata: Record +} + +/** Organization attachments remain private to their uploader across new login sessions. */ +export function organizationAttachmentBinding( + session: OrganizationAttachmentSession +): OrganizationAttachmentBinding { + const binding = session.metadata.organizationAttachment + if ( + session.purpose !== 'mothership_attachment' || + session.workspaceId !== null || + !isRecordLike(binding) || + typeof binding.organizationId !== 'string' || + !binding.organizationId || + binding.userId !== session.userId || + typeof binding.sessionId !== 'string' || + !binding.sessionId + ) { + throw new OrchestrationError('not_found', 'Attachment not found') + } + return { + organizationId: binding.organizationId, + userId: session.userId, + sessionId: binding.sessionId, + } +} + +/** A byte-transfer token cannot replace the session that initiated the upload. */ +export function assertOrganizationAttachmentControlBinding( + session: OrganizationAttachmentSession, + principal: Principal +): OrganizationAttachmentBinding { + const binding = organizationAttachmentBinding(session) + if ( + principal.kind !== 'session' || + principal.userId !== binding.userId || + principal.sessionId !== binding.sessionId + ) { + throw new OrchestrationError('not_found', 'Upload session not found') + } + return binding +} diff --git a/apps/sim/lib/uploads/shared/assistant-images.ts b/apps/sim/lib/uploads/shared/assistant-images.ts new file mode 100644 index 00000000000..deae05e9a5d --- /dev/null +++ b/apps/sim/lib/uploads/shared/assistant-images.ts @@ -0,0 +1,15 @@ +/** Inline image limits shared by Assistant upload, preview, and model preparation. */ +export const ASSISTANT_IMAGE_MAX_BYTES = 5 * 1024 * 1024 +export const ASSISTANT_IMAGE_MAX_COUNT = 5 +export const ASSISTANT_IMAGE_MAX_TOTAL_BYTES = ASSISTANT_IMAGE_MAX_BYTES * ASSISTANT_IMAGE_MAX_COUNT +export const ASSISTANT_IMAGE_CONTENT_TYPES = [ + 'image/jpeg', + 'image/png', + 'image/gif', + 'image/webp', +] as const +export const ASSISTANT_IMAGE_ACCEPT_ATTRIBUTE = ASSISTANT_IMAGE_CONTENT_TYPES.join(',') + +export function isAssistantImageType(contentType: string): boolean { + return ASSISTANT_IMAGE_CONTENT_TYPES.some((type) => type === contentType) +} diff --git a/apps/sim/lib/uploads/upload-session/application.test.ts b/apps/sim/lib/uploads/upload-session/application.test.ts index 70ae6b01c61..52f9cd31ea6 100644 --- a/apps/sim/lib/uploads/upload-session/application.test.ts +++ b/apps/sim/lib/uploads/upload-session/application.test.ts @@ -12,6 +12,12 @@ const mocks = vi.hoisted(() => ({ getPrincipalSession: vi.fn(), reauthorizeWorkspacePurpose: vi.fn(), getWorkspaceFile: vi.fn(), + authorizeOrganizationAttachment: vi.fn(), +})) + +vi.mock('@/lib/uploads/contexts/organization-assistant/application', () => ({ + authorizeOrganizationAttachmentControl: mocks.authorizeOrganizationAttachment, + createOrganizationAssistantAttachment: vi.fn(), })) vi.mock('@/lib/uploads/contexts/workspace', () => ({ @@ -43,7 +49,9 @@ vi.mock('@/app/api/files/uploads/purposes', () => ({ })) import { + abortInternalUploadSession, completeInternalUploadSession, + issueInternalUploadPartUrls, readWorkspaceUploadSession, } from '@/lib/uploads/upload-session/application' import type { UploadSessionRecord } from '@/lib/uploads/upload-session/service' @@ -58,6 +66,7 @@ const actor = { id: 'user-1', name: 'Ada', email: 'ada@example.com' } describe('upload session application', () => { beforeEach(() => { vi.clearAllMocks() + mocks.authorizeOrganizationAttachment.mockResolvedValue(undefined) const session = workspaceUploadSession() mocks.getOwnedSession.mockResolvedValue(session) mocks.finalizePurpose.mockResolvedValue({ @@ -91,6 +100,48 @@ describe('upload session application', () => { ) }) + it.each(['complete', 'abort', 'parts'] as const)( + 'rechecks organization membership before the %s control leg', + async (control) => { + const session = { + ...workspaceUploadSession(), + purpose: 'mothership_attachment' as const, + workspaceId: null, + } + mocks.getOwnedSession.mockResolvedValue(session) + const request = new NextRequest('http://localhost/api/files/uploads/upload-1/complete', { + headers: { host: 'localhost' }, + }) + const input = { uploadId: 'upload-1', uploadToken: 'upload-token', partNumbers: [1] } + if (control === 'complete') await completeInternalUploadSession(principal, input, request) + else if (control === 'abort') await abortInternalUploadSession(principal, input) + else await issueInternalUploadPartUrls(principal, input, request) + expect(mocks.assertAuthBinding).toHaveBeenCalledWith(session, principal) + expect(mocks.authorizeOrganizationAttachment).toHaveBeenCalledWith(principal, session) + expect(mocks.reauthorizeWorkspacePurpose).not.toHaveBeenCalled() + } + ) + + it('does not finalize when organization access is revoked after the session is claimed', async () => { + const session = { + ...workspaceUploadSession(), + purpose: 'mothership_attachment' as const, + workspaceId: null, + } + mocks.getOwnedSession.mockResolvedValue(session) + mocks.authorizeOrganizationAttachment + .mockResolvedValueOnce(undefined) + .mockRejectedValueOnce(new Error('Organization not found')) + await expect( + completeInternalUploadSession( + principal, + { uploadId: 'upload-1', uploadToken: 'upload-token' }, + new NextRequest('http://localhost/api/files/uploads/upload-1/complete') + ) + ).rejects.toThrow('Organization not found') + expect(mocks.finalizePurpose).not.toHaveBeenCalled() + }) + /** * The read is a control leg, so it re-authorizes the caller's present * workspace permission rather than trusting the session lookup alone. @@ -119,7 +170,11 @@ describe('upload session application', () => { * had created. */ it('returns the registered file once the session has completed', async () => { - const completed = { ...workspaceUploadSession(), status: 'completed' as const, completedFileId: 'file-1' } + const completed = { + ...workspaceUploadSession(), + status: 'completed' as const, + completedFileId: 'file-1', + } mocks.getOwnedSession.mockResolvedValue(completed) mocks.getPrincipalSession.mockResolvedValue(completed) mocks.getWorkspaceFile.mockResolvedValue({ id: 'file-1', name: 'file.txt' }) @@ -154,7 +209,11 @@ describe('upload session application', () => { * there was nothing. A failed read is not the same answer as no file. */ it('surfaces a failed file read instead of reporting the upload fileless', async () => { - const completed = { ...workspaceUploadSession(), status: 'completed' as const, completedFileId: 'file-1' } + const completed = { + ...workspaceUploadSession(), + status: 'completed' as const, + completedFileId: 'file-1', + } mocks.getOwnedSession.mockResolvedValue(completed) mocks.getPrincipalSession.mockResolvedValue(completed) mocks.getWorkspaceFile.mockRejectedValue(new Error('connection terminated')) @@ -169,7 +228,11 @@ describe('upload session application', () => { }) it('reads the completed file with throwOnError so a fault cannot read as absence', async () => { - const completed = { ...workspaceUploadSession(), status: 'completed' as const, completedFileId: 'file-1' } + const completed = { + ...workspaceUploadSession(), + status: 'completed' as const, + completedFileId: 'file-1', + } mocks.getOwnedSession.mockResolvedValue(completed) mocks.getPrincipalSession.mockResolvedValue(completed) mocks.getWorkspaceFile.mockResolvedValue({ id: 'file-1', name: 'file.txt' }) @@ -187,7 +250,11 @@ describe('upload session application', () => { /** A completed session whose file was since deleted has nothing to address. */ it('answers null when the completed file is gone', async () => { - const gone = { ...workspaceUploadSession(), status: 'completed' as const, completedFileId: 'file-1' } + const gone = { + ...workspaceUploadSession(), + status: 'completed' as const, + completedFileId: 'file-1', + } mocks.getOwnedSession.mockResolvedValue(gone) mocks.getPrincipalSession.mockResolvedValue(gone) mocks.getWorkspaceFile.mockResolvedValue(null) diff --git a/apps/sim/lib/uploads/upload-session/application.ts b/apps/sim/lib/uploads/upload-session/application.ts index 61929b1dd77..b7060d8d7b1 100644 --- a/apps/sim/lib/uploads/upload-session/application.ts +++ b/apps/sim/lib/uploads/upload-session/application.ts @@ -3,6 +3,10 @@ import type { CreateInternalFileUploadBody } from '@/lib/api/contracts/upload-se import type { OrchestrationRequestContext } from '@/lib/core/orchestration/types' import { OrchestrationError } from '@/lib/core/orchestration/types' import { loadActiveFolderPathIndex, resolveFolderPathFromIndex } from '@/lib/folders/queries' +import { + authorizeOrganizationAttachmentControl, + createOrganizationAssistantAttachment, +} from '@/lib/uploads/contexts/organization-assistant/application' import { getWorkspaceFile, type WorkspaceFileRecord } from '@/lib/uploads/contexts/workspace' import { abortUploadSession, @@ -82,6 +86,13 @@ export async function createInternalPurposeUploadSession( body: CreateInternalFileUploadBody, request: OrchestrationRequestContext ): Promise>> { + if (body.purpose === 'mothership_attachment' && body.organizationId) { + return createOrganizationAssistantAttachment(principal, { + ...body, + organizationId: body.organizationId, + localOrigin: requestOrigin(request), + }) + } return createPurposeUploadSession(principal, body, requestOrigin(request)) } @@ -97,6 +108,9 @@ export async function loadAuthorizedInternalUploadSession( userId: principalUserId(principal), }) if (session.purpose === 'workspace_file') assertUploadSessionAuthBinding(session, principal) + if (session.purpose === 'mothership_attachment' && session.workspaceId === null) { + assertUploadSessionAuthBinding(session, principal) + } return session } @@ -108,6 +122,8 @@ export async function issueInternalUploadPartUrls( const session = await loadAuthorizedInternalUploadSession(principal, input) if (session.purpose === 'workspace_file') { await reauthorizeWorkspaceUploadPurpose(principal, session, fileOperations.uploadParts) + } else if (session.purpose === 'mothership_attachment' && session.workspaceId === null) { + await authorizeOrganizationAttachmentControl(principal, session) } else { await reauthorizeUploadPurpose(principalUserId(principal), session) } @@ -127,6 +143,8 @@ export async function abortInternalUploadSession( const session = await loadAuthorizedInternalUploadSession(principal, input) if (session.purpose === 'workspace_file') { await reauthorizeWorkspaceUploadPurpose(principal, session, fileOperations.uploadCancel) + } else if (session.purpose === 'mothership_attachment' && session.workspaceId === null) { + await authorizeOrganizationAttachmentControl(principal, session) } else { await reauthorizeUploadPurpose(principalUserId(principal), session) } @@ -146,6 +164,8 @@ export async function completeInternalUploadSession( const authorize = async (claimed: UploadSessionRecord) => { if (claimed.purpose === 'workspace_file') { await reauthorizeWorkspaceUploadPurpose(principal, claimed, fileOperations.uploadComplete) + } else if (claimed.purpose === 'mothership_attachment' && claimed.workspaceId === null) { + await authorizeOrganizationAttachmentControl(principal, claimed) } else { await reauthorizeUploadPurpose(principalUserId(principal), claimed) } diff --git a/apps/sim/lib/uploads/upload-session/service.test.ts b/apps/sim/lib/uploads/upload-session/service.test.ts index 558b1fca924..e8fedafbc12 100644 --- a/apps/sim/lib/uploads/upload-session/service.test.ts +++ b/apps/sim/lib/uploads/upload-session/service.test.ts @@ -153,6 +153,64 @@ describe('upload sessions', () => { }) }) + it('binds organization images to the creating session and stores them without a workspace', async () => { + dbChainMockFns.returning.mockResolvedValueOnce([ + uploadRow({ + purpose: 'mothership_attachment', + workspaceId: null, + storageContext: 'mothership', + finalKey: 'assistant/org-1/user-1/upload-1/image.png', + contentType: 'image/png', + }), + ]) + await createUploadSession({ + id: 'upload-1', + userId: 'user-1', + purpose: 'mothership_attachment', + organizationId: 'org-1', + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + fileName: 'image.png', + contentType: 'image/png', + fileSize: 100, + metadata: { organizationAttachment: { organizationId: 'forged' } }, + }) + expect(dbChainMockFns.values).toHaveBeenCalledWith( + expect.objectContaining({ + workspaceId: null, + finalKey: 'assistant/org-1/user-1/upload-1/image.png', + metadata: { + organizationAttachment: { + organizationId: 'org-1', + userId: 'user-1', + sessionId: 'session-1', + }, + }, + }) + ) + expect(mockCreatePutTransfer).toHaveBeenCalledWith( + expect.objectContaining({ context: 'mothership', fileSize: 100 }) + ) + }) + + it.each([ + { contentType: 'text/html', fileSize: 100 }, + { contentType: 'image/png', fileSize: 5 * 1024 * 1024 + 1 }, + ])('rejects invalid organization images before storage initialization', async (file) => { + await expect( + createUploadSession({ + id: 'upload-1', + userId: 'user-1', + purpose: 'mothership_attachment', + organizationId: 'org-1', + principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, + fileName: 'image.png', + ...file, + }) + ).rejects.toThrow('Assistant attachments must be') + expect(mockCreatePutTransfer).not.toHaveBeenCalled() + expect(dbChainMockFns.values).not.toHaveBeenCalled() + }) + // Local storage stores an object's metadata sidecar beside it, under the // object's own name, so the whole key + suffix must fit one path component. // Three purposes built their key by hand and admitted a 255-character name diff --git a/apps/sim/lib/uploads/upload-session/service.ts b/apps/sim/lib/uploads/upload-session/service.ts index 2c9aed76a7c..af5c7f8358c 100644 --- a/apps/sim/lib/uploads/upload-session/service.ts +++ b/apps/sim/lib/uploads/upload-session/service.ts @@ -11,7 +11,7 @@ import { sha256Hex } from '@sim/security/hash' import { generateSecureToken } from '@sim/security/tokens' import { getErrorMessage } from '@sim/utils/errors' import { generateId } from '@sim/utils/id' -import { and, asc, eq, inArray, isNull, lt, or } from 'drizzle-orm' +import { and, asc, eq, inArray, isNull, lt, or, sql } from 'drizzle-orm' import { checkStorageQuotaForBillingContext, resolveStorageBillingContext, @@ -19,8 +19,13 @@ import { import { OrchestrationError } from '@/lib/core/orchestration/types' import { generateUniqueExecutionFileKey } from '@/lib/uploads/contexts/execution/utils' import { generateKnowledgeBaseFileKey } from '@/lib/uploads/contexts/knowledge-base/knowledge-base-file-manager' +import { assertOrganizationAttachmentControlBinding } from '@/lib/uploads/contexts/organization-assistant/binding' import { generateWorkspaceFileKey } from '@/lib/uploads/contexts/workspace' import { buildStorageKeySegment } from '@/lib/uploads/core/storage-key' +import { + ASSISTANT_IMAGE_MAX_BYTES, + isAssistantImageType, +} from '@/lib/uploads/shared/assistant-images' import { MAX_KNOWLEDGE_DOCUMENT_FILE_SIZE, MAX_WORKSPACE_FILE_SIZE, @@ -190,6 +195,12 @@ export type CreateUploadSessionParams = CreateUploadSessionBaseParams & } | { purpose: 'profile_picture'; workspaceId?: null } | { purpose: 'workspace_logo' | 'mothership_attachment'; workspaceId: string } + | { + purpose: 'mothership_attachment' + organizationId: string + principal: Principal + workspaceId?: never + } | { purpose: 'execution_attachment' workspaceId: string @@ -206,8 +217,21 @@ export async function createUploadSession( validateFile(params) const id = params.id ?? generateId() const uploadToken = generateSecureToken(32) - const workspaceId = params.purpose === 'profile_picture' ? null : params.workspaceId + const workspaceId = params.purpose === 'profile_picture' ? null : (params.workspaceId ?? null) const metadata = { ...(params.metadata ?? {}) } + if (params.purpose === 'mothership_attachment' && 'organizationId' in params) { + if (params.principal.kind !== 'session' || params.principal.userId !== params.userId) { + throw new UploadSessionError( + 'forbidden', + 'Organization attachments require the uploading session' + ) + } + metadata.organizationAttachment = { + organizationId: params.organizationId, + userId: params.userId, + sessionId: params.principal.sessionId, + } + } if (params.purpose === 'workspace_file' || params.purpose === 'knowledge_document') { if (!workspaceId) throw new Error(`${params.purpose} upload is missing workspaceId`) if (!params.principal) { @@ -500,6 +524,10 @@ export function assertUploadSessionAuthBinding( session: UploadSessionRecord, principal: Principal ): void { + if (session.purpose === 'mothership_attachment' && session.workspaceId === null) { + assertOrganizationAttachmentControlBinding(session, principal) + return + } if (!isPrincipalBoundUploadPurpose(session.purpose)) return const candidate = session.metadata.authBinding if (candidate === undefined) { @@ -913,6 +941,8 @@ export async function cleanupExpiredUploadSessions(): Promise<{ .where( and( inArray(uploadSession.status, ['completed', 'aborted', 'expired']), + /** Completed organization sessions are the durable private-image ownership records. */ + sql`NOT (${uploadSession.status} = 'completed' AND ${uploadSession.purpose} = 'mothership_attachment' AND ${uploadSession.workspaceId} IS NULL)`, lt(uploadSession.completedAt, terminalCutoff), or( isNull(uploadSession.processingLeaseId), @@ -1202,7 +1232,24 @@ function validateFile(params: CreateUploadSessionParams): void { if (params.fileSize > maximum) { throw new UploadSessionError('validation', `File size exceeds maximum of ${maximum} bytes`) } - if (params.purpose !== 'profile_picture' && !params.workspaceId.trim()) { + const organizationAttachment = + params.purpose === 'mothership_attachment' && 'organizationId' in params + if ( + organizationAttachment && + (!params.organizationId.trim() || + params.fileSize > ASSISTANT_IMAGE_MAX_BYTES || + !isAssistantImageType(params.contentType)) + ) { + throw new UploadSessionError( + 'validation', + 'Assistant attachments must be PNG, JPEG, GIF, or WebP images up to 5 MB' + ) + } + if ( + params.purpose !== 'profile_picture' && + !organizationAttachment && + !params.workspaceId?.trim() + ) { throw new UploadSessionError('validation', 'workspaceId must not be empty') } if (params.purpose === 'knowledge_document' && !params.knowledgeBaseId.trim()) { @@ -1231,7 +1278,10 @@ function requiresStorageQuota(purpose: UploadSessionPurpose): boolean { function isPrincipalBoundUploadPurpose(purpose: UploadSessionPurpose): boolean { return ( - purpose === 'workspace_file' || purpose === 'knowledge_document' || purpose === 'table_import' + purpose === 'workspace_file' || + purpose === 'knowledge_document' || + purpose === 'table_import' || + purpose === 'mothership_attachment' ) } @@ -1266,6 +1316,12 @@ function resolveUploadStorage( finalKey: `workspace-logos/${params.workspaceId}/${buildStorageKeySegment(`${id}-`, params.fileName)}`, } case 'mothership_attachment': + if ('organizationId' in params) { + return { + storageContext: 'mothership', + finalKey: `assistant/${params.organizationId}/${params.userId}/${id}/${buildStorageKeySegment('', params.fileName)}`, + } + } return { storageContext: 'mothership', finalKey: generateWorkspaceFileKey(params.workspaceId, params.fileName), diff --git a/apps/sim/lib/uploads/utils/file-utils.ts b/apps/sim/lib/uploads/utils/file-utils.ts index b00ac1ab7ee..a06539b5611 100644 --- a/apps/sim/lib/uploads/utils/file-utils.ts +++ b/apps/sim/lib/uploads/utils/file-utils.ts @@ -801,6 +801,7 @@ export function tryInferContextFromKey(key: string): StorageContext | null { if (key.startsWith('copilot/')) return 'copilot' if (key.startsWith('execution/')) return 'execution' if (key.startsWith('workspace/')) return 'workspace' + if (key.startsWith('assistant/')) return 'mothership' if (key.startsWith('profile-pictures/')) return 'profile-pictures' if (key.startsWith('og-images/')) return 'og-images' if (key.startsWith('workspace-logos/')) return 'workspace-logos' From 9432062d02a3ba9c49477af55e3fd835f442eb29 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Fri, 11 Sep 2026 11:35:17 -0700 Subject: [PATCH 2/2] fix(assistant): clean up orphaned image uploads --- .../uploads/upload-session/service.test.ts | 81 ++++++ .../sim/lib/uploads/upload-session/service.ts | 21 +- .../account-deletion-attachments.test.ts | 234 ++++++++++++++++++ apps/sim/lib/users/account-deletion.ts | 58 ++++- 4 files changed, 385 insertions(+), 9 deletions(-) create mode 100644 apps/sim/lib/users/account-deletion-attachments.test.ts diff --git a/apps/sim/lib/uploads/upload-session/service.test.ts b/apps/sim/lib/uploads/upload-session/service.test.ts index e8fedafbc12..d1712db1f66 100644 --- a/apps/sim/lib/uploads/upload-session/service.test.ts +++ b/apps/sim/lib/uploads/upload-session/service.test.ts @@ -1089,6 +1089,87 @@ describe('upload sessions', () => { expect.objectContaining({ key: FINAL_KEY, version: 'version-1' }) ) }) + + it('reclaims a completed Assistant image left behind by account deletion', async () => { + const image = uploadRow({ + purpose: 'mothership_attachment', + workspaceId: null, + storageContext: 'mothership', + finalKey: 'assistant/org-1/user-1/upload-1/image.png', + status: 'completed', + completedAt: new Date(Date.now() - 8 * 24 * 60 * 60 * 1000), + }) + queueTableRows(schemaMock.uploadSession, []) + queueTableRows(schemaMock.uploadSession, [image]) + mockHeadObject.mockResolvedValue(providerObject(sessionRecord(image), 'version-1')) + dbChainMockFns.returning + .mockResolvedValueOnce([image]) + .mockResolvedValueOnce([{ id: image.id }]) + + await expect(cleanupExpiredUploadSessions()).resolves.toEqual({ + expired: 0, + failed: 0, + purged: 1, + }) + expect(mockDeleteObjectVersion).toHaveBeenCalledWith({ + provider: 's3', + key: image.finalKey, + context: 'mothership', + version: 'version-1', + }) + expect(mockDeleteObjectVersion.mock.invocationCallOrder[0]).toBeLessThan( + dbChainMockFns.delete.mock.invocationCallOrder[0] + ) + }) + + it('retains orphan image ownership records when object deletion fails so cleanup can retry', async () => { + const image = uploadRow({ + purpose: 'mothership_attachment', + workspaceId: null, + storageContext: 'mothership', + status: 'completed', + completedAt: new Date(Date.now() - 8 * 24 * 60 * 60 * 1000), + }) + queueTableRows(schemaMock.uploadSession, []) + queueTableRows(schemaMock.uploadSession, [image]) + mockHeadObject.mockResolvedValue(providerObject(sessionRecord(image), 'version-1')) + mockDeleteObjectVersion.mockRejectedValueOnce(new Error('Storage unavailable')) + dbChainMockFns.returning.mockResolvedValueOnce([image]) + + await expect(cleanupExpiredUploadSessions()).resolves.toEqual({ + expired: 0, + failed: 1, + purged: 0, + }) + expect(dbChainMockFns.delete).not.toHaveBeenCalled() + expect(dbChainMockFns.set).toHaveBeenLastCalledWith( + expect.objectContaining({ + processingLeaseId: null, + processingLeaseExpiresAt: null, + error: 'Storage unavailable', + }) + ) + }) + + it('purges completed workspace attachment sessions without deleting their registered objects', async () => { + const attachment = uploadRow({ + purpose: 'mothership_attachment', + status: 'completed', + completedAt: new Date(Date.now() - 8 * 24 * 60 * 60 * 1000), + }) + queueTableRows(schemaMock.uploadSession, []) + queueTableRows(schemaMock.uploadSession, [attachment]) + dbChainMockFns.returning + .mockResolvedValueOnce([attachment]) + .mockResolvedValueOnce([{ id: attachment.id }]) + + await expect(cleanupExpiredUploadSessions()).resolves.toEqual({ + expired: 0, + failed: 0, + purged: 1, + }) + expect(mockDeleteObjectVersion).not.toHaveBeenCalled() + }) }) async function createWorkspaceUpload(fileSize: number) { diff --git a/apps/sim/lib/uploads/upload-session/service.ts b/apps/sim/lib/uploads/upload-session/service.ts index af5c7f8358c..5838abd8c74 100644 --- a/apps/sim/lib/uploads/upload-session/service.ts +++ b/apps/sim/lib/uploads/upload-session/service.ts @@ -5,7 +5,7 @@ import { requirePrincipalSubjectUserId, } from '@sim/auth/principal' import { db, dbFor } from '@sim/db' -import { uploadSession } from '@sim/db/schema' +import { organization, uploadSession, user } from '@sim/db/schema' import { safeCompare } from '@sim/security/compare' import { sha256Hex } from '@sim/security/hash' import { generateSecureToken } from '@sim/security/tokens' @@ -941,8 +941,17 @@ export async function cleanupExpiredUploadSessions(): Promise<{ .where( and( inArray(uploadSession.status, ['completed', 'aborted', 'expired']), - /** Completed organization sessions are the durable private-image ownership records. */ - sql`NOT (${uploadSession.status} = 'completed' AND ${uploadSession.purpose} = 'mothership_attachment' AND ${uploadSession.workspaceId} IS NULL)`, + /** Keep private images while both their uploader and organization exist. */ + sql`NOT ( + ${uploadSession.status} = 'completed' + AND ${uploadSession.purpose} = 'mothership_attachment' + AND ${uploadSession.workspaceId} IS NULL + AND EXISTS (SELECT 1 FROM ${user} WHERE ${user.id} = ${uploadSession.userId}) + AND EXISTS ( + SELECT 1 FROM ${organization} + WHERE ${organization.id} = ${uploadSession.metadata}->'organizationAttachment'->>'organizationId' + ) + )`, lt(uploadSession.completedAt, terminalCutoff), or( isNull(uploadSession.processingLeaseId), @@ -964,7 +973,11 @@ export async function cleanupExpiredUploadSessions(): Promise<{ candidate.status, cleanupDb ) - if (claimed.status === 'aborted' || claimed.status === 'expired') { + if ( + claimed.status === 'aborted' || + claimed.status === 'expired' || + (claimed.purpose === 'mothership_attachment' && claimed.workspaceId === null) + ) { await deleteOwnedFinalObject(claimed) } else if (claimed.status !== 'completed') { throw new Error(`Invalid terminal upload status ${claimed.status}`) diff --git a/apps/sim/lib/users/account-deletion-attachments.test.ts b/apps/sim/lib/users/account-deletion-attachments.test.ts new file mode 100644 index 00000000000..5b99873ead4 --- /dev/null +++ b/apps/sim/lib/users/account-deletion-attachments.test.ts @@ -0,0 +1,234 @@ +/** + * @vitest-environment node + */ +import { + dbChainMockFns, + hasMockCondition, + queueTableRows, + resetDbChainMock, + schemaMock, +} from '@sim/testing' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +const mocks = vi.hoisted(() => ({ + isSoleOwnerOfPaidOrganization: vi.fn(), + getPersonalSubscription: vi.fn(), + isUsingCloudStorage: vi.fn(), + deleteFiles: vi.fn(), +})) + +vi.mock('@/lib/billing/organizations/membership', () => ({ + isSoleOwnerOfPaidOrganization: mocks.isSoleOwnerOfPaidOrganization, +})) +vi.mock('@/lib/billing/core/plan', () => ({ + getHighestPriorityPersonalSubscription: mocks.getPersonalSubscription, +})) +vi.mock('@/lib/uploads', () => ({ + isUsingCloudStorage: mocks.isUsingCloudStorage, + StorageService: { deleteFiles: mocks.deleteFiles }, +})) +vi.mock('@/lib/workspaces/utils', () => ({ + reassignBilledAccountForUser: vi.fn(async () => ({ unresolved: [] })), + reassignOwnedWorkspacesForUser: vi.fn(async () => ({ unresolved: [] })), +})) +vi.mock('@/lib/table/events', () => ({ appendTableEvent: vi.fn() })) +vi.mock('@/lib/table/rows/executions', () => ({ + cancelPendingMarkersForGovernedSubject: vi.fn(async () => []), +})) + +import { deleteUserAccount } from '@/lib/users/account-deletion' + +const IMAGE_KEY = 'assistant/org-1/user-1/upload-1/photo.png' +const FAILED_IMAGE_KEY = 'assistant/org-2/user-1/upload-2/photo.png' +const NOW = new Date('2026-09-11T12:00:00Z') + +function imageDeletionFilter() { + return dbChainMockFns.where.mock.calls + .map(([condition]) => condition) + .find((condition) => + hasMockCondition( + condition, + (node) => node.type === 'inArray' && node.column === schemaMock.uploadSession.finalKey + ) + ) +} + +describe('account deletion of private organization Assistant images', () => { + beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() + vi.useFakeTimers() + vi.setSystemTime(NOW) + mocks.isSoleOwnerOfPaidOrganization.mockResolvedValue({ isBlocker: false }) + mocks.getPersonalSubscription.mockResolvedValue(null) + mocks.isUsingCloudStorage.mockReturnValue(true) + mocks.deleteFiles.mockResolvedValue({ deleted: 1, failed: [] }) + }) + + afterEach(() => { + vi.useRealTimers() + }) + + it.each([true, false])( + 'purges image objects and ownership records after deleting an account without workspaces (cloud: %s)', + async (cloudStorage) => { + mocks.isUsingCloudStorage.mockReturnValue(cloudStorage) + queueTableRows(schemaMock.uploadSession, [{ id: 'upload-1', key: IMAGE_KEY }]) + + const plan = await deleteUserAccount('user-1') + + expect(plan.workspacesToDelete).toEqual([]) + expect(mocks.deleteFiles).toHaveBeenCalledWith([IMAGE_KEY], 'mothership') + expect(dbChainMockFns.delete).toHaveBeenCalledWith(schemaMock.uploadSession) + const userDeleteIndex = dbChainMockFns.delete.mock.calls.findIndex( + ([table]) => table === schemaMock.user + ) + const imageDeleteIndex = dbChainMockFns.delete.mock.calls.findIndex( + ([table]) => table === schemaMock.uploadSession + ) + expect(dbChainMockFns.delete.mock.invocationCallOrder[userDeleteIndex]).toBeLessThan( + mocks.deleteFiles.mock.invocationCallOrder[0] + ) + expect(mocks.deleteFiles.mock.invocationCallOrder[0]).toBeLessThan( + dbChainMockFns.delete.mock.invocationCallOrder[imageDeleteIndex] + ) + } + ) + + it('scopes both collection and ownership deletion to this uploader’s completed organization images', async () => { + queueTableRows(schemaMock.uploadSession, [{ id: 'upload-1', key: IMAGE_KEY }]) + + await deleteUserAccount('user-1') + + const imageFilters = dbChainMockFns.where.mock.calls + .map(([condition]) => condition) + .filter((condition) => + hasMockCondition(condition, (node) => node.left === schemaMock.uploadSession.userId) + ) + expect(imageFilters).toHaveLength(2) + for (const filter of imageFilters) { + for (const [column, value] of [ + [schemaMock.uploadSession.userId, 'user-1'], + [schemaMock.uploadSession.purpose, 'mothership_attachment'], + [schemaMock.uploadSession.status, 'completed'], + ]) { + expect( + hasMockCondition( + filter, + (node) => node.type === 'eq' && node.left === column && node.right === value + ) + ).toBe(true) + } + expect( + hasMockCondition( + filter, + (node) => node.type === 'isNull' && node.column === schemaMock.uploadSession.workspaceId + ) + ).toBe(true) + } + }) + + it('retains ownership while an issued upload URL could recreate a purged object', async () => { + queueTableRows(schemaMock.uploadSession, [{ id: 'upload-1', key: IMAGE_KEY }]) + + await deleteUserAccount('user-1') + + expect( + hasMockCondition( + imageDeletionFilter(), + (node) => + node.type === 'lte' && + node.left === schemaMock.uploadSession.expiresAt && + node.right instanceof Date && + node.right.getTime() === NOW.getTime() + ) + ).toBe(true) + }) + + it('retains failed objects’ ownership records for the upload-session sweep', async () => { + queueTableRows(schemaMock.uploadSession, [ + { id: 'upload-1', key: IMAGE_KEY }, + { id: 'upload-2', key: FAILED_IMAGE_KEY }, + ]) + mocks.deleteFiles.mockResolvedValue({ + deleted: 1, + failed: [{ key: FAILED_IMAGE_KEY, error: 'Storage unavailable' }], + }) + + await deleteUserAccount('user-1') + + expect( + hasMockCondition( + imageDeletionFilter(), + (node) => + node.type === 'inArray' && + node.column === schemaMock.uploadSession.finalKey && + Array.isArray(node.values) && + node.values.length === 1 && + node.values[0] === IMAGE_KEY + ) + ).toBe(true) + }) + + it.each(['batch', 'object'])( + 'keeps ownership records when all %s deletions fail', + async (failure) => { + queueTableRows(schemaMock.uploadSession, [{ id: 'upload-1', key: IMAGE_KEY }]) + if (failure === 'batch') { + mocks.deleteFiles.mockRejectedValueOnce(new Error('Storage unavailable')) + } else { + mocks.deleteFiles.mockResolvedValueOnce({ + deleted: 0, + failed: [{ key: IMAGE_KEY, error: 'Storage unavailable' }], + }) + } + + await expect(deleteUserAccount('user-1')).resolves.toMatchObject({ blockers: [] }) + + expect(dbChainMockFns.delete).not.toHaveBeenCalledWith(schemaMock.uploadSession) + } + ) + + it('does not purge images or ownership records if account deletion rolls back', async () => { + queueTableRows(schemaMock.uploadSession, [{ id: 'upload-1', key: IMAGE_KEY }]) + dbChainMockFns.transaction.mockRejectedValueOnce(new Error('Transaction rolled back')) + + await expect(deleteUserAccount('user-1')).rejects.toThrow('Transaction rolled back') + + expect(mocks.deleteFiles).not.toHaveBeenCalled() + expect(dbChainMockFns.delete).not.toHaveBeenCalledWith(schemaMock.uploadSession) + }) + + it('leaves storage untouched when deletion is blocked for an active account', async () => { + mocks.getPersonalSubscription.mockResolvedValueOnce({ plan: 'pro' }) + + await expect(deleteUserAccount('user-1')).rejects.toMatchObject({ code: 'conflict' }) + + expect(dbChainMockFns.from).not.toHaveBeenCalledWith(schemaMock.uploadSession) + expect(mocks.deleteFiles).not.toHaveBeenCalled() + }) + + it('collects and purges image keys in bounded pages', async () => { + const firstPage = Array.from({ length: 1000 }, (_, index) => ({ + id: `upload-${String(index).padStart(4, '0')}`, + key: `assistant/org-1/user-1/upload-${index}/photo.png`, + })) + queueTableRows(schemaMock.uploadSession, firstPage) + queueTableRows(schemaMock.uploadSession, [{ id: 'upload-1000', key: IMAGE_KEY }]) + + await deleteUserAccount('user-1') + + expect(mocks.deleteFiles.mock.calls.map(([keys]) => keys.length)).toEqual([1000, 1]) + expect( + dbChainMockFns.where.mock.calls.some(([condition]) => + hasMockCondition( + condition, + (node) => + node.type === 'gt' && + node.left === schemaMock.uploadSession.id && + node.right === firstPage[999].id + ) + ) + ).toBe(true) + }) +}) diff --git a/apps/sim/lib/users/account-deletion.ts b/apps/sim/lib/users/account-deletion.ts index 447beba1b30..1b81d20988c 100644 --- a/apps/sim/lib/users/account-deletion.ts +++ b/apps/sim/lib/users/account-deletion.ts @@ -7,6 +7,7 @@ import { organization, permissions, tableRunDispatches, + uploadSession, user, workspaceFile, workspaceFiles, @@ -14,7 +15,7 @@ import { } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { formatQuotedNameList } from '@sim/utils/string' -import { and, eq, gt, inArray, isNotNull, ne, notExists, or, sql } from 'drizzle-orm' +import { and, eq, gt, inArray, isNotNull, isNull, lte, ne, notExists, or, sql } from 'drizzle-orm' import type { AccountDeletionBlocker, AccountDeletionPlan, @@ -415,7 +416,7 @@ export function extractProfilePictureKey(image: string | null): string | null { } /** - * Collects every stored object held by workspaces that go with the account. + * Collects the account's private images and stored objects in workspaces that go with it. * * This has to run *before* the rows are deleted: they disappear with the * workspace through `ON DELETE CASCADE`, and the retention sweep that normally @@ -429,6 +430,27 @@ async function collectAccountStorageKeys( workspaceIds: string[] ): Promise { const batches: StorageKeyBatch[] = [] + + await collectPages( + (afterId) => + db + .select({ id: uploadSession.id, key: uploadSession.finalKey }) + .from(uploadSession) + .where( + and( + eq(uploadSession.userId, userId), + eq(uploadSession.purpose, 'mothership_attachment'), + isNull(uploadSession.workspaceId), + eq(uploadSession.status, 'completed'), + gt(uploadSession.id, afterId) + ) + ) + .orderBy(uploadSession.id) + .limit(STORAGE_PAGE_SIZE), + batches, + () => 'mothership' + ) + if (!isUsingCloudStorage()) return batches const [profile] = await db @@ -497,7 +519,7 @@ async function collectAccountStorageKeys( * failure for work that cannot be undone. An orphaned object is recoverable from * the log; a deletion the caller believes failed is not. */ -async function purgeStorageObjects(batches: StorageKeyBatch[]): Promise { +async function purgeStorageObjects(userId: string, batches: StorageKeyBatch[]): Promise { for (const { context, keys } of batches) { if (keys.length === 0) continue try { @@ -509,8 +531,34 @@ async function purgeStorageObjects(batches: StorageKeyBatch[]): Promise { error, }) } + + if (context === 'mothership') { + const failedKeys = new Set(failed.map(({ key }) => key)) + const deletedImageKeys = keys.filter( + (key) => key.startsWith('assistant/') && !failedKeys.has(key) + ) + if (deletedImageKeys.length > 0) { + /** + * Keep ownership records while a signed PUT can recreate the object. + * The upload-session sweep retries these and failed object deletions + * after the deleted uploader and transfer expiry are confirmed. + */ + await db + .delete(uploadSession) + .where( + and( + eq(uploadSession.userId, userId), + eq(uploadSession.purpose, 'mothership_attachment'), + isNull(uploadSession.workspaceId), + eq(uploadSession.status, 'completed'), + lte(uploadSession.expiresAt, new Date()), + inArray(uploadSession.finalKey, deletedImageKeys) + ) + ) + } + } } catch (error) { - logger.error('Storage batch deletion failed during account deletion', { context, error }) + logger.error('Storage cleanup failed during account deletion', { context, error }) } } } @@ -743,7 +791,7 @@ export async function deleteUserAccount(userId: string): Promise