File size: 2,887 Bytes
dff2be9
 
 
67c6e06
 
 
 
 
 
 
 
dff2be9
67c6e06
 
 
 
 
 
 
 
 
 
 
 
 
 
9cd06b7
67c6e06
 
 
 
 
 
 
 
 
 
 
 
9cd06b7
dff2be9
9cd06b7
 
 
 
 
 
e10614a
dff2be9
 
 
 
 
 
9cd06b7
 
 
 
 
67c6e06
dff2be9
 
 
 
 
67c6e06
 
dff2be9
67c6e06
 
 
 
 
 
 
 
 
 
 
 
 
dff2be9
67c6e06
 
 
 
 
 
 
dff2be9
67c6e06
 
 
 
 
 
 
 
 
 
 
 
dff2be9
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
import { Client as MCPClient } from "@modelcontextprotocol/sdk/client/index.js";
import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";

export interface KeyValuePair {
  key: string;
  value: string;
}

export interface MCPServerConfig {
  url: string;
  type: "sse" | "stdio";
  command?: string;
  args?: string[];
  env?: KeyValuePair[];
  headers?: KeyValuePair[];
}

export interface MCPClientManager {
  tools: Record<string, any>;
  clients: any[];
  cleanup: () => Promise<void>;
}

/**
 * Initialize MCP clients for API calls
 * This uses the already running persistent HTTP or SSE servers
 */
export async function initializeMCPClients(
  mcpServers: MCPServerConfig[] = [],
  abortSignal?: AbortSignal
): Promise<MCPClientManager> {
  // Initialize tools
  let tools = {};
  const mcpClients: any[] = [];

  // Process each MCP server configuration
  for (const mcpServer of mcpServers) {
    try {
      const headers = mcpServer.headers?.reduce((acc, header) => {
        if (header.key) acc[header.key] = header.value || "";
        return acc;
      }, {} as Record<string, string>);

      // All servers are handled as HTTP or SSE
      // SSE is only when URL ends with /sse
      // which is the heuristic used by other clients

      const transport = mcpServer.url.endsWith("/sse")
        ? new SSEClientTransport(new URL(mcpServer.url), {
            requestInit: {
              headers,
            },
          })
        : new StreamableHTTPClientTransport(new URL(mcpServer.url), {
            requestInit: {
              headers,
            },
          });

      const mcpClient = new MCPClient({
        name: "mcp-chat-client",
        version: "0.1.0",
      });
      await mcpClient.connect(transport);
      mcpClients.push(mcpClient);

      const mcptools = await mcpClient.listTools();

      console.log(`MCP tools from ${mcpServer.url}:`, Object.keys(mcptools));

      // Add MCP tools to tools object
      tools = { ...tools, ...mcptools };
    } catch (error) {
      console.error("Failed to initialize MCP client:", error);
      // Continue with other servers instead of failing the entire request
    }
  }

  // Register cleanup for all clients if an abort signal is provided
  if (abortSignal && mcpClients.length > 0) {
    abortSignal.addEventListener("abort", async () => {
      await cleanupMCPClients(mcpClients);
    });
  }

  return {
    tools,
    clients: mcpClients,
    cleanup: async () => await cleanupMCPClients(mcpClients),
  };
}

async function cleanupMCPClients(clients: any[]): Promise<void> {
  // Clean up the MCP clients
  for (const client of clients) {
    try {
      await client.close();
    } catch (error) {
      console.error("Error closing MCP client:", error);
    }
  }
}