diff --git a/changelog.d/fix-image-media-type-sniff.fixed.md b/changelog.d/fix-image-media-type-sniff.fixed.md new file mode 100644 index 00000000..c2f6ced6 --- /dev/null +++ b/changelog.d/fix-image-media-type-sniff.fixed.md @@ -0,0 +1,4 @@ +- Image blocks whose declared `media_type` disagrees with their bytes (e.g. a + PNG labeled `image/webp` by an upstream surface) are corrected from the + magic bytes at compile time. The provider rejects any request with such a + mismatch, so one mislabeled image used to fail every subsequent turn. diff --git a/src/agent.ts b/src/agent.ts index 0855aafd..fc5c1eca 100644 --- a/src/agent.ts +++ b/src/agent.ts @@ -1,5 +1,6 @@ import type { Membrane, NormalizedMessage, NormalizedRequest, ContentBlock, YieldingStream } from '@animalabs/membrane'; import { isAbortedResponse } from '@animalabs/membrane'; +import { correctImageMediaTypes } from './image-media-type.js'; import { createHash } from 'node:crypto'; import type { CacheWireReceipt, KvUnifiedRequestHooks } from './kv-unified-wire.js'; import { @@ -525,7 +526,7 @@ export class Agent { async compileContext(budget?: TokenBudget): Promise { const result = await this.contextManager.compile(this.resolveBudget(budget)); if (!budget) this.settleRuntimeSettingsTransition(); - return result; + return { ...result, messages: correctImageMediaTypes(result.messages) }; } /** @@ -542,7 +543,7 @@ export class Agent { this.resolveBudget(budget), injections, opts as never, ); if (!budget) this.settleRuntimeSettingsTransition(); - return result; + return { ...result, messages: correctImageMediaTypes(result.messages) }; } // ========================================================================== diff --git a/src/image-media-type.ts b/src/image-media-type.ts new file mode 100644 index 00000000..ef24ddb8 --- /dev/null +++ b/src/image-media-type.ts @@ -0,0 +1,55 @@ +import type { ContentBlock, NormalizedMessage } from '@animalabs/membrane'; + +/** Detect an image's real format from its leading bytes. */ +export function sniffImageMediaTypeBase64(data: string): string | undefined { + const head = Buffer.from(data.slice(0, 32), 'base64'); + if (head.length >= 4 && head[0] === 0x89 && head[1] === 0x50 && head[2] === 0x4e && head[3] === 0x47) return 'image/png'; + if (head.length >= 3 && head[0] === 0xff && head[1] === 0xd8 && head[2] === 0xff) return 'image/jpeg'; + if (head.length >= 3 && head[0] === 0x47 && head[1] === 0x49 && head[2] === 0x46) return 'image/gif'; + if ( + head.length >= 12 && + head[0] === 0x52 && head[1] === 0x49 && head[2] === 0x46 && head[3] === 0x46 && + head[8] === 0x57 && head[9] === 0x45 && head[10] === 0x42 && head[11] === 0x50 + ) return 'image/webp'; + return undefined; +} + +/** + * Correct base64 image blocks whose declared media type doesn't match their + * bytes (e.g. Discord serving a PNG for an attachment it labels image/webp). + * The provider rejects the WHOLE request on a mismatch, so one mislabelled + * image in history would fail every turn. Returns new message/block objects + * for anything changed; stored history is never touched. + */ +export function correctImageMediaTypes(messages: NormalizedMessage[]): NormalizedMessage[] { + let changed = 0; + const fix = (blocks: ContentBlock[]): ContentBlock[] => { + let touched = false; + const out = blocks.map((block) => { + const b = block as { type: string; source?: { type?: string; data?: string; mediaType?: string }; content?: unknown }; + if (b.type === 'image' && b.source?.type === 'base64' && typeof b.source.data === 'string') { + const actual = sniffImageMediaTypeBase64(b.source.data); + if (actual && actual !== b.source.mediaType) { + touched = true; + changed++; + return { ...block, source: { ...b.source, mediaType: actual } } as ContentBlock; + } + } + if (b.type === 'tool_result' && Array.isArray(b.content)) { + const inner = fix(b.content as ContentBlock[]); + if (inner !== b.content) { + touched = true; + return { ...block, content: inner } as ContentBlock; + } + } + return block; + }); + return touched ? out : blocks; + }; + const result = messages.map((m) => { + const content = fix(m.content); + return content === m.content ? m : { ...m, content }; + }); + if (changed > 0) console.error(`[images] corrected media type on ${changed} image block(s) whose bytes disagreed with their label`); + return result; +} diff --git a/test/image-media-type.test.ts b/test/image-media-type.test.ts new file mode 100644 index 00000000..cf665b9b --- /dev/null +++ b/test/image-media-type.test.ts @@ -0,0 +1,25 @@ +import { test } from 'node:test'; +import assert from 'node:assert/strict'; +import { correctImageMediaTypes } from '../src/image-media-type.js'; + +const PNG_B64 = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4////fwAJ+wP9KobjigAAAABJRU5ErkJggg=='; + +test('corrects a PNG labelled image/webp, including inside tool_result, without mutating input', () => { + const img = { type: 'image', source: { type: 'base64', data: PNG_B64, mediaType: 'image/webp' } }; + const messages = [ + { participant: 'user', content: [{ type: 'text', text: 'look' }, img] }, + { participant: 'user', content: [{ type: 'tool_result', toolUseId: 't', content: [img] }] }, + ] as never; + const out = correctImageMediaTypes(messages) as unknown as Array<{ content: any[] }>; + assert.equal(out[0].content[1].source.mediaType, 'image/png'); + assert.equal(out[1].content[0].content[0].source.mediaType, 'image/png'); + assert.equal(img.source.mediaType, 'image/webp', 'stored block untouched'); +}); + +test('leaves correct and unknown images alone (same objects)', () => { + const messages = [{ participant: 'user', content: [ + { type: 'image', source: { type: 'base64', data: PNG_B64, mediaType: 'image/png' } }, + { type: 'image', source: { type: 'base64', data: 'AAAA', mediaType: 'image/png' } }, + ] }] as never; + assert.equal(correctImageMediaTypes(messages)[0], (messages as any)[0]); +});