@@ -4,7 +4,11 @@ import { WORKFLOW_LOADERS, WorkflowName, WORKFLOW_METADATA } from './generated-p
44import { ToolResponse } from '../types/common.js' ;
55import { PluginMeta } from './plugin-types.js' ;
66import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' ;
7- import { registerAndTrackTools } from '../utils/tool-registry.js' ;
7+ import {
8+ registerAndTrackTools ,
9+ removeTrackedTools ,
10+ isToolRegistered ,
11+ } from '../utils/tool-registry.js' ;
812import { ZodRawShape } from 'zod' ;
913
1014// Track enabled workflows and their tools for replacement functionality
@@ -30,29 +34,28 @@ function wrapHandlerWithExecutor(handler: ToolHandler) {
3034}
3135
3236/**
33- * Clear tracking of currently enabled workflows
34- * Note: MCP doesn't support removing tools, so this only clears our internal tracking.
35- * Tool replacement happens by overriding existing tool registrations.
37+ * Clear currently enabled workflows by actually removing registered tools
3638 */
3739export function clearEnabledWorkflows ( ) : void {
3840 if ( enabledTools . size === 0 ) {
39- log ( 'debug' , 'No tools to clear from tracking ' ) ;
41+ log ( 'debug' , 'No tools to clear' ) ;
4042 return ;
4143 }
4244
4345 const clearedWorkflows = Array . from ( enabledWorkflows ) ;
44- const clearedToolCount = enabledTools . size ;
46+ const toolNamesToRemove = Array . from ( enabledTools . keys ( ) ) ;
47+ const clearedToolCount = toolNamesToRemove . length ;
4548
46- log (
47- 'info' ,
48- `Clearing tracking for ${ clearedToolCount } tools from workflows: ${ clearedWorkflows . join ( ', ' ) } ` ,
49- ) ;
49+ log ( 'info' , `Removing ${ clearedToolCount } tools from workflows: ${ clearedWorkflows . join ( ', ' ) } ` ) ;
50+
51+ // Actually remove the registered tools using the tool registry
52+ const removedTools = removeTrackedTools ( toolNamesToRemove ) ;
5053
51- // Clear our tracking - tools will be overridden by new registrations
54+ // Clear our tracking
5255 enabledWorkflows . clear ( ) ;
5356 enabledTools . clear ( ) ;
5457
55- log ( 'debug ' , `✅ Cleared tracking for ${ clearedToolCount } tools` ) ;
58+ log ( 'info ' , `✅ Removed ${ removedTools . length } tools successfully ` ) ;
5659}
5760
5861/**
@@ -116,16 +119,22 @@ export async function enableWorkflows(
116119 callback : ( args : Record < string , unknown > ) => Promise < ToolResponse > ;
117120 } > = [ ] ;
118121
119- // Collect all tools from this workflow
122+ // Collect all tools from this workflow, filtering out already-registered tools
120123 for ( const toolKey of toolKeys ) {
121124 const tool = workflowModule [ toolKey ] as PluginMeta | undefined ;
122125
123126 if ( tool ?. name && typeof tool . handler === 'function' ) {
127+ // Always skip tools that are already registered (in all modes)
128+ if ( isToolRegistered ( tool . name ) ) {
129+ log ( 'debug' , `Skipping already registered tool: ${ tool . name } ` ) ;
130+ continue ;
131+ }
132+
124133 toolsToRegister . push ( {
125134 name : tool . name ,
126135 config : {
127136 description : tool . description ?? '' ,
128- inputSchema : tool . schema , // MCP SDK now handles complex types properly
137+ inputSchema : tool . schema ,
129138 } ,
130139 callback : wrapHandlerWithExecutor ( tool . handler as ToolHandler ) ,
131140 } ) ;
@@ -138,69 +147,29 @@ export async function enableWorkflows(
138147 }
139148 }
140149
141- // Use bulk registration with proper types - no runtime checking needed
142- try {
143- const availableTools = toolsToRegister . filter ( ( tool ) => {
144- // In testing/development, check for duplicate registrations
145- // The MCP SDK handles this internally, so this is just for logging
146- log ( 'debug' , `Preparing to register tool: ${ tool . name } ` ) ;
147- return true ;
148- } ) ;
149-
150- if ( availableTools . length > 0 ) {
151- log ( 'info' , `🚀 Enabling ${ availableTools . length } tools from '${ workflowName } ' workflow` ) ;
152-
153- // Convert to proper tool registration format, adapting callback signature
154- const toolRegistrations = availableTools . map ( ( tool ) => ( {
155- name : tool . name ,
156- config : {
157- description : tool . config . description ,
158- inputSchema : tool . config . inputSchema as unknown , // Cast to unknown for SDK interface
159- } ,
160- // Adapt callback to match SDK's expected signature (args, extra) => result
161- callback : ( args : unknown ) : Promise < ToolResponse > =>
162- tool . callback ( args as Record < string , unknown > ) ,
163- } ) ) ;
164-
165- // Use registerTools with proper types and tracking
166- const registeredTools = registerAndTrackTools ( server , toolRegistrations ) ;
167- log ( 'info' , `✅ Registered ${ registeredTools . length } tools from '${ workflowName } '` ) ;
168- // registerTools() automatically sends tool list change notification internally
169- } else {
170- log (
171- 'info' ,
172- `All ${ toolsToRegister . length } tools from '${ workflowName } ' were already registered` ,
173- ) ;
174- }
175- } catch ( error ) {
176- log ( 'error' , `Failed to register tools from '${ workflowName } ': ${ error } ` ) ;
177- // Fallback to simplified tool registration one at a time
150+ // Register all tools using bulk registration
151+ if ( toolsToRegister . length > 0 ) {
178152 log (
179153 'info' ,
180- `🚀 Fallback: Enabling ${ toolsToRegister . length } tools individually from '${ workflowName } ' workflow` ,
154+ `🚀 Registering ${ toolsToRegister . length } tools from '${ workflowName } ' workflow` ,
181155 ) ;
182- for ( const toolToRegister of toolsToRegister ) {
183- try {
184- // Use the simplified registerTools method with single tool to avoid type complexity
185- const singleToolRegistration = [
186- {
187- name : toolToRegister . name ,
188- config : {
189- description : toolToRegister . config . description ,
190- inputSchema : toolToRegister . config . inputSchema as unknown , // Cast to unknown for SDK interface
191- } ,
192- // Adapt callback to match SDK's expected signature
193- callback : ( args : unknown ) : Promise < ToolResponse > =>
194- toolToRegister . callback ( args as Record < string , unknown > ) ,
195- } ,
196- ] ;
197- registerAndTrackTools ( server , singleToolRegistration ) ;
198- log ( 'debug' , `Registered tool: ${ toolToRegister . name } ` ) ;
199- } catch ( toolError ) {
200- log ( 'error' , `Failed to register tool '${ toolToRegister . name } ': ${ toolError } ` ) ;
201- }
202- }
203- log ( 'info' , `✅ Registered ${ toolsToRegister . length } tools from '${ workflowName } '` ) ;
156+
157+ // Convert to proper tool registration format
158+ const toolRegistrations = toolsToRegister . map ( ( tool ) => ( {
159+ name : tool . name ,
160+ config : {
161+ description : tool . config . description ,
162+ inputSchema : tool . config . inputSchema as unknown ,
163+ } ,
164+ callback : ( args : unknown ) : Promise < ToolResponse > =>
165+ tool . callback ( args as Record < string , unknown > ) ,
166+ } ) ) ;
167+
168+ // Use bulk registration - no fallback needed with proper duplicate handling
169+ const registeredTools = registerAndTrackTools ( server , toolRegistrations ) ;
170+ log ( 'info' , `✅ Registered ${ registeredTools . length } tools from '${ workflowName } '` ) ;
171+ } else {
172+ log ( 'info' , `No new tools to register from '${ workflowName } ' (all already registered)` ) ;
204173 }
205174
206175 // Track the workflow as enabled
@@ -213,8 +182,7 @@ export async function enableWorkflows(
213182 }
214183 }
215184
216- // No manual notification needed - registerTools() handles notifications automatically
217- log ( 'debug' , 'Tool list change notifications handled automatically by registerTools()' ) ;
185+ // registerAndTrackTools() handles tool list change notifications automatically
218186
219187 log (
220188 'info' ,
0 commit comments