diff --git a/src/server/__tests__/conversation-attachments.test.ts b/src/server/__tests__/conversation-attachments.test.ts new file mode 100644 index 00000000..5662046c --- /dev/null +++ b/src/server/__tests__/conversation-attachments.test.ts @@ -0,0 +1,192 @@ +import { afterEach, beforeEach, describe, expect, mock, test } from 'bun:test' +import * as fs from 'node:fs/promises' +import * as os from 'node:os' +import * as path from 'node:path' + +let imageProcessorShouldThrow = false + +const getImageProcessorMock = mock(async () => { + if (imageProcessorShouldThrow) { + throw new Error('image processor unavailable') + } + return (input: Buffer) => { + let output = input + const instance = { + metadata: async () => ({ width: 3000, height: 4000, format: 'png' }), + resize: () => { + output = Buffer.from('resized-image') + return instance + }, + jpeg: () => { + output = Buffer.from('jpeg-image') + return instance + }, + png: () => { + output = Buffer.from('png-image') + return instance + }, + webp: () => { + output = Buffer.from('webp-image') + return instance + }, + toBuffer: async () => output, + } + return instance + } +}) + +mock.module('../../tools/FileReadTool/imageProcessor.js', () => ({ + getImageProcessor: getImageProcessorMock, +})) + +const { ConversationService } = await import('../services/conversationService.js') + +let tmpDir: string +let originalConfigDir: string | undefined + +beforeEach(async () => { + getImageProcessorMock.mockClear() + imageProcessorShouldThrow = false + tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), 'conversation-attachments-')) + originalConfigDir = process.env.CLAUDE_CONFIG_DIR + process.env.CLAUDE_CONFIG_DIR = tmpDir +}) + +afterEach(async () => { + if (originalConfigDir === undefined) delete process.env.CLAUDE_CONFIG_DIR + else process.env.CLAUDE_CONFIG_DIR = originalConfigDir + await fs.rm(tmpDir, { recursive: true, force: true }) +}) + +describe('ConversationService attachment materialization', () => { + test('inlines normalized image data attachments as SDK image blocks', async () => { + const svc = new ConversationService() + const sent: unknown[] = [] + const sessionId = 'session-image-normalize' + const original = Buffer.from('original-image') + + ;(svc as any).sessions.set(sessionId, { + sdkSocket: { + send(data: string) { + sent.push(JSON.parse(data)) + }, + }, + pendingOutbound: [], + }) + + const ok = await svc.sendMessage(sessionId, '这张图说了什么？', [ + { + type: 'image', + name: 'pasted-image.png', + mimeType: 'image/png', + data: `data:image/png;base64,${original.toString('base64')}`, + }, + ]) + + expect(ok).toBe(true) + expect(getImageProcessorMock).toHaveBeenCalled() + expect(sent).toHaveLength(1) + + const payload = sent[0] as { + message: { content: Array<{ type: string; text?: string; source?: { media_type?: string; data?: string } }> } + } + const textBlocks = payload.message.content.filter((block) => block.type === 'text') + const imageBlocks = payload.message.content.filter((block) => block.type === 'image') + expect(textBlocks[0]?.text).toBe('这张图说了什么？') + expect(textBlocks.some((block) => block.text?.includes('@"'))).toBe(false) + expect(imageBlocks).toHaveLength(1) + expect(imageBlocks[0]?.source?.media_type).toBe('image/png') + expect(imageBlocks[0]?.source?.data).toBe(Buffer.from('resized-image').toString('base64')) + + const metadataText = textBlocks.find((block) => block.text?.startsWith('[Image:'))?.text + const uploadPath = metadataText?.match(/source: ([^,\]]+)/)?.[1] + expect(uploadPath).toBeTruthy() + expect(uploadPath?.endsWith('.png')).toBe(true) + expect(await fs.readFile(uploadPath!)).toEqual(Buffer.from('resized-image')) + }) + + test('falls back to an upload path when image normalization cannot produce a block', async () => { + const svc = new ConversationService() + const sent: unknown[] = [] + const sessionId = 'session-image-fallback' + const original = createOversizedPngHeader() + imageProcessorShouldThrow = true + + ;(svc as any).sessions.set(sessionId, { + sdkSocket: { + send(data: string) { + sent.push(JSON.parse(data)) + }, + }, + pendingOutbound: [], + }) + + const ok = await svc.sendMessage(sessionId, '', [ + { + type: 'image', + name: 'pasted-image.png', + mimeType: 'image/png', + data: `data:image/png;base64,${original.toString('base64')}`, + }, + ]) + + expect(ok).toBe(true) + + const payload = sent[0] as { message: { content: Array<{ text: string }> } } + const text = payload.message.content[0]?.text ?? '' + const uploadPath = text.match(/@"([^"]+)"/)?.[1] + expect(uploadPath).toBeTruthy() + expect(uploadPath?.endsWith('.png')).toBe(true) + expect(await fs.readFile(uploadPath!)).toEqual(original) + }) + + test('inlines image file paths instead of asking the model to Read them first', async () => { + const svc = new ConversationService() + const sent: unknown[] = [] + const sessionId = 'session-image-path' + const original = Buffer.from('path-image') + const imagePath = path.join(tmpDir, 'screen.png') + await fs.writeFile(imagePath, original) + + ;(svc as any).sessions.set(sessionId, { + sdkSocket: { + send(data: string) { + sent.push(JSON.parse(data)) + }, + }, + pendingOutbound: [], + }) + + const ok = await svc.sendMessage(sessionId, '看这个截图', [ + { + type: 'file', + path: imagePath, + }, + ]) + + expect(ok).toBe(true) + + const payload = sent[0] as { + message: { content: Array<{ type: string; text?: string; source?: { media_type?: string; data?: string } }> } + } + const textBlocks = payload.message.content.filter((block) => block.type === 'text') + const imageBlocks = payload.message.content.filter((block) => block.type === 'image') + expect(textBlocks[0]?.text).toBe('看这个截图') + expect(textBlocks.some((block) => block.text?.includes('@"'))).toBe(false) + expect(textBlocks.some((block) => block.text?.includes(`source: ${imagePath}`))).toBe(true) + expect(imageBlocks).toHaveLength(1) + expect(imageBlocks[0]?.source?.media_type).toBe('image/png') + expect(imageBlocks[0]?.source?.data).toBe(Buffer.from('resized-image').toString('base64')) + }) +}) + +function createOversizedPngHeader(): Buffer { + const buffer = Buffer.alloc(24) + buffer[0] = 0x89 + buffer[1] = 0x50 + buffer[2] = 0x4e + buffer[3] = 0x47 + buffer.writeUInt32BE(3000, 16) + buffer.writeUInt32BE(4000, 20) + return buffer +} diff --git a/src/server/services/conversationService.ts b/src/server/services/conversationService.ts index dbad6054..5bccfb50 100644 --- a/src/server/services/conversationService.ts +++ b/src/server/services/conversationService.ts @@ -32,6 +32,11 @@ import { sanitizePath } from '../../utils/path.js' import { getProcessEnvWithTerminalShellEnvironment } from '../../utils/terminalShellEnvironment.js' import { attributionHeaderEnvForModel } from './attributionHeaderPolicy.js' import { buildNetworkEnvironment, loadNetworkSettings } from './networkSettings.js' +import { logError } from '../../utils/log.js' +import { + createImageMetadataText, + maybeResizeAndDownsampleImageBuffer, +} from '../../utils/imageResizer.js' const MAX_CAPTURED_PROCESS_LINES = 80 const MAX_CAPTURED_SDK_MESSAGES = 40 @@ -48,6 +53,14 @@ type AttachmentRef = { isDirectory?: boolean } +type UserContentBlock = Record + +type MaterializedAttachments = { + pathPrefix: string + imageBlocks: UserContentBlock[] + imageMetadataTexts: string[] +} + type SessionProcess = { proc: ReturnType outputCallbacks: Array<(msg: any) => void> @@ -393,16 +406,17 @@ export class ConversationService { return this.sessions.get(sessionId)?.initMessage ?? null } - sendMessage( + async sendMessage( sessionId: string, content: string, attachments?: AttachmentRef[], - ): boolean { + ): Promise { + const userContent = await this.buildUserContent(content, sessionId, attachments) return this.sendSdkMessage(sessionId, { type: 'user', message: { role: 'user', - content: this.buildUserContent(content, sessionId, attachments), + content: userContent, }, parent_tool_use_id: null, session_id: '', @@ -1424,26 +1438,43 @@ export class ConversationService { }) } - private buildUserContent( + private async buildUserContent( content: string, sessionId: string, attachments?: AttachmentRef[], - ): Array> { - const prefix = this.materializeAttachments(sessionId, attachments) + ): Promise { + const materialized = await this.materializeAttachments(sessionId, attachments) const trimmed = content.trim() - const text = prefix - ? `${prefix}${trimmed || 'Please analyze the attached files.'}`.trim() + const text = materialized.pathPrefix + ? `${materialized.pathPrefix}${trimmed || 'Please analyze the attached files.'}`.trim() : trimmed - return [{ type: 'text', text }] + const blocks: UserContentBlock[] = text + ? [{ type: 'text', text }] + : materialized.imageBlocks.length > 0 + ? [{ type: 'text', text: 'Please analyze the attached image.' }] + : [] + + blocks.push(...materialized.imageBlocks) + for (const metadataText of materialized.imageMetadataTexts) { + blocks.push({ type: 'text', text: metadataText }) + } + + return blocks.length > 0 ? blocks : [{ type: 'text', text: '' }] } - private materializeAttachments( + private async materializeAttachments( sessionId: string, attachments?: AttachmentRef[], - ): string { + ): Promise { + const empty = (): MaterializedAttachments => ({ + pathPrefix: '', + imageBlocks: [], + imageMetadataTexts: [], + }) + if (!attachments || attachments.length === 0) { - return '' + return empty() } const uploadDir = path.join( @@ -1451,10 +1482,20 @@ export class ConversationService { 'uploads', sessionId, ) - fs.mkdirSync(uploadDir, { recursive: true }) const savedPaths: string[] = [] + const imageBlocks: UserContentBlock[] = [] + const imageMetadataTexts: string[] = [] for (const attachment of attachments) { + if (this.shouldInlineImageAttachment(attachment)) { + const image = await this.materializeImageAttachment(attachment, uploadDir) + if (image) { + imageBlocks.push(image.block) + if (image.metadataText) imageMetadataTexts.push(image.metadataText) + continue + } + } + if (attachment.path) { savedPaths.push(attachment.path) continue @@ -1462,38 +1503,156 @@ export class ConversationService { if (!attachment.data) continue - const payload = this.parseAttachmentData(attachment.data) - if (!payload) continue + const parsed = this.parseAttachmentData(attachment.data) + if (!parsed) continue - const ext = this.getAttachmentExtension(attachment) + const ext = this.getAttachmentExtension({ + ...attachment, + mimeType: attachment.mimeType ?? parsed.mimeType, + }) const fileName = this.sanitizeAttachmentName(attachment.name, attachment.type, ext) - const outPath = path.join(uploadDir, `${crypto.randomUUID()}-${fileName}`) - fs.writeFileSync(outPath, payload) + const outPath = this.writeUploadAttachment(uploadDir, fileName, parsed.payload) savedPaths.push(outPath) } - if (savedPaths.length === 0) { - return '' + return { + pathPrefix: savedPaths.length > 0 + ? savedPaths.map((filePath) => `@"${filePath}"`).join(' ') + ' ' + : '', + imageBlocks, + imageMetadataTexts, } - - return savedPaths.map((filePath) => `@"${filePath}"`).join(' ') + ' ' } - private parseAttachmentData(data: string): Buffer | null { - const match = data.match(/^data:.*?;base64,(.*)$/) - const encoded = match ? match[1] : data + private parseAttachmentData(data: string): { payload: Buffer; mimeType?: string } | null { + const match = data.match(/^data:([^;,]+)?;base64,(.*)$/) + const encoded = match ? match[2] : data try { - return Buffer.from(encoded, 'base64') + return { + payload: Buffer.from(encoded ?? '', 'base64'), + mimeType: match?.[1], + } } catch { return null } } + private async materializeImageAttachment( + attachment: AttachmentRef, + uploadDir: string, + ): Promise<{ block: UserContentBlock; metadataText?: string } | null> { + const source = this.readImageAttachmentPayload(attachment) + if (!source) { + return null + } + + try { + const resized = await maybeResizeAndDownsampleImageBuffer( + source.payload, + source.payload.length, + source.ext, + ) + const normalizedExt = this.normalizeImageExtension(resized.mediaType) + const storedName = this.replaceFileExtension( + this.sanitizeAttachmentName(attachment.name, attachment.type, normalizedExt), + normalizedExt, + ) + const sourcePath = source.sourcePath ?? this.writeUploadAttachment( + uploadDir, + storedName, + resized.buffer, + ) + const metadataText = resized.dimensions + ? createImageMetadataText(resized.dimensions, sourcePath) + : sourcePath + ? `[Image source: ${sourcePath}]` + : undefined + + return { + block: { + type: 'image', + source: { + type: 'base64', + media_type: `image/${normalizedExt}`, + data: resized.buffer.toString('base64'), + }, + }, + metadataText: metadataText ?? undefined, + } + } catch (error) { + logError(error) + console.warn( + `[ConversationService] Failed to inline image attachment ${attachment.name ?? ''}; falling back to file path`, + ) + return null + } + } + + private readImageAttachmentPayload( + attachment: AttachmentRef, + ): { payload: Buffer; ext: string; sourcePath?: string } | null { + if (attachment.data) { + const parsed = this.parseAttachmentData(attachment.data) + if (!parsed) return null + return { + payload: parsed.payload, + ext: this.getAttachmentExtension({ + ...attachment, + mimeType: attachment.mimeType ?? parsed.mimeType, + }), + } + } + + if (!attachment.path || attachment.isDirectory) { + return null + } + + try { + return { + payload: fs.readFileSync(attachment.path), + ext: this.getAttachmentExtension(attachment), + sourcePath: attachment.path, + } + } catch (error) { + logError(error) + return null + } + } + + private shouldInlineImageAttachment(attachment: AttachmentRef): boolean { + if (attachment.isDirectory) return false + if (attachment.type === 'image') return true + if (attachment.mimeType?.startsWith('image/')) return true + const candidate = attachment.path ?? attachment.name ?? '' + return /\.(png|jpe?g|gif|webp)$/i.test(candidate) + } + + private writeUploadAttachment(uploadDir: string, fileName: string, payload: Buffer): string { + fs.mkdirSync(uploadDir, { recursive: true }) + const outPath = path.join(uploadDir, `${crypto.randomUUID()}-${fileName}`) + fs.writeFileSync(outPath, payload) + return outPath + } + + private normalizeImageExtension(ext: string): string { + const clean = ext.split('/').pop()?.split('+')[0]?.toLowerCase() || 'png' + return clean === 'jpg' ? 'jpeg' : clean + } + + private replaceFileExtension(fileName: string, ext: string): string { + const cleanExt = this.normalizeImageExtension(ext) + const base = fileName.replace(/\.[a-z0-9]+$/i, '') + return `${base}.${cleanExt}` + } + private getAttachmentExtension(attachment: AttachmentRef): string { const byName = attachment.name?.match(/\.([a-z0-9]+)$/i)?.[1] if (byName) return byName + const byPath = attachment.path?.match(/\.([a-z0-9]+)$/i)?.[1] + if (byPath) return byPath + const byMime = attachment.mimeType?.split('/')[1]?.split('+')[0] if (byMime) return byMime diff --git a/src/server/ws/handler.ts b/src/server/ws/handler.ts index 465abb41..0f6a8400 100644 --- a/src/server/ws/handler.ts +++ b/src/server/ws/handler.ts @@ -365,7 +365,7 @@ async function handleUserMessage( }, }) - const sent = conversationService.sendMessage( + const sent = await conversationService.sendMessage( sessionId, message.content, message.attachments