fix(mcp): clear prompt registry on refresh to prevent duplicates (#5385)

Co-authored-by: Jacob Richman <[email protected]>
Co-authored-by: Sandy Tao <[email protected]>

3 files changed, 40 insertions, 12 deletions

diff --git a/packages/core/src/prompts/prompt-registry.ts b/packages/core/src/prompts/prompt-registry.ts
index 56699130..a94183ac 100644
--- a/packages/core/src/prompts/prompt-registry.ts
+++ b/packages/core/src/prompts/prompt-registry.ts
@@ -53,4 +53,22 @@ export class PromptRegistry {
     }
     return serverPrompts.sort((a, b) => a.name.localeCompare(b.name));
   }
+
+  /**
+   * Clears all the prompts from the registry.
+   */
+  clear(): void {
+    this.prompts.clear();
+  }
+
+  /**
+   * Removes all prompts from a specific server.
+   */
+  removePromptsByServer(serverName: string): void {
+    for (const [name, prompt] of this.prompts.entries()) {
+      if (prompt.serverName === serverName) {
+        this.prompts.delete(name);
+      }
+    }
+  }
 }
diff --git a/packages/core/src/tools/tool-registry.test.ts b/packages/core/src/tools/tool-registry.test.ts
index de7c6309..88b23d84 100644
--- a/packages/core/src/tools/tool-registry.test.ts
+++ b/packages/core/src/tools/tool-registry.test.ts
@@ -172,6 +172,10 @@ describe('ToolRegistry', () => {
     );
     vi.spyOn(config, 'getMcpServers');
     vi.spyOn(config, 'getMcpServerCommand');
+    vi.spyOn(config, 'getPromptRegistry').mockReturnValue({
+      clear: vi.fn(),
+      removePromptsByServer: vi.fn(),
+    } as any);
     mockDiscoverMcpTools.mockReset().mockResolvedValue(undefined);
   });
 
@@ -353,7 +357,7 @@ describe('ToolRegistry', () => {
         mcpServerConfigVal,
         undefined,
         toolRegistry,
-        undefined,
+        config.getPromptRegistry(),
         false,
       );
     });
@@ -376,7 +380,7 @@ describe('ToolRegistry', () => {
         mcpServerConfigVal,
         undefined,
         toolRegistry,
-        undefined,
+        config.getPromptRegistry(),
         false,
       );
     });
diff --git a/packages/core/src/tools/tool-registry.ts b/packages/core/src/tools/tool-registry.ts
index 57627ee0..e60b8f74 100644
--- a/packages/core/src/tools/tool-registry.ts
+++ b/packages/core/src/tools/tool-registry.ts
@@ -150,6 +150,14 @@ export class ToolRegistry {
     this.tools.set(tool.name, tool);
   }
 
+  private removeDiscoveredTools(): void {
+    for (const tool of this.tools.values()) {
+      if (tool instanceof DiscoveredTool || tool instanceof DiscoveredMCPTool) {
+        this.tools.delete(tool.name);
+      }
+    }
+  }
+
   /**
    * Discovers tools from project (if available and configured).
    * Can be called multiple times to update discovered tools.
@@ -157,11 +165,9 @@ export class ToolRegistry {
    */
   async discoverAllTools(): Promise<void> {
     // remove any previously discovered tools
-    for (const tool of this.tools.values()) {
-      if (tool instanceof DiscoveredTool || tool instanceof DiscoveredMCPTool) {
-        this.tools.delete(tool.name);
-      }
-    }
+    this.removeDiscoveredTools();
+
+    this.config.getPromptRegistry().clear();
 
     await this.discoverAndRegisterToolsFromCommand();
 
@@ -182,11 +188,9 @@ export class ToolRegistry {
    */
   async discoverMcpTools(): Promise<void> {
     // remove any previously discovered tools
-    for (const tool of this.tools.values()) {
-      if (tool instanceof DiscoveredMCPTool) {
-        this.tools.delete(tool.name);
-      }
-    }
+    this.removeDiscoveredTools();
+
+    this.config.getPromptRegistry().clear();
 
     // discover tools using MCP servers, if configured
     await discoverMcpTools(
@@ -210,6 +214,8 @@ export class ToolRegistry {
       }
     }
 
+    this.config.getPromptRegistry().removePromptsByServer(serverName);
+
     const mcpServers = this.config.getMcpServers() ?? {};
     const serverConfig = mcpServers[serverName];
     if (serverConfig) {