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(
'generate_image',
{
title: 'Generate Image',
description: 'Generate an image from a text prompt using Gemini image models (Nano Banana Pro by default).',
inputSchema: {
prompt: z.string().min(1).describe('Text description of the image to generate'),
model: ImageModel.default('gemini-3-pro-image-preview').describe('Image generation model (Nano Banana Pro by default)'),
aspectRatio: AspectRatio.default('1:1').describe('Aspect ratio of the generated image'),
imageSize: ImageSize.default('1K').describe('Image resolution (1K or 2K)'),
},
annotations: {
readOnlyHint: true,
destructiveHint: false,
openWorldHint: true,
},
},
async ({ prompt, model, aspectRatio, imageSize }) => {
try {
const response = await ai.models.generateContent({
model,
contents: prompt,
config: {
responseModalities: ['TEXT', 'IMAGE'],
imageConfig: { aspectRatio, imageSize },
},
});
const image = extractImageFromResponse(response);
if (!image) {
return {
content: [{ type: 'text' as const, text: 'No image was generated. Try a different prompt.' }],
isError: true,
};
}
if (!validateImageSize(image.data)) {
return {
content: [{ type: 'text' as const, text: 'Generated image exceeds size limit. Try 1K imageSize or a simpler prompt.' }],
isError: true,
};
}
return {
content: [{ type: 'image' as const, data: image.data, mimeType: image.mimeType }],
};
} catch (error) {
return formatToolError(error);
}
},
);
}
