@retrivora-ai/rag-engine 1.9.7 → 1.9.8

package/src/core/mcp.ts

@@ -0,0 +1,261 @@
+import { spawn, ChildProcess } from 'child_process';
+import { MCPServerConfig } from '../config/RagConfig';
+
+interface MCPTool {
+  name: string;
+  description?: string;
+  inputSchema: any;
+}
+
+interface MCPCallResult {
+  content: Array<{ type: string; text?: string; [key: string]: any }>;
+  isError?: boolean;
+}
+
+export class MCPClient {
+  private childProcess?: ChildProcess;
+  private sseUrl?: string;
+  private messageIdCounter = 0;
+  private pendingRequests = new Map<number, { resolve: (val: any) => void; reject: (err: any) => void }>();
+  private stdioBuffer = '';
+  private initialized = false;
+
+  constructor(private config: MCPServerConfig) {}
+
+  private getNextMessageId(): number {
+    return ++this.messageIdCounter;
+  }
+
+  /**
+   * Connect and initialize the MCP server connection.
+   */
+  async connect(): Promise<void> {
+    if (this.initialized) return;
+
+    if (this.config.transport === 'stdio') {
+      await this.connectStdio();
+    } else {
+      await this.connectSSE();
+    }
+
+    // Call initialize
+    try {
+      const response = await this.sendJsonRpc('initialize', {
+        protocolVersion: '2024-11-05',
+        capabilities: {},
+        clientInfo: {
+          name: 'retrivora-sdk',
+          version: '1.0.0',
+        },
+      });
+
+      // Send initialized notification
+      await this.sendNotification('notifications/initialized');
+      this.initialized = true;
+      console.log(`[MCPClient] Connected and initialized server "${this.config.name}" successfully.`);
+    } catch (err: any) {
+      this.disconnect();
+      throw new Error(`[MCPClient] Failed to initialize server "${this.config.name}": ${err.message}`);
+    }
+  }
+
+  private async connectStdio(): Promise<void> {
+    const cmd = this.config.command;
+    if (!cmd) {
+      throw new Error(`[MCPClient] Command option is required for stdio transport on "${this.config.name}"`);
+    }
+
+    const args = this.config.args || [];
+    const env = { ...process.env, ...(this.config.env || {}) };
+
+    this.childProcess = spawn(cmd, args, { env, shell: true });
+
+    if (!this.childProcess || !this.childProcess.stdout || !this.childProcess.stdin) {
+      throw new Error(`[MCPClient] Failed to spawn stdio child process for "${this.config.name}"`);
+    }
+
+    this.childProcess.stdout.on('data', (data: Buffer) => {
+      this.stdioBuffer += data.toString('utf-8');
+      this.processStdioLines();
+    });
+
+    this.childProcess.stderr?.on('data', (data: Buffer) => {
+      console.warn(`[MCPClient] [stderr] [${this.config.name}]:`, data.toString('utf-8'));
+    });
+
+    this.childProcess.on('close', (code) => {
+      console.log(`[MCPClient] Connection closed for "${this.config.name}" with exit code: ${code}`);
+      this.disconnect();
+    });
+
+    // Wait a brief tick to ensure stdout is set up
+    await new Promise((resolve) => setTimeout(resolve, 100));
+  }
+
+  private processStdioLines() {
+    let newlineIndex: number;
+    while ((newlineIndex = this.stdioBuffer.indexOf('\n')) !== -1) {
+      const line = this.stdioBuffer.substring(0, newlineIndex).trim();
+      this.stdioBuffer = this.stdioBuffer.substring(newlineIndex + 1);
+
+      if (!line) continue;
+
+      try {
+        const payload = JSON.parse(line);
+        if (payload.id !== undefined) {
+          const req = this.pendingRequests.get(Number(payload.id));
+          if (req) {
+            this.pendingRequests.delete(Number(payload.id));
+            if (payload.error) {
+              req.reject(new Error(payload.error.message || JSON.stringify(payload.error)));
+            } else {
+              req.resolve(payload.result);
+            }
+          }
+        }
+      } catch (e) {
+        console.warn(`[MCPClient] Error parsing stdio JSON-RPC line from "${this.config.name}":`, e, 'Line content:', line);
+      }
+    }
+  }
+
+  private async connectSSE(): Promise<void> {
+    if (!this.config.url) {
+      throw new Error(`[MCPClient] URL is required for sse transport on "${this.config.name}"`);
+    }
+
+    // Set up standard SSE handshake / SSE connection
+    // In serverless contexts, we pull initial SSE transport configuration,
+    // and route HTTP POSTs to the client URL.
+    this.sseUrl = this.config.url;
+  }
+
+  private async sendJsonRpc(method: string, params: any): Promise<any> {
+    const id = this.getNextMessageId();
+    const request = {
+      jsonrpc: '2.0',
+      id,
+      method,
+      params,
+    };
+
+    if (this.config.transport === 'stdio') {
+      if (!this.childProcess || !this.childProcess.stdin) {
+        throw new Error(`[MCPClient] Server "${this.config.name}" is not connected (stdio process missing)`);
+      }
+      return new Promise((resolve, reject) => {
+        this.pendingRequests.set(id, { resolve, reject });
+        this.childProcess!.stdin!.write(JSON.stringify(request) + '\n', 'utf-8');
+      });
+    } else {
+      // SSE / HTTP POST fallback for cloud/serverless setups
+      if (!this.sseUrl) {
+        throw new Error(`[MCPClient] Server "${this.config.name}" SSE URL is not initialized`);
+      }
+      const response = await fetch(this.sseUrl, {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify(request),
+      });
+      if (!response.ok) {
+        throw new Error(`SSE JSON-RPC POST failed with status ${response.status}`);
+      }
+      const payload = await response.json();
+      if (payload.error) {
+        throw new Error(payload.error.message || JSON.stringify(payload.error));
+      }
+      return payload.result;
+    }
+  }
+
+  private async sendNotification(method: string, params?: any): Promise<void> {
+    const notification = {
+      jsonrpc: '2.0',
+      method,
+      params,
+    };
+
+    if (this.config.transport === 'stdio') {
+      if (this.childProcess?.stdin) {
+        this.childProcess.stdin.write(JSON.stringify(notification) + '\n', 'utf-8');
+      }
+    } else if (this.sseUrl) {
+      await fetch(this.sseUrl, {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify(notification),
+      }).catch(err => console.warn('[MCPClient] Failed to send notification over SSE:', err));
+    }
+  }
+
+  /**
+   * List all tools offered by the MCP server.
+   */
+  async listTools(): Promise<MCPTool[]> {
+    await this.connect();
+    const result = await this.sendJsonRpc('tools/list', {});
+    return result?.tools || [];
+  }
+
+  /**
+   * Execute a tool on the MCP server.
+   */
+  async callTool(name: string, args: Record<string, any> = {}): Promise<MCPCallResult> {
+    await this.connect();
+    return await this.sendJsonRpc('tools/call', { name, arguments: args });
+  }
+
+  disconnect() {
+    this.initialized = false;
+    this.pendingRequests.forEach(req => req.reject(new Error('MCPClient disconnected')));
+    this.pendingRequests.clear();
+    
+    if (this.childProcess) {
+      try {
+        this.childProcess.kill();
+      } catch { /* silent */ }
+      this.childProcess = undefined;
+    }
+  }
+}
+
+export class MCPRegistry {
+  private clients: MCPClient[] = [];
+
+  constructor(servers: MCPServerConfig[] = []) {
+    this.clients = servers.map(cfg => new MCPClient(cfg));
+  }
+
+  async getAllTools(): Promise<Array<{ serverName: string; tool: MCPTool }>> {
+    const list: Array<{ serverName: string; tool: MCPTool }> = [];
+    await Promise.all(
+      this.clients.map(async (client) => {
+        try {
+          const tools = await client.listTools();
+          tools.forEach(t => list.push({ serverName: client['config'].name, tool: t }));
+        } catch (e) {
+          console.warn(`[MCPRegistry] Failed to fetch tools from server "${client['config'].name}":`, e);
+        }
+      })
+    );
+    return list;
+  }
+
+  async callTool(toolName: string, args: Record<string, any> = {}): Promise<MCPCallResult> {
+    // Find client hosting this tool. (In standard MCP, tool names are unique per server block)
+    // We will search across all connected clients.
+    for (const client of this.clients) {
+      try {
+        const tools = await client.listTools();
+        if (tools.some(t => t.name === toolName)) {
+          return await client.callTool(toolName, args);
+        }
+      } catch { /* silent fallback */ }
+    }
+    throw new Error(`[MCPRegistry] Tool "${toolName}" not found on any registered MCP server.`);
+  }
+
+  disconnectAll() {
+    this.clients.forEach(c => c.disconnect());
+  }
+}