diff --git a/AGENTS.md b/AGENTS.md index da938fe..a513eee 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -49,7 +49,7 @@ These are cross-cutting; see leaf Intent Nodes for module-specific invariants. 2. **`mcp-client-id` required for /mcp** — POST, GET, DELETE all reject without header (400) 3. **Fail-secure** — Invalid inputs return empty objects/arrays, not errors 4. **Silent module failures** — Toolsets activate with partial tools if loaders fail -5. **Disable is state-only** — MCP SDK has no tool unregister; see `src/core/AGENTS.md` +5. **Disable unregisters tools** — `disableToolset` removes the toolset's tools from the MCP server via the SDK's `RegisteredTool.remove()`; see `src/core/AGENTS.md` ## Module Index diff --git a/README.md b/README.md index 22868f1..19adb4a 100644 --- a/README.md +++ b/README.md @@ -1081,7 +1081,7 @@ Meta-tools are registered based on mode: **DYNAMIC mode** (registered by default, or when `registerMetaTools` is true): - `enable_toolset` - Enable a toolset by name -- `disable_toolset` - Disable a toolset by name (state only; tools remain registered) +- `disable_toolset` - Disable a toolset by name (unregisters the toolset's tools, so clients can no longer list or call them) - `list_toolsets` - List available toolsets with active status - `describe_toolset` - Describe a specific toolset with definition and tools - `list_tools` - List currently registered tool names diff --git a/src/core/AGENTS.md b/src/core/AGENTS.md index f2d7d95..d8979e4 100644 --- a/src/core/AGENTS.md +++ b/src/core/AGENTS.md @@ -13,11 +13,11 @@ Central orchestration layer. Manages toolset lifecycle, tool registration, and s ## Invariants 1. **All tools go through ToolRegistry** — Collision detection happens here only. Bypassing it causes silent name conflicts. -2. **Disable is state-only** — MCP SDK has no `server.removeTool()`; disabled toolsets' tools remain callable by clients. `activeToolsets` tracks logical state only. +2. **Disable unregisters the tools** — `DynamicToolManager` keeps the `RegisteredTool` handle `registerTool()` returns for each tool and calls `remove()` on disable, then frees the names in ToolRegistry (`removeForToolset`) so the toolset can be enabled again. Disabled tools are neither listed nor callable. 3. **Notifications may fail silently** — "Not connected" errors are expected during SSE disconnect and are swallowed. 4. **Namespacing applied BEFORE collision check** — `mapAndValidate()` generates safe names first, then checks for conflicts. -5. **Toolset added to activeToolsets AFTER registration** — All tools must register successfully before toolset is marked active. Partial failures leave the toolset inactive but orphaned tools remain registered (MCP limitation). -6. **`enableToolsets()` batches notifications** — Calls `enableToolset(name, skipNotification=true)` per toolset, sends one `tools/list_changed` notification at the end. +5. **Toolset added to activeToolsets AFTER registration** — All tools must register successfully before toolset is marked active. A partial failure removes the tools already registered, so nothing is left callable and the toolset can be enabled again. +6. **`enableToolsets()` batches notifications** — Calls `enableToolset(name, skipNotification=true)` per toolset, sends one `tools/list_changed` notification at the end. The SDK also sends its own `tools/list_changed` for every `registerTool()` and `remove()` call on a connected server, so a disable of N tools sends N+1. 7. **JSON Schema `inputSchema` uses `registerTool` with a passthrough object** — The SDK rejects plain JSON Schema as a tool schema. Zod raw shapes and schemas are passed to `registerTool` as is. Clients see an open object schema for JSON Schema tools, not the declared one. ## Enable Toolset Flow @@ -38,7 +38,7 @@ Tools are registered with `McpServer.registerTool` (the 5-arg `server.tool()` ov ## Anti-patterns - Bypassing ToolRegistry for tool registration (causes collision issues) -- Expecting disable to unregister tools from MCP (it can't) +- Registering tools without keeping the `registerTool()` handle (disable and failed-enable rollback can't remove them) - Throwing on notification failures (they're expected in SSE disconnect) - Using `_meta` as a toolset key in the catalog (reserved for meta-tools — see `src/meta/AGENTS.md`) diff --git a/src/core/DynamicToolManager.ts b/src/core/DynamicToolManager.ts index c1ed0d5..9518138 100644 --- a/src/core/DynamicToolManager.ts +++ b/src/core/DynamicToolManager.ts @@ -1,4 +1,7 @@ -import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import type { + McpServer, + RegisteredTool, +} from "@modelcontextprotocol/sdk/server/mcp.js"; import { z } from "zod"; import type { ZodRawShapeCompat } from "@modelcontextprotocol/sdk/server/zod-compat.js"; import type { @@ -20,6 +23,8 @@ export class DynamicToolManager { private readonly toolRegistry: ToolRegistry; private readonly activeToolsets = new Set(); + /** SDK handles of each toolset's registered tools, used to unregister them. */ + private readonly toolHandles = new Map(); constructor(options: DynamicToolManagerOptions) { this.server = options.server; @@ -94,11 +99,8 @@ export class DynamicToolManager { return { success: false, message: policyCheck.message }; } - // Track tools registered for this enable operation to allow rollback - const registeredTools: string[] = []; - try { - const toolCount = await this.resolveAndRegisterTools(sanitized, registeredTools); + const toolCount = await this.resolveAndRegisterTools(sanitized); // Track state only after successful registration this.activeToolsets.add(sanitized); @@ -110,7 +112,7 @@ export class DynamicToolManager { return this.buildEnableResult(sanitized, toolCount); } catch (error) { - this.handlePartialFailure(sanitized, registeredTools); + this.removeToolsetTools(sanitized); return { success: false, message: `Failed to enable toolset '${sanitized}': ${ @@ -145,13 +147,9 @@ export class DynamicToolManager { /** * @param sanitized - The validated toolset name - * @param registeredTools - Mutable array tracking registered tool names for rollback * @returns The number of tools resolved */ - private async resolveAndRegisterTools( - sanitized: string, - registeredTools: string[] - ): Promise { + private async resolveAndRegisterTools(sanitized: string): Promise { const resolvedTools = await this.resolver.resolveToolsForToolsets( [sanitized], this.context @@ -161,7 +159,6 @@ export class DynamicToolManager { const mapped = this.toolRegistry.mapAndValidate(sanitized, resolvedTools); for (const tool of mapped) { this.registerSingleTool(tool, sanitized); - registeredTools.push(tool.name); } } @@ -184,20 +181,20 @@ export class DynamicToolManager { } /** - * @param sanitized - The toolset name that partially failed - * @param registeredTools - Tools that were registered before the failure + * Unregisters every tool of a toolset from the MCP server and the registry. + * Also rolls back a partially failed enable, so no orphaned tools stay callable. + * + * @param toolsetKey - The toolset whose tools to remove + * @returns The number of tools removed */ - private handlePartialFailure( - sanitized: string, - registeredTools: string[] - ): void { - if (registeredTools.length > 0) { - console.warn( - `Partial failure enabling toolset '${sanitized}'. ` + - `${registeredTools.length} tools were registered but toolset activation failed. ` + - `Tools remain registered due to MCP limitations: ${registeredTools.join(", ")}` - ); + private removeToolsetTools(toolsetKey: string): number { + const handles = this.toolHandles.get(toolsetKey) ?? []; + for (const handle of handles) { + handle.remove(); } + this.toolHandles.delete(toolsetKey); + this.toolRegistry.removeForToolset(toolsetKey); + return handles.length; } /** @@ -258,7 +255,7 @@ export class DynamicToolManager { ? (tool.inputSchema as ZodRawShapeCompat) : z.object({}).passthrough(); - this.server.registerTool( + const handle = this.server.registerTool( tool.name, { title: tool.annotations?.title, @@ -270,6 +267,10 @@ export class DynamicToolManager { return await tool.handler(args); } ); + this.toolHandles.set(toolsetKey, [ + ...(this.toolHandles.get(toolsetKey) ?? []), + handle, + ]); this.toolRegistry.addForToolset(toolsetKey, tool.name); } @@ -300,14 +301,14 @@ export class DynamicToolManager { }; } - // State-only disable; no unregistration support in MCP + const removedCount = this.removeToolsetTools(sanitized); this.activeToolsets.delete(sanitized); await this.notifyToolsChanged(); return { success: true, - message: `Toolset '${sanitized}' disabled successfully. Individual tools remain registered due to MCP limitations.`, + message: `Toolset '${sanitized}' disabled successfully. Removed ${removedCount} tools.`, }; } diff --git a/src/core/ToolRegistry.ts b/src/core/ToolRegistry.ts index f9f9936..cd0bf25 100644 --- a/src/core/ToolRegistry.ts +++ b/src/core/ToolRegistry.ts @@ -49,6 +49,18 @@ export class ToolRegistry { this.toolsetToNames.set(toolsetKey, set); } + /** + * Forgets every tool name registered for a toolset, so the names can be registered again. + * + * @param toolsetKey - The toolset whose tools to forget + */ + public removeForToolset(toolsetKey: string): void { + for (const name of this.toolsetToNames.get(toolsetKey) ?? []) { + this.names.delete(name); + } + this.toolsetToNames.delete(toolsetKey); + } + public mapAndValidate( toolsetKey: string, tools: McpToolDefinition[] diff --git a/src/meta/registerMetaTools.ts b/src/meta/registerMetaTools.ts index 60db1c1..812966d 100644 --- a/src/meta/registerMetaTools.ts +++ b/src/meta/registerMetaTools.ts @@ -60,7 +60,7 @@ export function registerMetaTools( server.registerTool( "disable_toolset", { - description: "Disable a toolset by name (state only)", + description: "Disable a toolset by name and remove its tools", inputSchema: { name: z.string().describe("Toolset name") }, annotations: { destructiveHint: true, idempotentHint: true }, }, diff --git a/tests/dynamicToolManager.test.ts b/tests/dynamicToolManager.test.ts index 3c126f6..7212d27 100644 --- a/tests/dynamicToolManager.test.ts +++ b/tests/dynamicToolManager.test.ts @@ -340,3 +340,100 @@ describe("DynamicToolManager with a real McpServer", () => { expect((res.content as any)[0].text).toBe("MSFT"); }); }); + +describe("DynamicToolManager disable removes tools", () => { + const ping = { + name: "ping", + description: "Ping", + inputSchema: {}, + handler: async () => ({ content: [{ type: "text", text: "pong" }] }), + }; + + // The SDK can't add the tools capability after connect, so, like the meta-tools + // do in a real server, register one tool before connecting. + function createServer() { + const server = new McpServer( + { name: "t", version: "1.0.0" }, + { capabilities: { tools: { listChanged: true } } } + ); + server.registerTool("meta", { inputSchema: {} }, async () => ({ content: [] })); + return server; + } + + async function connectManager(catalog: Record) { + const server = createServer(); + const resolver = new ModuleResolver({ catalog } as any); + const manager = new DynamicToolManager({ server, resolver }); + const [serverSide, clientSide] = InMemoryTransport.createLinkedPair(); + await server.connect(serverSide); + const client = new Client({ name: "c", version: "1.0.0" }); + await client.connect(clientSide); + return { manager, client }; + } + + async function listedNames(client: Client) { + return (await client.listTools()).tools.map((t) => t.name); + } + + it("makes a disabled toolset's tools unlisted and uncallable", async () => { + const { manager, client } = await connectManager({ + core: { name: "Core", description: "", tools: [ping] }, + }); + await manager.enableToolset("core"); + expect(await listedNames(client)).toEqual(["meta", "core.ping"]); + + const res = await manager.disableToolset("core"); + expect(res.success).toBe(true); + expect(res.message).toContain("Removed 1 tools"); + expect(await listedNames(client)).toEqual(["meta"]); + + const call = await client.callTool({ name: "core.ping", arguments: {} }); + expect(call.isError).toBe(true); + expect((call.content as any)[0].text).toMatch(/not found/i); + }); + + it("frees the tool names so the toolset can be enabled again", async () => { + const { manager, client } = await connectManager({ + core: { name: "Core", description: "", tools: [ping] }, + }); + await manager.enableToolset("core"); + await manager.disableToolset("core"); + expect(manager.getStatus().tools).toEqual([]); + expect(manager.getStatus().toolsetToTools).toEqual({}); + + const again = await manager.enableToolset("core"); + expect(again.success).toBe(true); + const res = await client.callTool({ name: "core.ping", arguments: {} }); + expect((res.content as any)[0].text).toBe("pong"); + }); + + it("removes only the disabled toolset's tools", async () => { + const { manager, client } = await connectManager({ + a: { name: "A", description: "", tools: [ping] }, + b: { name: "B", description: "", tools: [ping] }, + }); + + await manager.enableToolsets(["a", "b"]); + await manager.disableToolset("a"); + expect(await listedNames(client)).toEqual(["meta", "b.ping"]); + expect(manager.getActiveToolsets()).toEqual(["b"]); + }); + + it("rolls back already registered tools when enabling fails part-way", async () => { + // The second tool has the same name, so the SDK rejects its registration. + const tools = [ping, ping]; + const { manager, client } = await connectManager({ + core: { name: "Core", description: "", tools }, + }); + const res = await manager.enableToolset("core"); + expect(res.success).toBe(false); + expect(manager.isActive("core")).toBe(false); + expect(await listedNames(client)).toEqual(["meta"]); + expect(manager.getStatus().tools).toEqual([]); + + // The rolled-back names are free again once the toolset is fixed. + tools.pop(); + expect((await manager.enableToolset("core")).success).toBe(true); + expect(await listedNames(client)).toEqual(["meta", "core.ping"]); + }); +}); diff --git a/tests/e2e/dynamicMode.e2e.test.ts b/tests/e2e/dynamicMode.e2e.test.ts index b0abd62..ae39992 100644 --- a/tests/e2e/dynamicMode.e2e.test.ts +++ b/tests/e2e/dynamicMode.e2e.test.ts @@ -195,6 +195,15 @@ describe("DYNAMIC mode E2E", () => { const core = list.toolsets.find((ts) => ts.key === "core"); expect(core.active).toBe(false); + // Disabled tools are no longer listed and can no longer be called + expect(extractToolNames(await client.listTools())).not.toContain("core.ping"); + const staleCall = await client.callTool({ + name: "core.ping", + arguments: {}, + } as any); + expect(staleCall.isError).toBe(true); + expect(extractTextContent(staleCall)).toMatch(/not found/i); + await client.close(); }); diff --git a/tests/helpers/fakes.ts b/tests/helpers/fakes.ts index 2281ff7..fd3263e 100644 --- a/tests/helpers/fakes.ts +++ b/tests/helpers/fakes.ts @@ -20,13 +20,20 @@ export function createFakeMcpServer(options: { withNotifier?: boolean } = {}) { config: { description: string; inputSchema?: any; annotations?: RegisteredTool["annotations"] }, handler: (args: any) => Promise | any ) { - tools.push({ + const entry: RegisteredTool = { name, description: config.description, schema: config.inputSchema, annotations: config.annotations, handler, - }); + }; + tools.push(entry); + return { + remove() { + const index = tools.indexOf(entry); + if (index !== -1) tools.splice(index, 1); + }, + }; }, async connect(_transport: unknown): Promise {}, }; diff --git a/tests/toolRegistry.test.ts b/tests/toolRegistry.test.ts index e144037..729d4b9 100644 --- a/tests/toolRegistry.test.ts +++ b/tests/toolRegistry.test.ts @@ -218,4 +218,26 @@ describe("ToolRegistry", () => { expect(reg.has("nonexistent")).toBe(false); }); }); + + describe("removeForToolset() method", () => { + it("forgets the toolset's names so they can be added again", () => { + const reg = new ToolRegistry({ namespaceWithToolset: false }); + reg.addForToolset("a", "ping"); + reg.addForToolset("a", "pong"); + reg.addForToolset("b", "other"); + + reg.removeForToolset("a"); + expect(reg.has("ping")).toBe(false); + expect(reg.list()).toEqual(["other"]); + expect(reg.listByToolset()).toEqual({ b: ["other"] }); + expect(() => reg.addForToolset("a", "ping")).not.toThrow(); + }); + + it("ignores an unknown toolset", () => { + const reg = new ToolRegistry({ namespaceWithToolset: false }); + reg.add("ping"); + reg.removeForToolset("missing"); + expect(reg.list()).toEqual(["ping"]); + }); + }); });