DB Agent MCP Server

// MCP SDK関連 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js"; import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js"; // HTTPサーバー関連 import express, { type Request, type Response } from "express"; import cors from "cors"; // バリデーション・設定 import { z } from "zod"; import dotenv from "dotenv"; // ツール import { runQueryReadonlyTool, runQueryWriteTool } from "./tools/run_query"; // ミドルウェア import { authMiddleware, AUTH_CONFIG } from "./middleware/auth"; // 環境変数の読み込み dotenv.config(); // ============================================================================ // 設定 // ============================================================================ const SERVER_CONFIG = { name: "db-investigator", version: "0.0.1", port: process.env.PORT ? parseInt(process.env.PORT) : 3000, } as const; // MCPサーバーの初期化 const mcpServer = new McpServer({ name: SERVER_CONFIG.name, version: SERVER_CONFIG.version, }); // 読み取り専用ツールを登録（承認不要） mcpServer.registerTool( runQueryReadonlyTool.name, { description: runQueryReadonlyTool.description, inputSchema: z.object({ sql: z.string().describe("Read-only SQL query to execute (SELECT, SHOW, DESCRIBE, DESC, EXPLAIN only)"), }), }, async ({ sql }: { sql: string }) => { const result = await runQueryReadonlyTool.handler({ sql }); return { content: [ { type: "text" as const, text: JSON.stringify(result, null, 2), }, ], }; } ); // 書き込みツールを登録（承認必要） mcpServer.registerTool( runQueryWriteTool.name, { description: runQueryWriteTool.description, inputSchema: z.object({ sql: z.string().describe("Write SQL query to execute (INSERT, UPDATE, DELETE, CREATE, DROP, ALTER, etc.)"), }), annotations: { destructiveHint: true, readOnlyHint: false, }, }, async ({ sql }: { sql: string }) => { const result = await runQueryWriteTool.handler({ sql }); return { content: [ { type: "text" as const, text: JSON.stringify(result, null, 2), }, ], }; } ); // ============================================================================ // セッション管理 // ============================================================================ // セッションIDごとのtransportを保存 const transports = new Map<string, StreamableHTTPServerTransport>(); // Expressアプリケーションの作成 const app = express(); // CORSミドルウェアの設定 app.use(cors({ origin: "*", methods: ["GET", "POST", "DELETE", "OPTIONS"], allowedHeaders: ["Content-Type", "Authorization", "mcp-session-id", "X-API-Key"], })); // JSONボディパーサーの設定 app.use(express.json()); // MCPエンドポイントのハンドラー（共通処理） async function handleMcpRequest(req: Request, res: Response) { try { const parsedBody = req.body; const sessionId = req.headers["mcp-session-id"] as string; let transport: StreamableHTTPServerTransport | undefined; if (sessionId && transports.has(sessionId)) { // 既存のセッションを使用 transport = transports.get(sessionId); } else if (!sessionId && parsedBody && isInitializeRequest(parsedBody)) { // 新しい初期化リクエスト transport = new StreamableHTTPServerTransport({ sessionIdGenerator: () => crypto.randomUUID(), onsessioninitialized: (sid) => { console.log(`Session initialized with ID: ${sid}`); transports.set(sid, transport!); }, }); // セッション終了時のクリーンアップ transport.onclose = () => { const sid = transport!.sessionId; if (sid && transports.has(sid)) { console.log(`Transport closed for session ${sid}`); transports.delete(sid); } }; // MCPサーバーに接続 await mcpServer.connect(transport); await transport.handleRequest(req, res, parsedBody); return; } else { // 無効なリクエスト res.status(400).json({ jsonrpc: "2.0", error: { code: -32000, message: "Bad Request: No valid session ID provided", }, id: null, }); return; } // 既存のtransportでリクエストを処理 if (transport) { await transport.handleRequest(req, res, parsedBody); } } catch (error: any) { console.error("Request error:", error); if (!res.headersSent) { res.status(500).json({ error: error.message }); } } } // GETエンドポイント（SSEストリーム用）- 認証適用 app.get("/mcp", authMiddleware, handleMcpRequest); // POSTエンドポイント（JSON-RPCリクエスト用）- 認証適用 app.post("/mcp", authMiddleware, handleMcpRequest); // 404ハンドラー app.use((req: Request, res: Response) => { res.status(404).json({ error: "Not Found" }); }); // エラーハンドラー app.use((err: any, req: Request, res: Response, next: any) => { console.error("Server error:", err); if (!res.headersSent) { res.status(500).json({ error: err.message }); } }); // サーバーの起動 app.listen(SERVER_CONFIG.port, () => { console.log(`MCP Server running on http://localhost:${SERVER_CONFIG.port}/mcp`); if (AUTH_CONFIG.enabled) { console.log(`Authentication: ENABLED`); } else { console.log(`Authentication: DISABLED (for development only)`); } });