import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import type { GoogleGenAI } from '@google/genai';
import { z } from 'zod';
import { ImageModel, AspectRatio, ImageSize } from '../types.js';
import { formatToolError } from '../utils/errors.js';
import { validateImageSize, extractImageFromResponse } from '../utils/image.js';
export function register(server: McpServer, ai: GoogleGenAI): void {
server.registerTool(
'edit_image',
{
title: 'Edit Image',
description: 'Edit an image using a text prompt. Send a base64-encoded image and describe the desired changes.',
inputSchema: {
prompt: z.string().min(1).describe('Description of the edits to apply'),
image: z.string().min(1).describe('Base64-encoded source image'),
mimeType: z.enum(['image/png', 'image/jpeg', 'image/webp']).default('image/png').describe('MIME type of the source image'),
model: ImageModel.default('gemini-3-pro-image-preview').describe('Image model to use (Nano Banana Pro by default)'),
aspectRatio: AspectRatio.default('1:1').describe('Aspect ratio of the output image'),
imageSize: ImageSize.default('1K').describe('Output image resolution'),
},
annotations: {
readOnlyHint: true,
destructiveHint: false,
openWorldHint: true,
},
},
async ({ prompt, image, mimeType, model, aspectRatio, imageSize }) => {
try {
const response = await ai.models.generateContent({
model,
contents: [
{ text: prompt },
{ inlineData: { mimeType, data: image } },
],
config: {
responseModalities: ['TEXT', 'IMAGE'],
imageConfig: { aspectRatio, imageSize },
},
});
const result = extractImageFromResponse(response);
if (!result) {
return {
content: [{ type: 'text' as const, text: 'No edited image was produced. Try a different prompt.' }],
isError: true,
};
}
if (!validateImageSize(result.data)) {
return {
content: [{ type: 'text' as const, text: 'Edited image exceeds size limit. Try a smaller imageSize.' }],
isError: true,
};
}
return {
content: [{ type: 'image' as const, data: result.data, mimeType: result.mimeType }],
};
} catch (error) {
return formatToolError(error);
}
},
);
}
export function registerMulti(server: McpServer, ai: GoogleGenAI): void {
server.registerTool(
'edit_image_multi',
{
title: 'Edit Image (Multi-Reference)',
description: 'Edit or compose images using multiple reference images (up to 14). Uses gemini-3-pro-image-preview (Nano Banana Pro).',
inputSchema: {
prompt: z.string().min(1).describe('Description of the desired output'),
images: z.array(z.object({
data: z.string().min(1).describe('Base64-encoded image'),
mimeType: z.enum(['image/png', 'image/jpeg', 'image/webp']).default('image/png'),
})).min(1).max(14).describe('Array of reference images (1-14)'),
aspectRatio: AspectRatio.default('1:1').describe('Aspect ratio of the output image'),
imageSize: ImageSize.default('1K').describe('Output image resolution'),
},
annotations: {
readOnlyHint: true,
destructiveHint: false,
openWorldHint: true,
},
},
async ({ prompt, images, aspectRatio, imageSize }) => {
try {
const parts = [
{ text: prompt },
...images.map((img: { mimeType: string; data: string }) => ({
inlineData: { mimeType: img.mimeType, data: img.data },
})),
];
const response = await ai.models.generateContent({
model: 'gemini-3-pro-image-preview',
contents: parts,
config: {
responseModalities: ['TEXT', 'IMAGE'],
imageConfig: { aspectRatio, imageSize },
},
});
const result = extractImageFromResponse(response);
if (!result) {
return {
content: [{ type: 'text' as const, text: 'No image was produced. Try a different prompt.' }],
isError: true,
};
}
if (!validateImageSize(result.data)) {
return {
content: [{ type: 'text' as const, text: 'Output image exceeds size limit. Try a smaller imageSize.' }],
isError: true,
};
}
return {
content: [{ type: 'image' as const, data: result.data, mimeType: result.mimeType }],
};
} catch (error) {
return formatToolError(error);
}
},
);
}
