diff --git a/src/server/routes/mcp_route.ts b/src/server/routes/mcp_route.ts index 459ec30..7434f8a 100644 --- a/src/server/routes/mcp_route.ts +++ b/src/server/routes/mcp_route.ts @@ -3,18 +3,6 @@ import { v4 as uuidv4 } from "uuid"; const API_KEY = process.env.MCP_API_KEY ?? "super-secret-key"; const PORT = Number(process.env.PORT ?? 3000); -const PING_INTERVAL_MS = 25_000; - -// ===================== -// Store session & clients -// ===================== -type Client = { - id: string; - send: (data: string) => void; - close: () => void; -}; - -const sessions = new Map>(); // ===================== // Helper Functions @@ -28,48 +16,18 @@ function isAuthorized(headers: Headers) { return headers.get("x-api-key") === API_KEY; } -function formatSSE(event: string, data: any, id?: string) { - const payload = typeof data === "string" ? data : JSON.stringify(data); - const lines: string[] = []; - - if (id) lines.push(`id: ${id}`); - if (event) lines.push(`event: ${event}`); - - // Split data into multiple data: lines if needed - payload.split("\n").forEach(line => { - lines.push(`data: ${line}`); - }); - - lines.push(""); // Empty line to end the message - return lines.join("\n") + "\n"; -} - -function broadcast(sessionId: string, event: string, data: any) { - const clients = sessions.get(sessionId); - if (!clients) return 0; - const messageId = uuidv4(); - const message = formatSSE(event, data, messageId); - - for (const client of clients) { - try { - client.send(message); - } catch { - clients.delete(client); - } - } - return clients.size; -} - // ===================== // Tools Definition // ===================== type Tool = { name: string; description: string; - inputSchema?: { + inputSchema: { type: string; - properties?: Record; + properties: Record; required?: string[]; + additionalProperties?: boolean; + $schema?: string; }; run: (input?: any) => Promise; }; @@ -81,6 +39,8 @@ const tools: Tool[] = [ inputSchema: { type: "object", properties: {}, + additionalProperties: true, + $schema: "http://json-schema.org/draft-07/schema#", }, run: async () => ({ pong: Date.now() }), }, @@ -90,6 +50,8 @@ const tools: Tool[] = [ inputSchema: { type: "object", properties: {}, + additionalProperties: true, + $schema: "http://json-schema.org/draft-07/schema#", }, run: async () => ({ uuid: uuidv4() }), }, @@ -99,265 +61,299 @@ const tools: Tool[] = [ inputSchema: { type: "object", properties: { - message: { + input: { type: "string", description: "Message to echo back", }, }, + required: ["input"], + additionalProperties: true, + $schema: "http://json-schema.org/draft-07/schema#", }, run: async (input) => ({ echo: input }), }, + { + name: "Calculator", + description: "Useful for getting the result of a math expression. The input to this tool should be a valid mathematical expression that could be executed by a simple calculator.", + inputSchema: { + type: "object", + properties: { + input: { + type: "string", + }, + }, + required: ["input"], + additionalProperties: true, + $schema: "http://json-schema.org/draft-07/schema#", + }, + run: async (input) => { + try { + // Simple math evaluation (be careful in production!) + const result = Function(`"use strict"; return (${input.input})`)(); + return { result: String(result) }; + } catch (error: any) { + throw new Error(`Invalid expression: ${error.message}`); + } + }, + }, ]; +// ===================== +// MCP Protocol Types +// ===================== +type JSONRPCRequest = { + jsonrpc: "2.0"; + id: string | number; + method: string; + params?: any; +}; + +type JSONRPCResponse = { + jsonrpc: "2.0"; + id: string | number; + result?: any; + error?: { + code: number; + message: string; + data?: any; + }; +}; + +type JSONRPCNotification = { + jsonrpc: "2.0"; + method: string; + params?: any; +}; + +// ===================== +// MCP Handler +// ===================== +function handleMCPRequest(request: JSONRPCRequest): JSONRPCResponse { + const { id, method, params } = request; + + switch (method) { + case "initialize": + return { + jsonrpc: "2.0", + id, + result: { + protocolVersion: "2024-11-05", + capabilities: { + tools: {}, + }, + serverInfo: { + name: "elysia-mcp-server", + version: "1.0.0", + }, + }, + }; + + case "tools/list": + return { + jsonrpc: "2.0", + id, + result: { + tools: tools.map(({ name, description, inputSchema }) => ({ + name, + description, + inputSchema, + })), + }, + }; + + case "tools/call": + const toolName = params?.name; + const tool = tools.find((t) => t.name === toolName); + + if (!tool) { + return { + jsonrpc: "2.0", + id, + error: { + code: -32601, + message: `Tool '${toolName}' not found`, + }, + }; + } + + try { + // Note: This is synchronous for simplicity + // In real implementation, you'd need to handle async properly + let result: any; + tool.run(params?.arguments || {}).then((r) => (result = r)); + + return { + jsonrpc: "2.0", + id, + result: { + content: [ + { + type: "text", + text: JSON.stringify(result || { pending: true }), + }, + ], + }, + }; + } catch (error: any) { + return { + jsonrpc: "2.0", + id, + error: { + code: -32603, + message: error.message, + }, + }; + } + + case "ping": + return { + jsonrpc: "2.0", + id, + result: {}, + }; + + default: + return { + jsonrpc: "2.0", + id, + error: { + code: -32601, + message: `Method '${method}' not found`, + }, + }; + } +} + +async function handleMCPRequestAsync(request: JSONRPCRequest): Promise { + const { id, method, params } = request; + + if (method === "tools/call") { + const toolName = params?.name; + const tool = tools.find((t) => t.name === toolName); + + if (!tool) { + return { + jsonrpc: "2.0", + id, + error: { + code: -32601, + message: `Tool '${toolName}' not found`, + }, + }; + } + + try { + const result = await tool.run(params?.arguments || {}); + return { + jsonrpc: "2.0", + id, + result: { + content: [ + { + type: "text", + text: JSON.stringify(result), + }, + ], + }, + }; + } catch (error: any) { + return { + jsonrpc: "2.0", + id, + error: { + code: -32603, + message: error.message, + }, + }; + } + } + + // For other methods, use sync handler + return handleMCPRequest(request); +} + // ===================== // Server Initialization // ===================== export const MCPRoute = new Elysia() // ===================== - // SSE Stream - // ===================== - .get("/mcp/:sessionId", ({ params, set, request }) => { - const { sessionId } = params; - - // Check authorization for SSE connection - if (!isAuthorized(request.headers)) { - set.status = 401; - return { error: "Unauthorized" }; - } - - set.headers["Content-Type"] = "text/event-stream"; - set.headers["Cache-Control"] = "no-cache, no-transform"; - set.headers["Connection"] = "keep-alive"; - set.headers["Access-Control-Allow-Origin"] = "*"; - set.headers["X-Accel-Buffering"] = "no"; - - const stream = new TransformStream(); - const writer = stream.writable.getWriter(); - const encoder = new TextEncoder(); - - const client: Client = { - id: uuidv4(), - send: (data) => { - try { - writer.write(encoder.encode(data)); - } catch (e) { - console.error("Error writing to stream:", e); - } - }, - close: () => { - try { - writer.close(); - } catch (e) { - // Stream already closed - } - const set = sessions.get(sessionId); - if (set) { - set.delete(client); - if (set.size === 0) sessions.delete(sessionId); - } - }, - }; - - if (!sessions.has(sessionId)) sessions.set(sessionId, new Set()); - sessions.get(sessionId)!.add(client); - - // Send initial connection message - client.send(formatSSE("connected", { - sessionId, - clientId: client.id, - timestamp: Date.now() - })); - - // Send tools list on connection - client.send(formatSSE("tools", { - tools: tools.map(({ name, description, inputSchema }) => ({ - name, - description, - inputSchema - })) - })); - - // Setup ping interval - const ping = setInterval(() => { - try { - client.send(formatSSE("ping", { ts: Date.now() })); - } catch (e) { - clearInterval(ping); - } - }, PING_INTERVAL_MS); - - const readable = stream.readable; - - // Cleanup on connection close - readable.pipeTo(new WritableStream()).catch(() => { - clearInterval(ping); - client.close(); - }); - - return new Response(readable, { - headers: { - "Content-Type": "text/event-stream", - "Cache-Control": "no-cache, no-transform", - "Connection": "keep-alive", - "Access-Control-Allow-Origin": "*", - "X-Accel-Buffering": "no", - }, - status: 200, - }); - }) - - // ===================== - // MCP Session Status - // ===================== - .get("/mcp/:sessionId/status", ({ params, set }) => { - set.headers["Access-Control-Allow-Origin"] = "*"; - const clients = sessions.get(params.sessionId); - return { - sessionId: params.sessionId, - connected: clients?.size ?? 0, - clients: Array.from(clients || []).map(c => c.id), - }; - }) - - // ===================== - // MCP Broadcast + // MCP HTTP Streamable Endpoint // ===================== .post("/mcp/:sessionId", async ({ params, request, set }) => { - set.headers["Access-Control-Allow-Origin"] = "*"; - if (!isAuthorized(request.headers)) { - set.status = 401; - return { error: "Unauthorized" }; - } - - const body = await request.json(); - const event = body.event ?? "message"; - const data = body.data ?? body; - - const sentTo = broadcast(params.sessionId, event, data); - return { ok: true, sentTo }; - }) - - // ===================== - // Delete /mcp/:sessionId - // ===================== - .delete("/mcp/:sessionId", ({ params, request, set }) => { - set.headers["Access-Control-Allow-Origin"] = "*"; - if (!isAuthorized(request.headers)) { - set.status = 401; - return { error: "Unauthorized" }; - } - - const clients = sessions.get(params.sessionId); - if (clients) { - for (const c of clients) c.close(); - sessions.delete(params.sessionId); - } - return { ok: true }; - }) - - // ===================== - // Tools Introspection (Fixed path) - // ===================== - .get("/mcp/:sessionId/tools", ({ params, set, request }) => { + set.headers["Content-Type"] = "application/json"; set.headers["Access-Control-Allow-Origin"] = "*"; - // Optional: Check auth if needed + // Optional: Check authorization // if (!isAuthorized(request.headers)) { // set.status = 401; // return { error: "Unauthorized" }; // } - return { - tools: tools.map(({ name, description, inputSchema }) => ({ - name, - description, - inputSchema - })) - }; + try { + const body = await request.json(); + + // Handle single request + if (!Array.isArray(body)) { + const response = await handleMCPRequestAsync(body as JSONRPCRequest); + return response; + } + + // Handle batch requests + const responses = await Promise.all( + body.map((req) => handleMCPRequestAsync(req as JSONRPCRequest)) + ); + return responses; + } catch (error: any) { + set.status = 400; + return { + jsonrpc: "2.0", + id: null, + error: { + code: -32700, + message: "Parse error", + data: error.message, + }, + }; + } }) - // Alternative global tools endpoint - .get("/tools", ({ set }) => { + // ===================== + // Simple tools list endpoint (for debugging) + // ===================== + .get("/mcp/:sessionId/tools", ({ set }) => { set.headers["Access-Control-Allow-Origin"] = "*"; return { - tools: tools.map(({ name, description, inputSchema }) => ({ + data: tools.map(({ name, description, inputSchema }) => ({ name, + value: name, description, - inputSchema - })) + inputSchema, + })), }; }) // ===================== - // Run Tool (Fixed path) + // Session Status // ===================== - .post("/mcp/:sessionId/tools/:toolName", async ({ params, request, set }) => { + .get("/mcp/:sessionId/status", ({ params, set }) => { set.headers["Access-Control-Allow-Origin"] = "*"; - if (!isAuthorized(request.headers)) { - set.status = 401; - return { error: "Unauthorized" }; - } - - const tool = tools.find((t) => t.name === params.toolName); - if (!tool) { - set.status = 404; - return { error: `Tool '${params.toolName}' not found` }; - } - - try { - const body = await request.json().catch(() => ({})); - const result = await tool.run(body); - - // Broadcast tool execution result to session - broadcast(params.sessionId, "tool_result", { - tool: tool.name, - result, - timestamp: Date.now(), - }); - - return { - ok: true, - tool: tool.name, - result - }; - } catch (error: any) { - set.status = 500; - return { - error: "Tool execution failed", - message: error.message - }; - } + return { + sessionId: params.sessionId, + status: "active", + timestamp: Date.now(), + }; }) - // Alternative global tool execution - .post("/tools/:toolName", async ({ params, request, set }) => { + // ===================== + // Health Check + // ===================== + .get("/health", ({ set }) => { set.headers["Access-Control-Allow-Origin"] = "*"; - if (!isAuthorized(request.headers)) { - set.status = 401; - return { error: "Unauthorized" }; - } - - const tool = tools.find((t) => t.name === params.toolName); - if (!tool) { - set.status = 404; - return { error: `Tool '${params.toolName}' not found` }; - } - - try { - const body = await request.json().catch(() => ({})); - const result = await tool.run(body); - return { - ok: true, - tool: tool.name, - result - }; - } catch (error: any) { - set.status = 500; - return { - error: "Tool execution failed", - message: error.message - }; - } + return { + status: "ok", + timestamp: Date.now(), + tools: tools.length, + }; }) // ===================== @@ -365,24 +361,16 @@ export const MCPRoute = new Elysia() // ===================== .options("/mcp/:sessionId", ({ set }) => { set.headers["Access-Control-Allow-Origin"] = "*"; - set.headers["Access-Control-Allow-Methods"] = "GET,POST,DELETE,OPTIONS"; - set.headers["Access-Control-Allow-Headers"] = "Content-Type,X-API-Key,Authorization"; + set.headers["Access-Control-Allow-Methods"] = "GET,POST,OPTIONS"; + set.headers["Access-Control-Allow-Headers"] = "Content-Type,Authorization,X-API-Key"; set.status = 204; return ""; }) .options("/mcp/:sessionId/tools", ({ set }) => { set.headers["Access-Control-Allow-Origin"] = "*"; - set.headers["Access-Control-Allow-Methods"] = "GET,POST,OPTIONS"; - set.headers["Access-Control-Allow-Headers"] = "Content-Type,X-API-Key,Authorization"; - set.status = 204; - return ""; - }) - - .options("/mcp/:sessionId/tools/:toolName", ({ set }) => { - set.headers["Access-Control-Allow-Origin"] = "*"; - set.headers["Access-Control-Allow-Methods"] = "POST,OPTIONS"; - set.headers["Access-Control-Allow-Headers"] = "Content-Type,X-API-Key,Authorization"; + set.headers["Access-Control-Allow-Methods"] = "GET,OPTIONS"; + set.headers["Access-Control-Allow-Headers"] = "Content-Type,Authorization,X-API-Key"; set.status = 204; return ""; }); \ No newline at end of file