Skip to content

Commit 6bccb71

Browse files
j15zj15z
committed
fix(media): validate upload paths before generation
Allow media tools to consume chat-scoped uploads directly while keeping outputs under files/. Fail malformed or unresolved declarations before provider work and regenerate the Mothership-derived tool schemas.
1 parent 1b82b97 commit 6bccb71

11 files changed

Lines changed: 869 additions & 553 deletions

File tree

‎apps/sim/lib/copilot/generated/tool-catalog-v1.ts‎

Lines changed: 52 additions & 229 deletions
Large diffs are not rendered by default.

‎apps/sim/lib/copilot/generated/tool-schemas-v1.ts‎

Lines changed: 52 additions & 253 deletions
Large diffs are not rendered by default.

‎apps/sim/lib/copilot/tools/handlers/upload-file-reader.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,17 @@ function toWorkspaceFileRecord(row: typeof workspaceFiles.$inferSelect): Workspa
6161
}
6262
}
6363

64+
/**
65+
* Resolve a chat-scoped upload to the same file-record shape used by workspace files.
66+
*/
67+
export async function resolveChatUpload(
68+
filename: string,
69+
chatId: string
70+
): Promise<WorkspaceFileRecord | null> {
71+
const row = await findMothershipUploadRowByChatAndName(chatId, filename)
72+
return row ? toWorkspaceFileRecord(row) : null
73+
}
74+
6475
/**
6576
* Resolve a mothership upload row by VFS name (the collision-disambiguated `displayName`
6677
* for new rows, or `originalName` for legacy rows that predate the column). Prefers an

‎apps/sim/lib/copilot/tools/server/image/generate-image.test.ts‎

Lines changed: 211 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,211 @@
1+
/**
2+
* @vitest-environment node
3+
*/
4+
import { beforeEach, describe, expect, it, vi } from 'vitest'
5+
6+
const mocks = vi.hoisted(() => ({
7+
fetchWorkspaceFileBuffer: vi.fn(),
8+
generateContent: vi.fn(),
9+
resolveChatUpload: vi.fn(),
10+
resolveWorkspaceFileReference: vi.fn(),
11+
validateWorkspaceFileWriteTarget: vi.fn(),
12+
writeWorkspaceFileByPath: vi.fn(),
13+
}))
14+
15+
vi.mock('@google/genai', () => ({
16+
GoogleGenAI: class {
17+
models = { generateContent: mocks.generateContent }
18+
},
19+
}))
20+
21+
vi.mock('@/lib/core/config/api-keys', () => ({
22+
getRotatingApiKey: vi.fn(() => 'test-api-key'),
23+
}))
24+
25+
vi.mock('@/lib/copilot/tools/handlers/upload-file-reader', () => ({
26+
resolveChatUpload: mocks.resolveChatUpload,
27+
}))
28+
29+
vi.mock('@/lib/uploads/contexts/workspace/workspace-file-manager', () => ({
30+
fetchWorkspaceFileBuffer: mocks.fetchWorkspaceFileBuffer,
31+
resolveWorkspaceFileReference: mocks.resolveWorkspaceFileReference,
32+
}))
33+
34+
vi.mock('@/lib/copilot/vfs/resource-writer', () => ({
35+
validateWorkspaceFileWriteTarget: mocks.validateWorkspaceFileWriteTarget,
36+
writeWorkspaceFileByPath: mocks.writeWorkspaceFileByPath,
37+
}))
38+
39+
import { generateImageServerTool } from '@/lib/copilot/tools/server/image/generate-image'
40+
41+
const CONTEXT = {
42+
userId: 'user-1',
43+
workspaceId: 'workspace-1',
44+
chatId: 'chat-1',
45+
}
46+
47+
const IMAGE_RECORD = {
48+
id: 'wf_image',
49+
workspaceId: 'workspace-1',
50+
name: 'person.png',
51+
key: 'uploads/person.png',
52+
path: '/api/files/serve/uploads%2Fperson.png',
53+
size: 3,
54+
type: 'image/png',
55+
uploadedBy: 'user-1',
56+
uploadedAt: new Date('2026-01-01T00:00:00Z'),
57+
updatedAt: new Date('2026-01-01T00:00:00Z'),
58+
storageContext: 'mothership' as const,
59+
}
60+
61+
describe('generateImageServerTool', () => {
62+
beforeEach(() => {
63+
vi.clearAllMocks()
64+
mocks.fetchWorkspaceFileBuffer.mockResolvedValue(Buffer.from('image'))
65+
mocks.validateWorkspaceFileWriteTarget.mockResolvedValue({})
66+
mocks.generateContent.mockResolvedValue({
67+
candidates: [
68+
{ content: { parts: [{ inlineData: { data: 'aW1hZ2U=', mimeType: 'image/png' } }] } },
69+
],
70+
})
71+
mocks.writeWorkspaceFileByPath.mockResolvedValue({
72+
id: 'wf_result',
73+
name: 'result.png',
74+
vfsPath: 'files/result.png',
75+
mode: 'create',
76+
downloadUrl: '/download/result.png',
77+
})
78+
})
79+
80+
it('loads a chat upload when inputs.files uses an uploads/ path', async () => {
81+
mocks.resolveChatUpload.mockResolvedValue(IMAGE_RECORD)
82+
83+
const result = await generateImageServerTool.execute(
84+
{
85+
prompt: 'Preserve this person and put them in a library',
86+
inputs: { files: [{ path: 'uploads/person.png' }] },
87+
outputs: { files: [{ path: 'files/result.png' }] },
88+
},
89+
CONTEXT
90+
)
91+
92+
expect(result.success).toBe(true)
93+
expect(mocks.resolveChatUpload).toHaveBeenCalledWith('person.png', 'chat-1')
94+
expect(mocks.resolveWorkspaceFileReference).not.toHaveBeenCalled()
95+
expect(mocks.fetchWorkspaceFileBuffer).toHaveBeenCalledWith(IMAGE_RECORD)
96+
})
97+
98+
it('fails before generation when an input path cannot be resolved', async () => {
99+
mocks.resolveWorkspaceFileReference.mockResolvedValue(null)
100+
101+
const result = await generateImageServerTool.execute(
102+
{
103+
prompt: 'Edit the reference image',
104+
inputs: { files: [{ path: 'files/missing.png' }] },
105+
outputs: { files: [{ path: 'files/result.png' }] },
106+
},
107+
CONTEXT
108+
)
109+
110+
expect(result).toEqual({
111+
success: false,
112+
message: 'Failed to generate image: Input file not found: files/missing.png',
113+
})
114+
expect(mocks.generateContent).not.toHaveBeenCalled()
115+
expect(mocks.writeWorkspaceFileByPath).not.toHaveBeenCalled()
116+
})
117+
118+
it('fails before generation when outputs.files targets uploads/', async () => {
119+
const result = await generateImageServerTool.execute(
120+
{
121+
prompt: 'Create a portrait',
122+
outputs: { files: [{ path: 'uploads/result.png' }] },
123+
},
124+
CONTEXT
125+
)
126+
127+
expect(result).toEqual({
128+
success: false,
129+
message:
130+
'Failed to generate image: Media output paths must start with "files/"; uploads/ paths are read-only: uploads/result.png',
131+
})
132+
expect(mocks.generateContent).not.toHaveBeenCalled()
133+
expect(mocks.writeWorkspaceFileByPath).not.toHaveBeenCalled()
134+
})
135+
136+
it('fails before generation when an explicit output path is empty', async () => {
137+
mocks.validateWorkspaceFileWriteTarget.mockRejectedValue(
138+
new Error('Workspace file paths must start with "files/"')
139+
)
140+
141+
const result = await generateImageServerTool.execute(
142+
{
143+
prompt: 'Create a portrait',
144+
outputs: { files: [{ path: '' }] },
145+
},
146+
CONTEXT
147+
)
148+
149+
expect(result.success).toBe(false)
150+
expect(mocks.generateContent).not.toHaveBeenCalled()
151+
expect(mocks.writeWorkspaceFileByPath).not.toHaveBeenCalled()
152+
})
153+
154+
it('rejects additional output declarations before generation', async () => {
155+
const result = await generateImageServerTool.execute(
156+
{
157+
prompt: 'Create a portrait',
158+
outputs: {
159+
files: [{ path: 'files/one.png' }, { path: 'uploads/two.png' }],
160+
},
161+
},
162+
CONTEXT
163+
)
164+
165+
expect(result).toEqual({
166+
success: false,
167+
message: 'Failed to generate image: Output requires exactly one file; received 2',
168+
})
169+
expect(mocks.validateWorkspaceFileWriteTarget).not.toHaveBeenCalled()
170+
expect(mocks.generateContent).not.toHaveBeenCalled()
171+
})
172+
173+
it('rejects explicitly empty input and output containers before generation', async () => {
174+
const emptyInputs = await generateImageServerTool.execute(
175+
{ prompt: 'Edit the supplied image', inputs: {} },
176+
CONTEXT
177+
)
178+
const emptyOutputs = await generateImageServerTool.execute(
179+
{ prompt: 'Create a portrait', outputs: {} },
180+
CONTEXT
181+
)
182+
183+
expect(emptyInputs).toEqual({
184+
success: false,
185+
message: 'Failed to generate image: Input requires at least one file',
186+
})
187+
expect(emptyOutputs).toEqual({
188+
success: false,
189+
message: 'Failed to generate image: Output requires exactly one file; received 0',
190+
})
191+
expect(mocks.generateContent).not.toHaveBeenCalled()
192+
})
193+
194+
it('fails before generation when canonical output preflight rejects a files/ target', async () => {
195+
mocks.validateWorkspaceFileWriteTarget.mockRejectedValue(
196+
new Error('File already exists at files/result.png')
197+
)
198+
199+
const result = await generateImageServerTool.execute(
200+
{
201+
prompt: 'Create a portrait',
202+
outputs: { files: [{ path: 'files/result.png' }] },
203+
},
204+
CONTEXT
205+
)
206+
207+
expect(result.success).toBe(false)
208+
expect(mocks.generateContent).not.toHaveBeenCalled()
209+
expect(mocks.writeWorkspaceFileByPath).not.toHaveBeenCalled()
210+
})
211+
})

‎apps/sim/lib/copilot/tools/server/image/generate-image.ts‎

Lines changed: 35 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,20 @@
11
import { GoogleGenAI, type Part } from '@google/genai'
22
import { createLogger } from '@sim/logger'
3-
import { getErrorMessage, toError } from '@sim/utils/errors'
3+
import { getErrorMessage } from '@sim/utils/errors'
44
import { GenerateImage } from '@/lib/copilot/generated/tool-catalog-v1'
55
import {
66
assertServerToolNotAborted,
77
type BaseServerTool,
88
type ServerToolContext,
99
} from '@/lib/copilot/tools/server/base-tool'
10+
import {
11+
prepareMediaOutput,
12+
requireMediaFileDeclarations,
13+
resolveMediaInputFile,
14+
} from '@/lib/copilot/tools/server/media/file-paths'
1015
import { writeWorkspaceFileByPath } from '@/lib/copilot/vfs/resource-writer'
1116
import { getRotatingApiKey } from '@/lib/core/config/api-keys'
12-
import {
13-
fetchWorkspaceFileBuffer,
14-
resolveWorkspaceFileReference,
15-
} from '@/lib/uploads/contexts/workspace/workspace-file-manager'
17+
import { fetchWorkspaceFileBuffer } from '@/lib/uploads/contexts/workspace/workspace-file-manager'
1618

1719
const logger = createLogger('GenerateImageTool')
1820

@@ -74,6 +76,12 @@ export const generateImageServerTool: BaseServerTool<GenerateImageArgs, Generate
7476
}
7577

7678
try {
79+
const outputFile = await prepareMediaOutput({
80+
output: params.outputs,
81+
workspaceId,
82+
userId: context.userId,
83+
})
84+
7785
const apiKey = getRotatingApiKey('gemini')
7886
const ai = new GoogleGenAI({ apiKey })
7987

@@ -82,34 +90,29 @@ export const generateImageServerTool: BaseServerTool<GenerateImageArgs, Generate
8290

8391
const parts: Part[] = []
8492

85-
const referencePaths = params.inputs?.files?.map((file) => file.path) ?? []
93+
const referencePaths = params.inputs
94+
? requireMediaFileDeclarations(params.inputs.files, 'Input').map((file) => file.path)
95+
: []
8696

8797
if (referencePaths.length) {
8898
for (const filePath of referencePaths) {
89-
try {
90-
const fileRecord = await resolveWorkspaceFileReference(workspaceId, filePath)
91-
if (fileRecord) {
92-
const buffer = await fetchWorkspaceFileBuffer(fileRecord)
93-
const base64 = buffer.toString('base64')
94-
const mime = fileRecord.type || 'image/png'
95-
parts.push({
96-
inlineData: { mimeType: mime, data: base64 },
97-
})
98-
logger.info('Loaded reference image', {
99-
filePath,
100-
name: fileRecord.name,
101-
size: buffer.length,
102-
mimeType: mime,
103-
})
104-
} else {
105-
logger.warn('Reference file not found, skipping', { filePath })
106-
}
107-
} catch (err) {
108-
logger.warn('Failed to load reference image, skipping', {
109-
filePath,
110-
error: toError(err).message,
111-
})
112-
}
99+
const fileRecord = await resolveMediaInputFile({
100+
workspaceId,
101+
chatId: context.chatId,
102+
path: filePath,
103+
})
104+
const buffer = await fetchWorkspaceFileBuffer(fileRecord)
105+
const base64 = buffer.toString('base64')
106+
const mime = fileRecord.type || 'image/png'
107+
parts.push({
108+
inlineData: { mimeType: mime, data: base64 },
109+
})
110+
logger.info('Loaded reference image', {
111+
filePath,
112+
name: fileRecord.name,
113+
size: buffer.length,
114+
mimeType: mime,
115+
})
113116
}
114117
}
115118

@@ -161,8 +164,7 @@ export const generateImageServerTool: BaseServerTool<GenerateImageArgs, Generate
161164
}
162165

163166
const ext = mimeType.includes('jpeg') || mimeType.includes('jpg') ? '.jpg' : '.png'
164-
const outputFile = params.outputs?.files?.[0]
165-
const outputPath = outputFile?.path || `files/generated-image${ext}`
167+
const resolvedOutputPath = outputFile?.path || `files/generated-image${ext}`
166168
const imageBuffer = Buffer.from(imageBase64, 'base64')
167169
const mode = outputFile?.mode ?? 'create'
168170

@@ -171,7 +173,7 @@ export const generateImageServerTool: BaseServerTool<GenerateImageArgs, Generate
171173
workspaceId,
172174
userId: context.userId,
173175
target: {
174-
path: outputPath,
176+
path: resolvedOutputPath,
175177
mode,
176178
mimeType: outputFile?.mimeType,
177179
},

‎apps/sim/lib/copilot/tools/server/media/ffmpeg.ts‎

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,13 @@ import {
66
type BaseServerTool,
77
type ServerToolContext,
88
} from '@/lib/copilot/tools/server/base-tool'
9+
import {
10+
prepareMediaOutput,
11+
resolveMediaInputFile,
12+
} from '@/lib/copilot/tools/server/media/file-paths'
913
import { writeWorkspaceFileByPath } from '@/lib/copilot/vfs/resource-writer'
1014
import { type FfmpegOperation, type MediaFile, runFfmpegOperation } from '@/lib/media/ffmpeg'
11-
import {
12-
fetchWorkspaceFileBuffer,
13-
resolveWorkspaceFileReference,
14-
} from '@/lib/uploads/contexts/workspace/workspace-file-manager'
15+
import { fetchWorkspaceFileBuffer } from '@/lib/uploads/contexts/workspace/workspace-file-manager'
1516

1617
const logger = createLogger('FfmpegTool')
1718

@@ -81,12 +82,19 @@ export const ffmpegServerTool: BaseServerTool<FfmpegArgs, FfmpegResult> = {
8182
}
8283

8384
try {
85+
const outputFile = await prepareMediaOutput({
86+
output: params.outputs,
87+
workspaceId,
88+
userId: context.userId,
89+
})
90+
8491
const mediaFiles: MediaFile[] = []
8592
for (const filePath of inputPaths) {
86-
const fileRecord = await resolveWorkspaceFileReference(workspaceId, filePath)
87-
if (!fileRecord) {
88-
return { success: false, message: `Input file not found: ${filePath}` }
89-
}
93+
const fileRecord = await resolveMediaInputFile({
94+
workspaceId,
95+
chatId: context.chatId,
96+
path: filePath,
97+
})
9098
const buffer = await fetchWorkspaceFileBuffer(fileRecord)
9199
mediaFiles.push({
92100
buffer,
@@ -123,7 +131,6 @@ export const ffmpegServerTool: BaseServerTool<FfmpegArgs, FfmpegResult> = {
123131
return { success: false, message: `ffmpeg ${params.operation} produced no output` }
124132
}
125133

126-
const outputFile = params.outputs?.files?.[0]
127134
const outputPath = outputFile?.path || `files/ffmpeg-${params.operation}.${result.ext}`
128135
const mode = outputFile?.mode ?? 'create'
129136

0 commit comments

Comments
 (0)