From 64a8049438bf854f9ecdac86ebd900e3aad6eea3 Mon Sep 17 00:00:00 2001 From: Arun Venmany Date: Thu, 9 Apr 2026 19:00:02 +0530 Subject: [PATCH] bug fixes --- packages/toolpack-sdk/src/mcp/client.ts | 89 +++++++++--------- .../src/tools/mcp-tools/index.test.ts | 85 ++++++++++++++++- .../toolpack-sdk/src/tools/mcp-tools/index.ts | 92 ++++++++++++------- 3 files changed, 184 insertions(+), 82 deletions(-) diff --git a/packages/toolpack-sdk/src/mcp/client.ts b/packages/toolpack-sdk/src/mcp/client.ts index d499950..03aa9e3 100644 --- a/packages/toolpack-sdk/src/mcp/client.ts +++ b/packages/toolpack-sdk/src/mcp/client.ts @@ -79,6 +79,10 @@ export class McpClient extends EventEmitter { return this._connected && this.process !== null; } + private async initializeServer(): Promise { + await this.request('initialize', { client: 'toolpack-sdk' }); + } + // ====================================================================== // Connection // ====================================================================== @@ -88,62 +92,55 @@ export class McpClient extends EventEmitter { throw new McpConnectionError('Client is shutting down'); } - return new Promise((resolve, reject) => { - try { - this.buffer = ''; - this.process = spawn(this.config.command, this.config.args || [], { - env: { ...process.env, ...this.config.env }, - stdio: ['pipe', 'pipe', 'pipe'], // stdin, stdout, stderr - }); - - if (!this.process.stdout || !this.process.stdin) { - throw new McpConnectionError('Failed to spawn MCP server: stdout/stdin unavailable'); - } - - this.process.stdout.on('data', (data: Buffer) => { - this.handleData(data); - }); + this.buffer = ''; + this.process = spawn(this.config.command, this.config.args || [], { + env: { ...process.env, ...this.config.env }, + stdio: ['pipe', 'pipe', 'pipe'], // stdin, stdout, stderr + }); - // Route child stderr through logger instead of inheriting - // (inherited stderr corrupts Ink TUI rendering) - if (this.process.stderr) { - this.process.stderr.on('data', (data: Buffer) => { - logWarn(`[MCP server stderr] ${data.toString().trim()}`); - }); - } + if (!this.process.stdout || !this.process.stdin) { + throw new McpConnectionError('Failed to spawn MCP server: stdout/stdin unavailable'); + } - this.process.on('error', (err) => { - this._connected = false; - this.emit('error', err); - }); + this.process.stdout.on('data', (data: Buffer) => { + this.handleData(data); + }); - this.process.on('exit', (code) => { - const wasConnected = this._connected; - this._connected = false; - this.process = null; + // Route child stderr through logger instead of inheriting + // (inherited stderr corrupts Ink TUI rendering) + if (this.process.stderr) { + this.process.stderr.on('data', (data: Buffer) => { + logWarn(`[MCP server stderr] ${data.toString().trim()}`); + }); + } - // Reject all pending requests - this.rejectAllPending( - new McpConnectionError(`MCP server exited with code ${code}`, code) - ); + this.process.on('error', (err) => { + this._connected = false; + this.emit('error', err); + }); - this.emit('close', code); + this.process.on('exit', (code) => { + const wasConnected = this._connected; + this._connected = false; + this.process = null; - // Auto-reconnect on unexpected crash - if (wasConnected && !this._shuttingDown && this.autoReconnect) { - this.attemptReconnect(); - } - }); + // Reject all pending requests + this.rejectAllPending( + new McpConnectionError(`MCP server exited with code ${code}`, code) + ); - this._connected = true; - this._reconnectAttempts = 0; + this.emit('close', code); - // Give it a moment to start - setTimeout(resolve, 500); - } catch (error) { - reject(error); + // Auto-reconnect on unexpected crash + if (wasConnected && !this._shuttingDown && this.autoReconnect) { + this.attemptReconnect(); } }); + + this._reconnectAttempts = 0; + + await this.initializeServer(); + this._connected = true; } // ====================================================================== diff --git a/packages/toolpack-sdk/src/tools/mcp-tools/index.test.ts b/packages/toolpack-sdk/src/tools/mcp-tools/index.test.ts index f72d90a..26abe15 100644 --- a/packages/toolpack-sdk/src/tools/mcp-tools/index.test.ts +++ b/packages/toolpack-sdk/src/tools/mcp-tools/index.test.ts @@ -23,6 +23,10 @@ vi.mock('../../mcp/client.js', async () => { } async request(method: string): Promise { + if (method === 'initialize') { + return { initialized: true }; + } + if (method === 'tools/list') { return { tools: [ @@ -291,6 +295,35 @@ describe('McpToolManager', () => { const serverTools = toolsAfter.filter(t => t.name.includes('test-server')); expect(serverTools).toHaveLength(0); }); + + it('should only remove exact server tools when disconnecting', async () => { + const config: McpToolsConfig = { + servers: [], + }; + + manager = new McpToolManager(config); + + await manager.connectServer({ + name: 'test', + command: 'node', + args: ['test.js'], + }); + + await manager.connectServer({ + name: 'test-server', + command: 'node', + args: ['test.js'], + }); + + await manager.disconnectServer('test'); + + const toolsAfter = manager.getToolDefinitions(); + const exactServerTools = toolsAfter.filter(t => t.name.startsWith('mcp.test-server.')); + expect(exactServerTools.length).toBeGreaterThan(0); + + const removedTools = toolsAfter.filter(t => t.name.startsWith('mcp.test.')); + expect(removedTools.every(t => t.name.startsWith('mcp.test-server.'))).toBe(true); + }); it('should disconnect from all servers', async () => { const config: McpToolsConfig = { @@ -355,6 +388,48 @@ describe('McpToolManager', () => { expect(result).toBeDefined(); expect(typeof result).toBe('string'); }); + + it('should refresh tools after reconnecting', async () => { + const config: McpToolsConfig = { + servers: [ + { + name: 'test-server', + command: 'node', + args: ['test.js'], + }, + ], + }; + + manager = new McpToolManager(config); + await manager.connectServer(config.servers[0]); + + const client = (manager as any).clients.get('test-server'); + client.request = async (method: string): Promise => { + if (method === 'tools/list') { + return { + tools: [ + { + name: 'updated_tool', + description: 'Updated tool after reconnect', + inputSchema: { type: 'object', properties: {}, required: [] }, + }, + ], + }; + } + + if (method === 'initialize') { + return { initialized: true }; + } + throw new Error(`Unknown method: ${method}`); + }; + + client.emit('reconnected', { attempt: 1 }); + await new Promise(resolve => setTimeout(resolve, 0)); + + const tools = manager.getToolDefinitions(); + expect(tools.some(t => t.name.includes('updated_tool'))).toBe(true); + expect(tools.some(t => t.name.includes('test_tool'))).toBe(false); + }); }); }); @@ -395,10 +470,10 @@ describe('createMcpToolProject', () => { ], }; - const project = await createMcpToolProject(config); + const project = await createMcpToolProject(config) as any; - expect((project as any)._mcpManager).toBeDefined(); - expect((project as any)._mcpManager).toBeInstanceOf(McpToolManager); + expect(project.mcpManager).toBeDefined(); + expect(project.mcpManager).toBeInstanceOf(McpToolManager); // Cleanup await disconnectMcpToolProject(project); @@ -437,7 +512,7 @@ describe('createMcpToolProject', () => { expect(project.tools.length).toBeGreaterThan(0); - const manager = (project as any)._mcpManager; + const manager = (project as any).mcpManager; expect(manager.getConnectedServers()).toHaveLength(2); // Cleanup @@ -458,7 +533,7 @@ describe('disconnectMcpToolProject', () => { }; const project = await createMcpToolProject(config); - const manager = (project as any)._mcpManager; + const manager = (project as any).mcpManager; expect(manager.getConnectedServers()).toHaveLength(1); diff --git a/packages/toolpack-sdk/src/tools/mcp-tools/index.ts b/packages/toolpack-sdk/src/tools/mcp-tools/index.ts index 076383d..c6932b2 100644 --- a/packages/toolpack-sdk/src/tools/mcp-tools/index.ts +++ b/packages/toolpack-sdk/src/tools/mcp-tools/index.ts @@ -43,6 +43,7 @@ export class McpToolManager { private clients = new Map(); private serverConfigs = new Map(); private toolDefinitions = new Map(); + private toolOwners = new Map(); constructor(private config: McpToolsConfig) {} @@ -67,21 +68,9 @@ export class McpToolManager { // Connect to the server await client.connect(); - // Discover available tools - const toolsResponse = await client.request('tools/list'); - const mcpTools: McpTool[] = toolsResponse?.tools || []; - - logInfo(`[MCP] Discovered ${mcpTools.length} tools from ${name}`); - - // Convert MCP tools to Toolpack tools - const prefix = toolPrefix || `mcp.${name}.`; - for (const mcpTool of mcpTools) { - const toolDef = this.convertMcpTool(mcpTool, name, prefix, client); - this.toolDefinitions.set(toolDef.name, toolDef); - } - this.clients.set(name, client); this.serverConfigs.set(name, serverConfig); + await this.discoverServerTools(name, client, serverConfig); } catch (error) { logError(`[MCP] Failed to connect to ${name}: ${error}`); @@ -113,9 +102,10 @@ export class McpToolManager { logInfo(`[MCP] Disconnecting from ${name}`); // Remove tools from this server - for (const [toolName, toolDef] of this.toolDefinitions) { - if (toolDef.name.includes(name)) { + for (const [toolName, owner] of this.toolOwners) { + if (owner === name) { this.toolDefinitions.delete(toolName); + this.toolOwners.delete(toolName); } } @@ -195,31 +185,74 @@ export class McpToolManager { * Set up event handlers for an MCP client */ private setupClientEvents(client: McpClient, serverName: string): void { - // McpClient extends EventEmitter, so we can use .on() - (client as any).on('error', (error: any) => { + client.on('error', (error: any) => { logError(`[MCP] ${serverName} error: ${error}`); }); - (client as any).on('close', (code: any) => { + client.on('close', (code: any) => { logWarn(`[MCP] ${serverName} closed with code ${code}`); }); - (client as any).on('reconnecting', ({ attempt, max }: any) => { + client.on('reconnecting', ({ attempt, max }: any) => { logInfo(`[MCP] ${serverName} reconnecting (${attempt}/${max})`); }); - (client as any).on('reconnected', ({ attempt }: any) => { + client.on('reconnected', ({ attempt }: any) => { logInfo(`[MCP] ${serverName} reconnected after ${attempt} attempts`); + this.refreshServerTools(serverName).catch(err => { + logError(`[MCP] ${serverName} tool refresh failed after reconnect: ${err}`); + }); }); - (client as any).on('reconnect_failed', (attempts: any) => { + client.on('reconnect_failed', (attempts: any) => { logError(`[MCP] ${serverName} failed to reconnect after ${attempts} attempts`); }); - (client as any).on('notification', (message: any) => { + client.on('notification', (message: any) => { logInfo(`[MCP] ${serverName} notification: ${JSON.stringify(message)}`); }); } + + private removeServerToolDefinitions(serverName: string): void { + for (const [toolName, owner] of this.toolOwners) { + if (owner === serverName) { + this.toolDefinitions.delete(toolName); + this.toolOwners.delete(toolName); + } + } + } + + private async discoverServerTools( + serverName: string, + client: McpClient, + serverConfig: McpServerConfig + ): Promise { + const toolsResponse = await client.request('tools/list'); + const mcpTools: McpTool[] = toolsResponse?.tools || []; + + logInfo(`[MCP] Discovered ${mcpTools.length} tools from ${serverName}`); + + this.removeServerToolDefinitions(serverName); + + const prefix = serverConfig.toolPrefix || `mcp.${serverName}.`; + for (const mcpTool of mcpTools) { + const toolDef = this.convertMcpTool(mcpTool, serverName, prefix, client); + this.toolDefinitions.set(toolDef.name, toolDef); + this.toolOwners.set(toolDef.name, serverName); + } + } + + private async refreshServerTools(serverName: string): Promise { + const client = this.clients.get(serverName); + const serverConfig = this.serverConfigs.get(serverName); + if (!client || !serverConfig) return; + + await this.discoverServerTools(serverName, client, serverConfig); + } +} + +export interface McpToolProject extends ToolProject { + mcpManager: McpToolManager; } /** @@ -253,7 +286,7 @@ export class McpToolManager { */ export async function createMcpToolProject( config: McpToolsConfig -): Promise { +): Promise { const manager = new McpToolManager(config); // Connect to all servers @@ -262,7 +295,7 @@ export async function createMcpToolProject( // Get all tool definitions const tools = manager.getToolDefinitions(); - const project: ToolProject = { + const project: McpToolProject = { manifest: { key: 'mcp-tools', name: 'mcp-tools', @@ -274,20 +307,17 @@ export async function createMcpToolProject( tools: tools.map(t => t.name), }, tools, + mcpManager: manager, }; - // Store manager for lifecycle management (non-standard property) - (project as any)._mcpManager = manager; - return project; } /** * Disconnect all MCP servers in a tool project */ -export async function disconnectMcpToolProject(project: ToolProject): Promise { - const manager = (project as any)._mcpManager as McpToolManager | undefined; - if (manager) { - await manager.disconnectAll(); +export async function disconnectMcpToolProject(project: ToolProject | McpToolProject): Promise { + if ('mcpManager' in project) { + await project.mcpManager.disconnectAll(); } } \ No newline at end of file