feat: mcp pool (#1)

.env.example

@@ -1,16 +0,0 @@
-# Nostr
-PRIVATE_KEY=your_private_key_here
-RELAY_URLS=wss://relay.damus.io,wss://relay.nostr.band,wss://nos.lol
-
-# MCP Service Info
-MCP_SERVICE_NAME="DVM MCP Bridge"
-MCP_SERVICE_ABOUT="MCP-enabled DVM providing AI and computational tools"
-
-# MCP Client Connection
-MCP_CLIENT_NAME="DVM MCP Bridge Client"
-MCP_CLIENT_VERSION="1.0.0"
-MCP_SERVER_COMMAND=bun  # The command to start the MCP server
-MCP_SERVER_ARGS=run,src/external-mcp-server.ts  # Comma-separated args to pass to the server command
-
-# Optional: Comma-separated list of allowed pubkeys. If empty, allows all.
-ALLOWED_PUBKEYS=

.gitignore

@@ -178,3 +178,4 @@ dist
 codebase*
 .aider*
 .prettierrc
+config.yml

bun.lock

@@ -9,6 +9,7 @@
         "dotenv": "^16.4.7",
         "nostr-tools": "^2.10.4",
         "prettier": "^3.4.2",
+        "yaml": "^2.7.0",
       },
       "devDependencies": {
         "@types/bun": "latest",
@@ -81,6 +82,8 @@
 
     "unpipe": ["unpipe@1.0.0", "", {}, "sha512-pjy2bYhSsufwWlKwPc+l3cN7+wuJlK6uz0YdJEOlQDbl6jo/YlPi4mb8agUkVC8BF7V8NuzeyPNqRksA3hztKQ=="],
 
+    "yaml": ["yaml@2.7.0", "", { "bin": { "yaml": "bin.mjs" } }, "sha512-+hSoy/QHluxmC9kCIJyL/uyFmLmc+e5CFR5Wa+bpIhIj85LVb9ZH2nVnqrHoSvKogwODv0ClqZkmiSSaIH5LTA=="],
+
     "zod": ["zod@3.24.1", "", {}, "sha512-muH7gBL9sI1nciMZV67X5fTKKBLtwpZ5VBp1vsOQzj1MhrBZ4wlVCm3gedKZWLp0Oyel8sIGfeiz54Su+OVT+A=="],
 
     "zod-to-json-schema": ["zod-to-json-schema@3.24.1", "", { "peerDependencies": { "zod": "^3.24.1" } }, "sha512-3h08nf3Vw3Wl3PK+q3ow/lIil81IT2Oa7YpQyUUDsEWbXveMesdfK1xBd2RhCkynwZndAxixji/7SYJJowr62w=="],

config.example.yml

@@ -0,0 +1,35 @@
+# Nostr Configuration
+nostr:
+  # Your private key in hex format (32-byte hex string)
+  privateKey: "your_private_key_here"
+  # List of relay URLs (must start with ws:// or wss://)
+  relayUrls:
+    - "wss://relay1.com"
+    - "wss://relay2.net"
+
+# MCP Service Configuration
+mcp:
+  # Service information
+  name: "DVM MCP Bridge"
+  about: "MCP-enabled DVM providing AI and computational tools"
+  # Required client information
+  clientName: "DVM MCP Bridge Client"
+  clientVersion: "1.0.0"
+  
+  # MCP Servers Configuration accepts multiple servers
+  servers:
+    - name: "server1"
+      command: "node"
+      args: 
+        - "run"
+        - "src/external-mcp-server1.ts"
+    
+    - name: "server2"
+      command: "python"
+      args:
+        - "src/external-mcp-server2.py"
+
+# Optional: Whitelist Configuration
+# whitelist:
+  # List of allowed public keys (leave empty for no restrictions)
+  # allowedPubkeys: []

package.json

@@ -8,7 +8,8 @@
     "dev": "bun --watch src/dvm-bridge.ts",
     "start": "bun run src/dvm-bridge.ts",
     "typecheck": "tsc --noEmit",
-    "lint": "bun run typecheck && bun run format"
+    "lint": "bun run typecheck && bun run format",
+    "test": "bun test"
   },
   "devDependencies": {
     "@types/bun": "latest"
@@ -21,6 +22,7 @@
     "@noble/hashes": "^1.7.1",
     "dotenv": "^16.4.7",
     "nostr-tools": "^2.10.4",
-    "prettier": "^3.4.2"
+    "prettier": "^3.4.2",
+    "yaml": "^2.7.0"
   }
 }

src/config.ts

@@ -1,6 +1,7 @@
-import { config } from 'dotenv';
+import { parse } from 'yaml';
 import { join } from 'path';
-import { existsSync } from 'fs';
+import { existsSync, readFileSync } from 'fs';
+import type { MCPServerConfig } from './types';
 
 interface NostrConfig {
   privateKey: string;
@@ -12,8 +13,7 @@ interface MCPConfig {
   about: string;
   clientName: string;
   clientVersion: string;
-  serverCommand: string;
-  serverArgs: string[];
+  servers: MCPServerConfig[];
 }
 
 interface WhitelistConfig {
@@ -26,75 +26,120 @@ interface AppConfig {
   whitelist: WhitelistConfig;
 }
 
-const envPath = join(process.cwd(), '.env');
-if (!existsSync(envPath)) {
-  throw new Error(
-    'No .env file found. Please create one based on .env.example'
-  );
-}
-
-const result = config();
-
-if (result.error) {
-  throw new Error(`Error loading .env file: ${result.error.message}`);
-}
-
+const CONFIG_PATH = join(process.cwd(), 'config.yml');
 const HEX_KEYS_REGEX = /^(?:[0-9a-fA-F]{64})$/;
 
-function requireEnvVar(name: string): string {
-  const value = process.env[name];
-  if (!value) {
+if (!existsSync(CONFIG_PATH)) {
   throw new Error(
-      `Missing required environment variable: ${name}. Check your .env file`
+    'No config.yml file found. Please create one based on config.example.yml'
   );
+}
+
+function loadConfig(): AppConfig {
+  try {
+    const configFile = readFileSync(CONFIG_PATH, 'utf8');
+    const rawConfig = parse(configFile);
+
+    const config: AppConfig = {
+      nostr: {
+        privateKey: validateRequiredField(
+          rawConfig.nostr?.privateKey,
+          'nostr.privateKey'
+        ),
+        relayUrls: validateRelayUrls(rawConfig.nostr?.relayUrls),
+      },
+      mcp: {
+        name: getConfigValue(rawConfig.mcp?.name, 'DVM MCP Bridge'),
+        about: getConfigValue(
+          rawConfig.mcp?.about,
+          'MCP-enabled DVM providing AI and computational tools'
+        ),
+        clientName: validateRequiredField(
+          rawConfig.mcp?.clientName,
+          'mcp.clientName'
+        ),
+        clientVersion: validateRequiredField(
+          rawConfig.mcp?.clientVersion,
+          'mcp.clientVersion'
+        ),
+        servers: validateMCPServers(rawConfig.mcp?.servers),
+      },
+      whitelist: {
+        allowedPubkeys: rawConfig.whitelist?.allowedPubkeys
+          ? new Set(
+              rawConfig.whitelist.allowedPubkeys.map((pk: string) => pk.trim())
+            )
+          : undefined,
+      },
+    };
+
+    if (!HEX_KEYS_REGEX.test(config.nostr.privateKey)) {
+      throw new Error('privateKey must be a 32-byte hex string');
+    }
+
+    return config;
+  } catch (error) {
+    throw new Error(`Failed to load config: ${error}`);
+  }
+}
+
+function validateRequiredField(value: any, fieldName: string): string {
+  if (!value) {
+    throw new Error(`Missing required config field: ${fieldName}`);
   }
   return value;
 }
 
-function getEnvVar(name: string, defaultValue: string): string {
-  return process.env[name] || defaultValue;
+function getConfigValue(
+  value: string | undefined,
+  defaultValue: string
+): string {
+  return value || defaultValue;
 }
 
-export const CONFIG: AppConfig = {
-  nostr: {
-    privateKey: requireEnvVar('PRIVATE_KEY'),
-    relayUrls: requireEnvVar('RELAY_URLS')
-      .split(',')
-      .map((url) => url.trim()),
-  },
-  mcp: {
-    name: getEnvVar('MCP_SERVICE_NAME', 'DVM MCP Bridge'),
-    about: getEnvVar(
-      'MCP_SERVICE_ABOUT',
-      'MCP-enabled DVM providing AI and computational tools'
-    ),
-    clientName: requireEnvVar('MCP_CLIENT_NAME'),
-    clientVersion: requireEnvVar('MCP_CLIENT_VERSION'),
-    serverCommand: requireEnvVar('MCP_SERVER_COMMAND'),
-    serverArgs: requireEnvVar('MCP_SERVER_ARGS').split(','),
-  },
-  whitelist: {
-    allowedPubkeys: process.env.ALLOWED_PUBKEYS
-      ? new Set(process.env.ALLOWED_PUBKEYS.split(',').map((pk) => pk.trim()))
-      : undefined,
-  },
-};
-
-if (!HEX_KEYS_REGEX.test(CONFIG.nostr.privateKey)) {
-  throw new Error('PRIVATE_KEY must be a 32-byte hex string');
-}
-
-CONFIG.nostr.relayUrls.forEach((url) => {
-  try {
-    new URL(url);
-    if (!url.startsWith('ws://') && !url.startsWith('wss://')) {
-      throw new Error(`Relay URL must start with ws:// or wss://: ${url}`);
+function validateRelayUrls(urls: any): string[] {
+  if (!Array.isArray(urls) || urls.length === 0) {
+    throw new Error(
+      'At least one relay URL must be provided in nostr.relayUrls'
+    );
   }
+
+  return urls.map((url: string) => {
+    try {
+      const trimmedUrl = url.trim();
+      new URL(trimmedUrl);
+      if (!trimmedUrl.startsWith('ws://') && !trimmedUrl.startsWith('wss://')) {
+        throw new Error(
+          `Relay URL must start with ws:// or wss://: ${trimmedUrl}`
+        );
+      }
+      return trimmedUrl;
     } catch (error) {
       throw new Error(`Invalid relay URL: ${url}`);
     }
-});
-
-if (CONFIG.nostr.relayUrls.length === 0) {
-  throw new Error('At least one relay URL must be provided in RELAY_URLS');
+  });
 }
+
+function validateMCPServers(servers: any): MCPServerConfig[] {
+  if (!Array.isArray(servers) || servers.length === 0) {
+    throw new Error(
+      'At least one MCP server must be configured in mcp.servers'
+    );
+  }
+
+  return servers.map((server: any, index: number) => {
+    if (!server.name || !server.command || !Array.isArray(server.args)) {
+      throw new Error(
+        `Invalid MCP server configuration at index ${index}. Required fields: name, command, args[]`
+      );
+    }
+
+    return {
+      name: server.name,
+      command: server.command,
+      args: server.args,
+    };
+  });
+}
+
+export const CONFIG = loadConfig();

src/dvm-bridge.ts

@@ -1,22 +1,22 @@
-import { MCPClientHandler } from './mcp-client';
 import { NostrAnnouncer } from './nostr/announcer';
 import { RelayHandler } from './nostr/relay';
 import { keyManager } from './nostr/keys';
 import relayHandler from './nostr/relay';
 import type { Event } from 'nostr-tools/pure';
 import { CONFIG } from './config';
+import { MCPPool } from './mcp-pool';
 
 export class DVMBridge {
-  private mcpClient: MCPClientHandler;
+  private mcpPool: MCPPool;
   private nostrAnnouncer: NostrAnnouncer;
   private relayHandler: RelayHandler;
   private isRunning: boolean = false;
 
   constructor() {
     console.log('Initializing DVM Bridge...');
-    this.mcpClient = new MCPClientHandler();
+    this.mcpPool = new MCPPool(CONFIG.mcp.servers);
     this.relayHandler = relayHandler;
-    this.nostrAnnouncer = new NostrAnnouncer(this.mcpClient);
+    this.nostrAnnouncer = new NostrAnnouncer(this.mcpPool);
   }
 
   private isWhitelisted(pubkey: string): boolean {
@@ -31,12 +31,13 @@ export class DVMBridge {
       console.log('Bridge is already running');
       return;
     }
-    try {
-      console.log('Connecting to MCP server...');
-      await this.mcpClient.connect();
 
-      const tools = await this.mcpClient.listTools();
-      console.log('Available MCP tools:', tools);
+    try {
+      console.log('Connecting to MCP servers...');
+      await this.mcpPool.connect();
+
+      const tools = await this.mcpPool.listTools();
+      console.log(`Available MCP tools across all servers: ${tools.length}`);
 
       console.log('Announcing service to Nostr network...');
       await this.nostrAnnouncer.updateAnnouncement();
@@ -59,7 +60,7 @@ export class DVMBridge {
 
     console.log('Stopping DVM Bridge...');
     try {
-      await this.mcpClient.disconnect();
+      await this.mcpPool.disconnect();
       this.relayHandler.cleanup();
       this.isRunning = false;
       console.log('DVM Bridge stopped successfully');
@@ -76,7 +77,7 @@ export class DVMBridge {
           const command = event.tags.find((tag) => tag[0] === 'c')?.[1];
 
           if (command === 'list-tools') {
-            const tools = await this.mcpClient.listTools();
+            const tools = await this.mcpPool.listTools();
             const response = keyManager.signEvent({
               ...keyManager.createEventTemplate(6910),
               content: JSON.stringify({
@@ -103,7 +104,7 @@ export class DVMBridge {
             await this.relayHandler.publishEvent(processingStatus);
 
             try {
-              const result = await this.mcpClient.callTool(
+              const result = await this.mcpPool.callTool(
                 jobRequest.name,
                 jobRequest.parameters
               );

src/mcp-client.ts

@@ -1,17 +1,17 @@
 import { Client } from '@modelcontextprotocol/sdk/client/index.js';
 import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
 import { CONFIG } from './config';
-// TODO: Add connection to multiple mcp servers and router
+import type { MCPServerConfig } from './types';
+
 export class MCPClientHandler {
   private client: Client;
   private transport: StdioClientTransport;
 
-  constructor() {
+  constructor(config: MCPServerConfig) {
     this.transport = new StdioClientTransport({
-      command: CONFIG.mcp.serverCommand,
-      args: CONFIG.mcp.serverArgs,
+      command: config.command,
+      args: config.args,
     });
-
     this.client = new Client(
       {
         name: CONFIG.mcp.clientName,

src/mcp-pool.ts

@@ -0,0 +1,47 @@
+import type { Tool } from '@modelcontextprotocol/sdk/types.js';
+import { MCPClientHandler } from './mcp-client';
+import type { MCPServerConfig } from './types';
+
+export class MCPPool {
+  private clients: Map<string, MCPClientHandler> = new Map();
+  private toolRegistry: Map<string, MCPClientHandler> = new Map();
+
+  constructor(serverConfigs: MCPServerConfig[]) {
+    serverConfigs.forEach((config) => {
+      const client = new MCPClientHandler(config);
+      this.clients.set(config.name, client);
+    });
+  }
+
+  async connect() {
+    await Promise.all(
+      Array.from(this.clients.values()).map((client) => client.connect())
+    );
+  }
+
+  async listTools(): Promise<Tool[]> {
+    const allTools: Tool[] = [];
+    for (const client of this.clients.values()) {
+      const tools = await client.listTools();
+      tools.forEach((tool) => {
+        this.toolRegistry.set(tool.name, client);
+      });
+      allTools.push(...tools);
+    }
+    return allTools;
+  }
+
+  async callTool(name: string, args: Record<string, any>) {
+    const client = this.toolRegistry.get(name);
+    if (!client) {
+      throw new Error(`No MCP server found for tool: ${name}`);
+    }
+    return await client.callTool(name, args);
+  }
+
+  async disconnect() {
+    await Promise.all(
+      Array.from(this.clients.values()).map((client) => client.disconnect())
+    );
+  }
+}

src/nostr/announcer.ts

@@ -1,16 +1,16 @@
 import { CONFIG } from '../config';
 import { RelayHandler } from './relay';
 import { keyManager } from './keys';
-import { MCPClientHandler } from '../mcp-client';
 import relayHandler from './relay';
+import type { MCPPool } from '../mcp-pool';
 
 export class NostrAnnouncer {
   private relayHandler: RelayHandler;
-  private mcpClient: MCPClientHandler;
+  private mcpPool: MCPPool;
 
-  constructor(mcpClient: MCPClientHandler) {
+  constructor(mcpPool: MCPPool) {
     this.relayHandler = relayHandler;
-    this.mcpClient = mcpClient;
+    this.mcpPool = mcpPool;
   }
 
   async announceRelayList() {
@@ -25,24 +25,24 @@ export class NostrAnnouncer {
   }
 
   async announceService() {
-    const toolsResult = await this.mcpClient.listTools();
+    const tools = await this.mcpPool.listTools();
     const event = keyManager.signEvent({
       ...keyManager.createEventTemplate(31990),
       content: JSON.stringify({
         name: CONFIG.mcp.name,
         about: CONFIG.mcp.about,
-        tools: toolsResult,
+        tools: tools,
       }),
       tags: [
         ['d', 'dvm-announcement'],
         ['k', '5910'],
         ['capabilities', 'mcp-1.0'],
         ['t', 'mcp'],
-        ...toolsResult.map((tool) => ['t', tool.name]),
+        ...tools.map((tool) => ['t', tool.name]),
       ],
     });
     await this.relayHandler.publishEvent(event);
-    console.log(`Announced service with ${toolsResult.length} tools`);
+    console.log(`Announced service with ${tools.length} tools`);
   }
 
   async updateAnnouncement() {

src/nostr/keys.ts

@@ -1,10 +1,8 @@
 import { hexToBytes } from '@noble/hashes/utils';
 import { getPublicKey, finalizeEvent } from 'nostr-tools/pure';
-import type { Event } from 'nostr-tools/pure';
+import type { Event, UnsignedEvent } from 'nostr-tools/pure';
 import { CONFIG } from '../config';
 
-export type UnsignedEvent = Omit<Event, 'sig' | 'id'>;
-
 export const createKeyManager = (privateKeyHex: string) => {
   const privateKeyBytes = hexToBytes(privateKeyHex);
   const pubkey = getPublicKey(privateKeyBytes);

src/tests/keys.test.ts

@@ -0,0 +1,25 @@
+import { expect, test, describe } from 'bun:test';
+import { createKeyManager } from '../nostr/keys';
+
+describe('KeyManager', () => {
+  const testPrivateKey =
+    'd4d4d7aae7857054596c4c0976b22a73acac3a10d30bf56db35ee038bbf0dd44';
+  const keyManager = createKeyManager(testPrivateKey);
+
+  test('should create valid event template', () => {
+    const template = keyManager.createEventTemplate(1);
+    expect(template.kind).toBe(1);
+    expect(template.pubkey).toBeDefined();
+    expect(template.created_at).toBeNumber();
+    expect(template.tags).toBeArray();
+    expect(template.content).toBe('');
+  });
+
+  test('should sign events correctly', () => {
+    const template = keyManager.createEventTemplate(1);
+    const signedEvent = keyManager.signEvent(template);
+    expect(signedEvent.id).toBeDefined();
+    expect(signedEvent.sig).toBeDefined();
+    expect(signedEvent.pubkey).toBe(keyManager.pubkey);
+  });
+});

src/tests/mcp-pool.test.ts

@@ -0,0 +1,48 @@
+import { expect, test, describe, beforeAll, afterAll } from 'bun:test';
+import { MCPPool } from '../mcp-pool';
+import { join } from 'path';
+import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js';
+import { createMockServer } from './mock-server';
+
+describe('MCPPool', () => {
+  let mcpPool: MCPPool;
+  let transports: any[] = [];
+  const serverNames = ['server1', 'server2', 'server3', 'server4'];
+  beforeAll(async () => {
+    const mockServerPath = join(import.meta.dir, 'mock-server.ts');
+
+    const serverConfigs = serverNames.map((name) => ({
+      name,
+      command: 'bun',
+      args: ['run', mockServerPath, name],
+    }));
+
+    const servers = await Promise.all(
+      serverNames.map((name) => createMockServer(name))
+    );
+    transports = servers.map((s) => s.transport);
+
+    mcpPool = new MCPPool(serverConfigs);
+    await mcpPool.connect();
+  });
+
+  afterAll(async () => {
+    await mcpPool.disconnect();
+    await Promise.all(transports.map((t) => t.close()));
+  });
+
+  test('should list tools from all servers', async () => {
+    const tools = await mcpPool.listTools();
+    expect(tools.length).toEqual(serverNames.length);
+    expect(tools.map((t) => t.name).sort()).toEqual(
+      serverNames.map((name) => `${name}-echo`).sort()
+    );
+  });
+
+  test('should call tool on correct server', async () => {
+    const result = (await mcpPool.callTool('server1-echo', {
+      text: 'test message',
+    })) as CallToolResult;
+    expect(result.content[0].text).toBe('[server1] test message');
+  });
+});

src/tests/mock-server.ts

@@ -0,0 +1,30 @@
+import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
+import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
+import { z } from 'zod';
+
+export const createMockServer = async (name: string) => {
+  const server = new McpServer({
+    name: `Mock ${name}`,
+    version: '1.0.0',
+  });
+
+  server.tool(
+    `${name}-echo`,
+    `Echo tool for ${name}`,
+    {
+      text: z.string(),
+    },
+    async ({ text }) => ({
+      content: [{ type: 'text' as const, text: `[${name}] ${text}` }],
+    })
+  );
+
+  const transport = new StdioServerTransport();
+  await server.connect(transport);
+
+  return { server, transport };
+};
+
+if (import.meta.path === Bun.main) {
+  await createMockServer(process.argv[2] || 'default');
+}

src/types.ts

@@ -0,0 +1,5 @@
+export interface MCPServerConfig {
+  name: string;
+  command: string;
+  args: string[];
+}