From 8e99c26dd8d2bea622909d4ec0c017828397ebb8 Mon Sep 17 00:00:00 2001 From: luisfelipe-alt Date: Tue, 9 Jun 2026 18:51:02 +0000 Subject: [PATCH] fix(core): implement atomic update in MCP tool discovery (#27619) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Gal Zahavi <38544478+galz10@users.noreply.github.com> --- packages/core/src/tools/mcp-client.test.ts | 16 +++++++++++----- packages/core/src/tools/mcp-client.ts | 1 + 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/packages/core/src/tools/mcp-client.test.ts b/packages/core/src/tools/mcp-client.test.ts index 4aa9bef289..4d682c52c2 100644 --- a/packages/core/src/tools/mcp-client.test.ts +++ b/packages/core/src/tools/mcp-client.test.ts @@ -1244,9 +1244,15 @@ describe('mcp-client', () => { await client.disconnect(); expect(mockedClient.close).toHaveBeenCalledOnce(); - expect(mockedToolRegistry.removeMcpToolsByServer).toHaveBeenCalledOnce(); - expect(mockedPromptRegistry.removePromptsByServer).toHaveBeenCalledOnce(); - expect(resourceRegistry.removeResourcesByServer).toHaveBeenCalledOnce(); + expect(mockedToolRegistry.removeMcpToolsByServer).toHaveBeenCalledWith( + 'test-server', + ); + expect(mockedPromptRegistry.removePromptsByServer).toHaveBeenCalledWith( + 'test-server', + ); + expect(resourceRegistry.removeResourcesByServer).toHaveBeenCalledWith( + 'test-server', + ); }); }); @@ -1570,8 +1576,8 @@ describe('mcp-client', () => { // Trigger notification - should fail internally but catch the error await notificationCallback(); - // Should try to remove tools - expect(mockedToolRegistry.removeMcpToolsByServer).toHaveBeenCalled(); + // Should NOT try to remove tools because discovery failed (atomic refresh) + expect(mockedToolRegistry.removeMcpToolsByServer).not.toHaveBeenCalled(); // Should NOT emit success feedback expect(coreEvents.emitFeedback).not.toHaveBeenCalledWith( diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts index a50549e435..47048526d2 100644 --- a/packages/core/src/tools/mcp-client.ts +++ b/packages/core/src/tools/mcp-client.ts @@ -1404,6 +1404,7 @@ export async function discoverTools( error, mcpServerName, ); + throw error; } return []; }