Spaces:
Running
Running
File size: 2,240 Bytes
67c6e06 |
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 |
import { experimental_createMCPClient as createMCPClient } from 'ai';
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 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 {
// All servers are handled as SSE
const transport = {
type: 'sse' as const,
url: mcpServer.url,
headers: mcpServer.headers?.reduce((acc, header) => {
if (header.key) acc[header.key] = header.value || '';
return acc;
}, {} as Record<string, string>)
};
const mcpClient = await createMCPClient({ transport });
mcpClients.push(mcpClient);
const mcptools = await mcpClient.tools();
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);
}
}
} |