UNPKG

@yepcode/mcp-server

Version:

MCP server for YepCode

287 lines (284 loc) 13 kB
import { Server } from "@modelcontextprotocol/sdk/server/index.js"; import { ListToolsRequestSchema, CallToolRequestSchema, ErrorCode, McpError, ListResourcesRequestSchema, ListResourceTemplatesRequestSchema, } from "@modelcontextprotocol/sdk/types.js"; import { zodToJsonSchema } from "zod-to-json-schema"; import { YepCodeEnv, YepCodeApiManager, YepCodeRun, Execution, } from "@yepcode/run"; import dotenv from "dotenv"; import { RunCodeSchema, buildRunCodeSchema, SetEnvVarSchema, RemoveEnvVarSchema, RunProcessSchema, GetExecutionSchema, } from "./types.js"; import { getVersion, isEmpty } from "./utils.js"; import Logger from "./logger.js"; const RUN_PROCESS_TOOL_NAME_PREFIX = "run_ycp_"; const RUN_PROCESS_TOOL_TAG = "mcp-tool"; dotenv.config(); class YepCodeMcpServer extends Server { constructor(config, { logsToStderr = false, disableRunCodeTool = false, skipRunCodeCleanup = false, } = {}) { super({ name: "yepcode-mcp-server", version: getVersion(), }, { capabilities: { tools: {}, resources: {}, resourceTemplates: {}, }, }); this.disableRunCodeTool = disableRunCodeTool; this.skipRunCodeCleanup = skipRunCodeCleanup; this.setupHandlers(); this.setupErrorHandling(); try { this.yepCodeRun = new YepCodeRun(config); this.yepCodeEnv = new YepCodeEnv(config); this.yepCodeApi = YepCodeApiManager.getInstance(config); this.logger = new Logger(this.yepCodeApi.getClientId(), { logsToStderr, }); this.logger.info("YepCode initialized successfully"); } catch (error) { this.logger = new Logger("YepCodeMcpServer", { logsToStderr, }); this.logger.error("Exception while initializing YepCode", error); throw new McpError(ErrorCode.InternalError, "Exception while initializing YepCode. Have you set the YEPCODE_API_TOKEN environment variable?"); } } setupErrorHandling() { this.onerror = (error) => { this.logger.error("[MCP Error]", error); }; process.on("SIGINT", async () => { this.logger.info("Received SIGINT, shutting down"); await this.close(); process.exit(0); }); } setupHandlers() { this.setupToolHandlers(); this.setupResourceHandlers(); this.setupTemplateHandlers(); } setupResourceHandlers() { this.setRequestHandler(ListResourcesRequestSchema, async () => ({ resources: [], // Empty list as we don't have any resources })); } setupTemplateHandlers() { this.setRequestHandler(ListResourceTemplatesRequestSchema, async () => ({ resourceTemplates: [], // Empty list as we don't have any templates })); } async handleToolRequest(schema, request, handler) { const parsed = schema.safeParse(request.params.arguments); if (!parsed.success) { this.logger.error("Invalid request arguments", parsed.error); throw new McpError(ErrorCode.InvalidParams, "Invalid request arguments"); } try { this.logger.info(`Handling tool request: ${request.params.name}`); const result = await handler(parsed.data); return { isError: false, content: [ { type: "text", text: JSON.stringify({ ...result, }, null, 2), }, ], }; } catch (error) { const errorMessage = error instanceof Error ? error.message : "Unknown error occurred"; this.logger.error(`Error in tool handler: ${request.params.name}`, error); return { isError: true, content: [ { type: "text", text: JSON.stringify({ error: errorMessage, }, null, 2), }, ], }; } } async executionResult(executionId) { const execution = new Execution({ yepCodeApi: this.yepCodeApi, executionId, }); await execution.waitForDone(); return { executionId, logs: execution.logs, processId: execution.processId, status: execution.status, timeline: execution.timeline, ...(execution.returnValue && { returnValue: execution.returnValue, }), ...(execution.error && { error: execution.error }), }; } setupToolHandlers() { this.setRequestHandler(ListToolsRequestSchema, async () => { this.logger.info(`Handling ListTools request`); const envVars = await this.yepCodeEnv.getEnvVars(); const tools = this.disableRunCodeTool ? [] : [ { name: "run_code", description: `Execute LLM-generated code safely in YepCode’s secure, production-grade sandboxes. This tool is ideal when your AI agent needs to handle tasks that don’t have a predefined tool available — but could be solved by writing and running a custom script. It supports external dependencies (NPM or PyPI), so it’s perfect for: • Complex data transformations • API calls to services not yet integrated • Custom logic implementations • One-off utility scripts Tip: First try to find a tool that matches your task, but if not available, try generating the code and running it here!`, inputSchema: zodToJsonSchema(buildRunCodeSchema(envVars.map((envVar) => envVar.key))), }, { name: "set_env_var", description: "Set a YepCode environment variable to be available for future code executions", inputSchema: zodToJsonSchema(SetEnvVarSchema), }, { name: "remove_env_var", description: "Remove a YepCode environment variable", inputSchema: zodToJsonSchema(RemoveEnvVarSchema), }, { name: "get_execution", description: "Get the status, result, logs, timeline, etc. of a YepCode execution", inputSchema: zodToJsonSchema(GetExecutionSchema), }, ]; let page = 0; let limit = 100; while (true) { const processes = await this.yepCodeApi.getProcesses({ page, limit }); this.logger.info(`Found ${processes?.data?.length} processes`); if (!processes.data) { break; } tools.push(...processes.data .filter((process) => process.tags?.includes(RUN_PROCESS_TOOL_TAG)) .map((process) => { const inputSchema = zodToJsonSchema(RunProcessSchema); if (!isEmpty(process.parametersSchema)) { inputSchema.properties.parameters = process.parametersSchema; } else { delete inputSchema.properties.parameters; } let toolName = `${RUN_PROCESS_TOOL_NAME_PREFIX}${process.slug}`; if (toolName.length > 60) { toolName = `${RUN_PROCESS_TOOL_NAME_PREFIX}${process.id}`; } return { name: toolName, description: `${process.name}${process.description ? ` - ${process.description}` : ""}`, inputSchema, }; })); if (!processes.hasNextPage) { break; } page++; } this.logger.info(`Found ${tools.length} tools: ${tools .map((tool) => tool.name) .join(", ")}`); return { tools, }; }); this.setRequestHandler(CallToolRequestSchema, async (request) => { this.logger.info(`Received CallTool request for: ${request.params.name}`); if (request.params.name.startsWith(RUN_PROCESS_TOOL_NAME_PREFIX)) { const processId = request.params.name.replace(RUN_PROCESS_TOOL_NAME_PREFIX, ""); return this.handleToolRequest(RunProcessSchema, request, async (data) => { const { synchronousExecution = true, parameters, ...options } = data; const { executionId } = await this.yepCodeApi.executeProcessAsync(processId, parameters, { ...options, initiatedBy: "@yepcode/mcp-server", }); if (!synchronousExecution) { return { executionId, }; } return await this.executionResult(executionId); }); } switch (request.params.name) { case "run_code": if (this.disableRunCodeTool) { this.logger.error("Run code tool is disabled"); throw new McpError(ErrorCode.MethodNotFound, "Run code tool is disabled"); } return this.handleToolRequest(RunCodeSchema, request, async (data) => { const { code, options } = data; const logs = []; let executionError; let returnValue; this.logger.info("Running code with YepCode", { codeLength: code.length, options, }); const execution = await this.yepCodeRun.run(code, { removeOnDone: !this.skipRunCodeCleanup, ...options, initiatedBy: "@yepcode/mcp-server", onLog: (log) => { logs.push(log); }, onError: (error) => { executionError = error.message; this.logger.error("YepCode execution error", error); }, onFinish: (value) => { returnValue = value; this.logger.info("YepCode execution finished", { hasReturnValue: value !== undefined, }); }, }); await execution.waitForDone(); return { logs, returnValue, ...(executionError && { error: executionError }), }; }); case "set_env_var": return this.handleToolRequest(SetEnvVarSchema, request, async (data) => { const { key, value, isSensitive } = data; this.logger.info(`Setting environment variable: ${key}`, { isSensitive, }); await this.yepCodeEnv.setEnvVar(key, value, isSensitive); return {}; }); case "remove_env_var": return this.handleToolRequest(RemoveEnvVarSchema, request, async (data) => { this.logger.info(`Removing environment variable: ${data.key}`); await this.yepCodeEnv.delEnvVar(data.key); return {}; }); case "get_execution": return this.handleToolRequest(GetExecutionSchema, request, async (data) => { return await this.executionResult(data.executionId); }); default: this.logger.error(`Unknown tool requested: ${request.params.name}`); throw new McpError(ErrorCode.MethodNotFound, `Unknown tool: ${request.params.name}`); } }); } } export default YepCodeMcpServer;