import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import type { GoogleGenAI } from '@google/genai';
import { z } from 'zod';
import { TextModel } from '../types.js';
import { formatToolError } from '../utils/errors.js';
export function register(server: McpServer, ai: GoogleGenAI): void {
server.registerTool(
'code_execution',
{
title: 'Code Execution',
description: 'Execute Python code in a sandboxed environment. Gemini generates and runs code, returning both the code and results.',
inputSchema: {
prompt: z.string().min(1).describe('Describe what code to write and execute'),
model: TextModel.default('gemini-2.5-flash').describe('Gemini model to use'),
temperature: z.number().min(0).max(2).optional().describe('Sampling temperature'),
maxOutputTokens: z.number().min(1).optional().describe('Maximum output tokens'),
},
annotations: {
readOnlyHint: true,
destructiveHint: false,
openWorldHint: true,
},
},
async ({ prompt, model, temperature, maxOutputTokens }) => {
try {
const response = await ai.models.generateContent({
model,
contents: prompt,
config: {
temperature,
maxOutputTokens,
tools: [{ codeExecution: {} }],
},
});
const parts = response.candidates?.[0]?.content?.parts ?? [];
const sections: string[] = [];
for (const part of parts) {
if (part.text) {
sections.push(part.text);
} else if (part.executableCode) {
sections.push(`\`\`\`python\n${part.executableCode.code}\n\`\`\``);
} else if (part.codeExecutionResult) {
const outcome = part.codeExecutionResult.outcome ?? 'UNKNOWN';
sections.push(`**Execution result** (${outcome}):\n\`\`\`\n${part.codeExecutionResult.output}\n\`\`\``);
}
}
return {
content: [{ type: 'text' as const, text: sections.join('\n\n') }],
};
} catch (error) {
return formatToolError(error);
}
},
);
}
