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 { // 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 { // 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) {