Skip to content

Commit 70219f2

Browse files
committed
fix(core): sniff MCP image MIME types
1 parent 9e5599c commit 70219f2

4 files changed

Lines changed: 320 additions & 11 deletions

File tree

packages/core/src/tools/mcp-tool.test.ts

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -721,6 +721,89 @@ describe('DiscoveredMCPTool', () => {
721721
);
722722
});
723723

724+
it('should correct the MIME type of a mismatched image content block', async () => {
725+
const params = { param: 'imageMismatch' };
726+
const webpData = Buffer.from([
727+
0x52, 0x49, 0x46, 0x46,
728+
0x00, 0x00, 0x00, 0x00,
729+
0x57, 0x45, 0x42, 0x50,
730+
0x00, 0x00,
731+
]).toString('base64');
732+
733+
mockCallTool.mockResolvedValue(
734+
createSdkResponse(serverToolName, {
735+
content: [
736+
{
737+
type: 'image',
738+
data: webpData,
739+
mimeType: 'image/png',
740+
},
741+
],
742+
}),
743+
);
744+
745+
const invocation = tool.build(params);
746+
const toolResult = await invocation.execute({
747+
abortSignal: new AbortController().signal,
748+
});
749+
750+
expect(toolResult.llmContent).toEqual([
751+
{
752+
text: `[Tool '${serverToolName}' provided the following image data with mime-type: image/webp]`,
753+
},
754+
{
755+
inlineData: {
756+
mimeType: 'image/webp',
757+
data: webpData,
758+
},
759+
},
760+
]);
761+
expect(toolResult.returnDisplay).toBe('[Image: image/webp]');
762+
});
763+
764+
it('should correct the MIME type of a mismatched resource block blob', async () => {
765+
const params = { param: 'resourceMismatch' };
766+
const webpData = Buffer.from([
767+
0x52, 0x49, 0x46, 0x46,
768+
0x00, 0x00, 0x00, 0x00,
769+
0x57, 0x45, 0x42, 0x50,
770+
0x00, 0x00,
771+
]).toString('base64');
772+
773+
mockCallTool.mockResolvedValue(
774+
createSdkResponse(serverToolName, {
775+
content: [
776+
{
777+
type: 'resource',
778+
resource: {
779+
uri: 'file:///path/to/image.png',
780+
blob: webpData,
781+
mimeType: 'image/png',
782+
},
783+
},
784+
],
785+
}),
786+
);
787+
788+
const invocation = tool.build(params);
789+
const toolResult = await invocation.execute({
790+
abortSignal: new AbortController().signal,
791+
});
792+
793+
expect(toolResult.llmContent).toEqual([
794+
{
795+
text: `[Tool '${serverToolName}' provided the following embedded resource with mime-type: image/webp]`,
796+
},
797+
{
798+
inlineData: {
799+
mimeType: 'image/webp',
800+
data: webpData,
801+
},
802+
},
803+
]);
804+
expect(toolResult.returnDisplay).toBe('[Embedded Resource: image/webp]');
805+
});
806+
724807
describe('AbortSignal support', () => {
725808
const MOCK_TOOL_DELAY = 1000;
726809
const ABORT_DELAY = 50;

packages/core/src/tools/mcp-tool.ts

Lines changed: 28 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import type { MessageBus } from '../confirmation-bus/message-bus.js';
2424
import type { McpContext } from './mcp-client.js';
2525

2626
import { wrapUntrusted } from '../utils/textUtils.js';
27+
import { validateAndCorrectMimeType } from '../utils/imageMimeDetection.js';
2728

2829
/**
2930
* The separator used to qualify MCP tool names with their server prefix.
@@ -457,15 +458,19 @@ function transformImageAudioBlock(
457458
block: McpMediaBlock,
458459
toolName: string,
459460
): Part[] {
461+
const correctedMimeType =
462+
block.type === 'image'
463+
? validateAndCorrectMimeType(block.mimeType, block.data)
464+
: block.mimeType;
460465
return [
461466
{
462467
text: `[Tool '${toolName}' provided the following ${
463468
block.type
464-
} data with mime-type: ${block.mimeType}]`,
469+
} data with mime-type: ${correctedMimeType}]`,
465470
},
466471
{
467472
inlineData: {
468-
mimeType: block.mimeType,
473+
mimeType: correctedMimeType,
469474
data: block.data,
470475
},
471476
},
@@ -481,14 +486,18 @@ function transformResourceBlock(
481486
return { text: wrapUntrusted(resource.text) };
482487
}
483488
if (resource?.blob) {
484-
const mimeType = resource.mimeType || 'application/octet-stream';
489+
const declaredMimeType = resource.mimeType || 'application/octet-stream';
490+
const correctedMimeType = validateAndCorrectMimeType(
491+
declaredMimeType,
492+
resource.blob,
493+
);
485494
return [
486495
{
487-
text: `[Tool '${toolName}' provided the following embedded resource with mime-type: ${mimeType}]`,
496+
text: `[Tool '${toolName}' provided the following embedded resource with mime-type: ${correctedMimeType}]`,
488497
},
489498
{
490499
inlineData: {
491-
mimeType,
500+
mimeType: correctedMimeType,
492501
data: resource.blob,
493502
},
494503
},
@@ -562,19 +571,27 @@ function getStringifiedResultForDisplay(rawResponse: Part[]): string {
562571
switch (block.type) {
563572
case 'text':
564573
return block.text;
565-
case 'image':
566-
return `[Image: ${block.mimeType}]`;
574+
case 'image': {
575+
const correctedMimeType = validateAndCorrectMimeType(
576+
block.mimeType,
577+
block.data,
578+
);
579+
return `[Image: ${correctedMimeType}]`;
580+
}
567581
case 'audio':
568582
return `[Audio: ${block.mimeType}]`;
569583
case 'resource_link':
570584
return `[Link to ${block.title || block.name}: ${block.uri}]`;
571-
case 'resource':
585+
case 'resource': {
572586
if (block.resource?.text) {
573587
return block.resource.text;
574588
}
575-
return `[Embedded Resource: ${
576-
block.resource?.mimeType || 'unknown type'
577-
}]`;
589+
const declaredMimeType = block.resource?.mimeType || 'unknown type';
590+
const correctedMimeType = block.resource?.blob
591+
? validateAndCorrectMimeType(declaredMimeType, block.resource.blob)
592+
: declaredMimeType;
593+
return `[Embedded Resource: ${correctedMimeType}]`;
594+
}
578595
default:
579596
return `[Unknown content type: ${(block as { type: string }).type}]`;
580597
}
Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
/**
2+
* @license
3+
* Copyright 2026 Google LLC
4+
* SPDX-License-Identifier: Apache-2.0
5+
*/
6+
7+
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
8+
import {
9+
detectImageMimeType,
10+
validateAndCorrectMimeType,
11+
} from './imageMimeDetection.js';
12+
import { debugLogger } from './debugLogger.js';
13+
14+
describe('imageMimeDetection', () => {
15+
// Base64 helper strings representing magic bytes of various formats
16+
const pngBase64 = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0x00, 0x00]).toString('base64');
17+
const webpBase64 = Buffer.from([
18+
0x52, 0x49, 0x46, 0x46, // RIFF
19+
0x00, 0x00, 0x00, 0x00, // Size
20+
0x57, 0x45, 0x42, 0x50, // WEBP
21+
0x00, 0x00
22+
]).toString('base64');
23+
const jpegBase64 = Buffer.from([0xff, 0xd8, 0xff, 0xe0, 0x00, 0x00]).toString('base64');
24+
const gifBase64 = Buffer.from([0x47, 0x49, 0x46, 0x38, 0x39, 0x61, 0x00, 0x00]).toString('base64');
25+
const txtBase64 = Buffer.from('Hello world!').toString('base64');
26+
27+
describe('detectImageMimeType', () => {
28+
it('should correctly detect PNG files', () => {
29+
expect(detectImageMimeType(pngBase64)).toBe('image/png');
30+
});
31+
32+
it('should correctly detect WebP files', () => {
33+
expect(detectImageMimeType(webpBase64)).toBe('image/webp');
34+
});
35+
36+
it('should correctly detect JPEG files', () => {
37+
expect(detectImageMimeType(jpegBase64)).toBe('image/jpeg');
38+
});
39+
40+
it('should correctly detect GIF files', () => {
41+
expect(detectImageMimeType(gifBase64)).toBe('image/gif');
42+
});
43+
44+
it('should return null for non-image / unknown files', () => {
45+
expect(detectImageMimeType(txtBase64)).toBeNull();
46+
});
47+
48+
it('should return null for empty or invalid base64 data', () => {
49+
expect(detectImageMimeType('')).toBeNull();
50+
expect(detectImageMimeType('abc')).toBeNull(); // too short to decode 4 bytes
51+
});
52+
53+
it('should handle whitespace/newlines in base64 data', () => {
54+
const formattedPng = pngBase64.match(/.{1,4}/g)?.join('\n') || pngBase64;
55+
expect(detectImageMimeType(formattedPng)).toBe('image/png');
56+
});
57+
58+
it('should handle leading whitespace/newlines in base64 data', () => {
59+
expect(detectImageMimeType(' \n\r ' + pngBase64)).toBe('image/png');
60+
});
61+
});
62+
63+
describe('validateAndCorrectMimeType', () => {
64+
let warnSpy: ReturnType<typeof vi.spyOn>;
65+
66+
beforeEach(() => {
67+
warnSpy = vi.spyOn(debugLogger, 'warn').mockImplementation(() => {});
68+
});
69+
70+
afterEach(() => {
71+
vi.restoreAllMocks();
72+
});
73+
74+
it('should return original mime type if there is no mismatch', () => {
75+
expect(validateAndCorrectMimeType('image/png', pngBase64)).toBe('image/png');
76+
expect(warnSpy).not.toHaveBeenCalled();
77+
});
78+
79+
it('should return corrected mime type and log a warning if there is a mismatch', () => {
80+
// e.g. Figma WebP labeled as PNG
81+
expect(validateAndCorrectMimeType('image/png', webpBase64)).toBe('image/webp');
82+
expect(warnSpy).toHaveBeenCalledTimes(1);
83+
expect(warnSpy.mock.calls[0][0]).toContain(
84+
'Image MIME type mismatch: declared as image/png but detected as image/webp'
85+
);
86+
});
87+
88+
it('should fallback to declared mime type if image format is not detected', () => {
89+
expect(validateAndCorrectMimeType('application/octet-stream', txtBase64)).toBe(
90+
'application/octet-stream'
91+
);
92+
expect(warnSpy).not.toHaveBeenCalled();
93+
});
94+
95+
it('should NOT correct mime type if the declared type is not an image or generic', () => {
96+
const textStartingWithGif = Buffer.from('GIF is an image format, not a text file').toString('base64');
97+
expect(validateAndCorrectMimeType('text/plain', textStartingWithGif)).toBe('text/plain');
98+
expect(warnSpy).not.toHaveBeenCalled();
99+
});
100+
});
101+
});
Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
/**
2+
* @license
3+
* Copyright 2026 Google LLC
4+
* SPDX-License-Identifier: Apache-2.0
5+
*/
6+
7+
import { debugLogger } from './debugLogger.js';
8+
9+
/**
10+
* Detects the actual image MIME type by inspecting the magic bytes (file signature)
11+
* of base64 encoded data. Supports PNG, WebP, JPEG, and GIF.
12+
*
13+
* @param base64Data The base64 encoded string of the image.
14+
* @returns The detected MIME type (e.g. 'image/webp') or null if not detected/unsupported.
15+
*/
16+
export function detectImageMimeType(base64Data: string): string | null {
17+
if (!base64Data) {
18+
return null;
19+
}
20+
try {
21+
// Take a small prefix of the base64 data and strip whitespace.
22+
// 48 characters of base64 yields up to 36 bytes of binary data.
23+
const cleanPrefix = base64Data.trimStart().slice(0, 48).replace(/\s+/g, '');
24+
const buffer = Buffer.from(cleanPrefix, 'base64');
25+
if (buffer.length < 4) {
26+
return null;
27+
}
28+
29+
// PNG: 89 50 4E 47
30+
if (
31+
buffer[0] === 0x89 &&
32+
buffer[1] === 0x50 &&
33+
buffer[2] === 0x4e &&
34+
buffer[3] === 0x47
35+
) {
36+
return 'image/png';
37+
}
38+
39+
// WebP: RIFF (bytes 0-3) and WEBP (bytes 8-11)
40+
if (
41+
buffer.length >= 12 &&
42+
buffer[0] === 0x52 && // R
43+
buffer[1] === 0x49 && // I
44+
buffer[2] === 0x46 && // F
45+
buffer[3] === 0x46 && // F
46+
buffer[8] === 0x57 && // W
47+
buffer[9] === 0x45 && // E
48+
buffer[10] === 0x42 && // B
49+
buffer[11] === 0x50 // P
50+
) {
51+
return 'image/webp';
52+
}
53+
54+
// JPEG: FF D8 FF
55+
if (
56+
buffer[0] === 0xff &&
57+
buffer[1] === 0xd8 &&
58+
buffer[2] === 0xff
59+
) {
60+
return 'image/jpeg';
61+
}
62+
63+
// GIF: GIF (47 49 46)
64+
if (
65+
buffer[0] === 0x47 && // G
66+
buffer[1] === 0x49 && // I
67+
buffer[2] === 0x46 // F
68+
) {
69+
return 'image/gif';
70+
}
71+
} catch {
72+
// Return null on any decode errors
73+
}
74+
return null;
75+
}
76+
77+
/**
78+
* Validates the declared MIME type against the actual format detected from the base64 data.
79+
* If a mismatch is detected, it logs a warning and returns the corrected MIME type.
80+
* Otherwise, it returns the declared MIME type.
81+
*
82+
* @param declaredMimeType The declared MIME type.
83+
* @param base64Data The base64 encoded string of the image.
84+
* @returns The corrected or original MIME type.
85+
*/
86+
export function validateAndCorrectMimeType(
87+
declaredMimeType: string,
88+
base64Data: string,
89+
): string {
90+
const lowerDeclared = declaredMimeType.toLowerCase();
91+
const isImageOrGeneric =
92+
lowerDeclared.startsWith('image/') ||
93+
lowerDeclared === 'application/octet-stream' ||
94+
lowerDeclared === 'unknown type';
95+
96+
if (!isImageOrGeneric) {
97+
return declaredMimeType;
98+
}
99+
100+
const detectedType = detectImageMimeType(base64Data);
101+
if (detectedType && detectedType !== declaredMimeType) {
102+
debugLogger.warn(
103+
`Image MIME type mismatch: declared as ${declaredMimeType} but detected as ${detectedType}`,
104+
);
105+
return detectedType;
106+
}
107+
return declaredMimeType;
108+
}

0 commit comments

Comments
 (0)