Add MCP server

packages/mcp/eslint.config.mjs

@@ -0,0 +1,20 @@
+import { base } from '@internal/eslint-config';
+import path from 'node:path';
+import { fileURLToPath } from 'node:url';
+
+const __dirname = path.dirname(fileURLToPath(import.meta.url));
+
+export default [
+  {
+    ignores: ['eslint.config.mjs']
+  },
+  ...base,
+  {
+    languageOptions: {
+      parserOptions: {
+        projectService: true,
+        tsconfigRootDir: __dirname
+      }
+    }
+  }
+];

packages/mcp/package.json

@@ -0,0 +1,27 @@
+{
+  "name": "@dbagent/mcp",
+  "version": "0.0.1",
+  "private": true,
+  "exports": {
+    ".": "./src/index.ts"
+  },
+  "scripts": {
+    "lint": "eslint ./src",
+    "tsc": "tsc --noEmit"
+  },
+  "dependencies": {
+    "fastmcp": "^1.18.0",
+    "pg": "8.13.3",
+    "typescript": "^5.7.3",
+    "zod": "^3.24.2"
+  },
+  "devDependencies": {
+    "@internal/tsconfig": "workspace:*",
+    "@types/node": "^22.13.5",
+    "@types/pg": "^8.11.11"
+  },
+  "peerDependencies": {
+    "@modelcontextprotocol/sdk": "^1.5.0",
+    "ai": "^4.1.45"
+  }
+}

packages/mcp/src/index.ts

@@ -0,0 +1 @@
+export { createToolSet } from './utils/ai-sdk';

packages/mcp/src/servers/postgres/index.ts

@@ -0,0 +1,83 @@
+#!/usr/bin/env node
+
+import { FastMCP } from 'fastmcp';
+import pg from 'pg';
+import { z } from 'zod';
+
+async function buildServer(databaseUrl: string) {
+  const server = new FastMCP({
+    name: 'PostgreSQL',
+    version: '0.1.0'
+  });
+
+  const resourceBaseUrl = new URL(databaseUrl);
+  resourceBaseUrl.protocol = 'postgres:';
+  resourceBaseUrl.password = '';
+
+  const client = new pg.Pool({
+    connectionString: databaseUrl
+  });
+
+  server.addTool({
+    name: 'get-slow-queries',
+    description: `Get a list of slow queries formatted as a JSON array.
+            Contains how many times the query was called, the max execution time in seconds,
+            the mean execution time in seconds, the total execution time (all calls together)
+            in seconds, and the query itself.`,
+    parameters: z.object({
+      thresholdMs: z.number().optional().default(5000)
+    }),
+    execute: async ({ thresholdMs }) => {
+      const result = await client.query(
+        `
+                SELECT 
+                calls,
+                round(max_exec_time/1000) max_exec_secs,
+                round(mean_exec_time/1000) mean_exec_secs, 
+                round(total_exec_time/1000) total_exec_secs,
+                query 
+                FROM pg_stat_statements 
+                WHERE max_exec_time > $1 
+                ORDER BY total_exec_time DESC 
+                LIMIT 10;`,
+        [thresholdMs]
+      );
+
+      return JSON.stringify(result.rows);
+    }
+  });
+
+  server.addTool({
+    name: 'explain-query',
+    description: `Run explain on a a query. Returns the explain plan as received from PostgreSQL.
+            The query needs to be complete, it cannot contain $1, $2, etc. If you need to, replace the parameters with your own made up values.
+            It's very important that $1, $2, etc. are not passed to this tool. Use the tool describeTable to get the types of the columns.
+            If you know the schema, pass it in as well.`,
+    parameters: z.object({
+      schema: z.string(),
+      query: z.string()
+    }),
+    execute: async ({ schema, query }) => {
+      if (query.includes('$1') || query.includes('$2') || query.includes('$3') || query.includes('$4')) {
+        return 'The query seems to contain placeholders ($1, $2, etc). Replace them with actual values and try again.';
+      }
+      await client.query('BEGIN');
+      try {
+        await client.query(`SET search_path TO ${schema}`);
+        const explainQuery = `EXPLAIN ${query}`;
+        console.log(schema);
+        console.log(explainQuery);
+        const result = await client.query(explainQuery);
+        console.log(result.rows);
+        return result.rows.map((row: { [key: string]: string }) => row['QUERY PLAN']).join('\n');
+      } catch (error) {
+        console.error('Error explaining query', error);
+        return 'I could not run EXPLAIN on that query. Try a different method.';
+      } finally {
+        await client.query('ROLLBACK');
+      }
+    }
+  });
+
+  return server;
+}

packages/mcp/src/utils/ai-sdk.ts

@@ -0,0 +1,72 @@
+import { Client } from '@modelcontextprotocol/sdk/client/index.js';
+import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
+import { jsonSchema, type Tool } from 'ai';
+
+type ToolSetConfig = {
+  mcpServers: {
+    [key: string]: {
+      command: string;
+      args: string[];
+      env?: Record<string, string>;
+    };
+  };
+  onToolCall?: <Result>(serverName: string, toolName: string, args: any, result: Result) => void;
+};
+
+type ToolSet = {
+  tools: {
+    [key: string]: Tool;
+  };
+  clients: {
+    [key: string]: Client;
+  };
+};
+
+export async function createToolSet(config: ToolSetConfig): Promise<ToolSet> {
+  const toolset: ToolSet = {
+    tools: {},
+    clients: {}
+  };
+
+  for (const [serverName, serverConfig] of Object.entries(config.mcpServers)) {
+    const transport = new StdioClientTransport({
+      ...serverConfig,
+      stderr: process.stderr
+    });
+
+    const client = new Client(
+      {
+        name: `${serverName}-client`,
+        version: '1.0.0'
+      },
+      {
+        capabilities: {}
+      }
+    );
+
+    toolset.clients[serverName] = client;
+    await client.connect(transport);
+
+    const { tools } = await client.listTools();
+    for (const tool of tools) {
+      const toolName = tool.name === serverName ? tool.name : `${serverName}_${tool.name}`;
+
+      toolset.tools[toolName] = {
+        description: tool.description || '',
+        parameters: jsonSchema(tool.inputSchema as any),
+        execute: async (args) => {
+          const result = await client.callTool({
+            name: tool.name,
+            arguments: args
+          });
+
+          config.onToolCall?.(serverName, toolName, args, result);
+
+          return JSON.stringify(result);
+        }
+      };
+    }
+  }
+
+  return toolset;
+}

packages/mcp/tsconfig.json

@@ -0,0 +1,6 @@
+{
+  "$schema": "https://json.schemastore.org/tsconfig",
+  "extends": "@internal/tsconfig/base.json",
+  "include": ["."],
+  "exclude": ["dist", "build", "node_modules"]
+}