fix(workflow-executor): match MCP tools by config id

packages/ai-proxy/src/forest-integration-client.ts

@@ -14,6 +14,7 @@ export type CustomConfig = ZendeskConfig | KolarConfig | SnowflakeConfig;
 export type ForestIntegrationName = 'Zendesk' | 'Kolar' | 'Snowflake';
 
 export interface ForestIntegrationConfig {
+  id?: string;
   integrationName: ForestIntegrationName;
   config: CustomConfig;
   isForestConnector: true;
@@ -39,16 +40,16 @@ export default class ForestIntegrationClient implements ToolProvider {
   async loadTools(): Promise<RemoteTool[]> {
     const tools: RemoteTool[] = [];
 
-    this.configs.forEach(({ integrationName, config }) => {
+    this.configs.forEach(({ id: mcpServerId, integrationName, config }) => {
       switch (integrationName) {
         case 'Zendesk':
-          tools.push(...getZendeskTools(config as ZendeskConfig));
+          tools.push(...getZendeskTools(config as ZendeskConfig, mcpServerId));
           break;
         case 'Kolar':
-          tools.push(...getKolarTools(config as KolarConfig));
+          tools.push(...getKolarTools(config as KolarConfig, mcpServerId));
           break;
         case 'Snowflake':
-          tools.push(...getSnowflakeTools(config as SnowflakeConfig));
+          tools.push(...getSnowflakeTools(config as SnowflakeConfig, mcpServerId));
           break;
         default:
           this.logger?.('Warn', `Unsupported integration: ${integrationName}`);

packages/ai-proxy/src/integrations/kolar/tools.ts

@@ -11,7 +11,7 @@ export interface KolarConfig {
   apiKey: string;
 }
 
-export default function getKolarTools(config: KolarConfig): RemoteTool[] {
+export default function getKolarTools(config: KolarConfig, mcpServerId?: string): RemoteTool[] {
   const { baseUrl, headers } = getKolarConfig(config);
 
   return [
@@ -23,6 +23,7 @@ export default function getKolarTools(config: KolarConfig): RemoteTool[] {
     tool =>
       new ServerRemoteTool({
         sourceId: 'kolar',
+        mcpServerId,
         tool,
       }),
   );

packages/ai-proxy/src/integrations/snowflake/tools.ts

@@ -15,7 +15,10 @@ export interface SnowflakeConfig {
   defaultRole?: string;
 }
 
-export default function getSnowflakeTools(config: SnowflakeConfig): RemoteTool[] {
+export default function getSnowflakeTools(
+  config: SnowflakeConfig,
+  mcpServerId?: string,
+): RemoteTool[] {
   const headers = getSnowflakeAuthHeaders(config);
   const baseUrl = getSnowflakeBaseUrl(config.accountIdentifier);
 
@@ -27,6 +30,7 @@ export default function getSnowflakeTools(config: SnowflakeConfig): RemoteTool[]
     tool =>
       new ServerRemoteTool({
         sourceId: 'snowflake',
+        mcpServerId,
         tool,
       }),
   );

packages/ai-proxy/src/integrations/zendesk/tools.ts

@@ -15,7 +15,7 @@ export interface ZendeskConfig {
   apiToken: string;
 }
 
-export default function getZendeskTools(config: ZendeskConfig): RemoteTool[] {
+export default function getZendeskTools(config: ZendeskConfig, mcpServerId?: string): RemoteTool[] {
   const { baseUrl, headers } = getZendeskConfig(config);
 
   return [
@@ -29,6 +29,7 @@ export default function getZendeskTools(config: ZendeskConfig): RemoteTool[] {
     tool =>
       new ServerRemoteTool({
         sourceId: 'zendesk',
+        mcpServerId,
         tool,
       }),
   );

packages/ai-proxy/src/mcp-client.ts

@@ -8,23 +8,29 @@ import McpServerRemoteTool from './mcp-server-remote-tool';
 
 export type McpServers = MultiServerMCPClient['config']['mcpServers'];
 
-export type McpServerConfig = MultiServerMCPClient['config']['mcpServers'][string];
+export type McpServerConfig = MultiServerMCPClient['config']['mcpServers'][string] & {
+  id?: string;
+};
 
 export type McpConfiguration = {
-  configs: McpServers;
+  configs: Record<string, McpServerConfig>;
 } & Omit<MultiServerMCPClient['config'], 'mcpServers'>;
 
 export default class McpClient implements ToolProvider {
   private readonly mcpClients: Record<string, MultiServerMCPClient> = {};
+  private readonly mcpServerIdsByName: Record<string, string | undefined> = {};
   private readonly logger?: Logger;
 
   constructor(config: McpConfiguration, logger?: Logger) {
     this.logger = logger;
     // split the config into several clients to be more resilient
     // if a mcp server is down, the others will still work
     Object.entries(config.configs).forEach(([name, serverConfig]) => {
+      const { id: mcpServerId, ...rest } = serverConfig as McpServerConfig &
+        Record<string, unknown>;
+      this.mcpServerIdsByName[name] = mcpServerId;
       this.mcpClients[name] = new MultiServerMCPClient({
-        mcpServers: { [name]: serverConfig },
+        mcpServers: { [name]: rest as McpServerConfig },
         ...config,
       });
     });
@@ -39,7 +45,12 @@ export default class McpClient implements ToolProvider {
         try {
           const loadedTools = (await client.getTools()) ?? [];
           const extendedTools = loadedTools.map(
-            tool => new McpServerRemoteTool({ tool, sourceId: name }),
+            tool =>
+              new McpServerRemoteTool({
+                tool,
+                sourceId: name,
+                mcpServerId: this.mcpServerIdsByName[name],
+              }),
           );
           tools.push(...extendedTools);
         } catch (error) {

packages/ai-proxy/src/mcp-server-remote-tool.ts

@@ -3,7 +3,11 @@ import type { StructuredToolInterface } from '@langchain/core/tools';
 import RemoteTool from './remote-tool';
 
 export default class McpServerRemoteTool<ToolType = unknown> extends RemoteTool {
-  constructor(options: { tool: StructuredToolInterface<ToolType>; sourceId?: string }) {
+  constructor(options: {
+    tool: StructuredToolInterface<ToolType>;
+    sourceId?: string;
+    mcpServerId?: string;
+  }) {
     super({ ...options, sourceType: 'mcp-server' });
   }
 }

packages/ai-proxy/src/remote-tool.ts

@@ -4,15 +4,20 @@ export default abstract class RemoteTool<ToolType = unknown> {
   base: StructuredToolInterface<ToolType>;
   sourceId: string;
   sourceType: string;
+  // Shared across McpServerRemoteTool and ServerRemoteTool because the orchestrator
+  // stores Forest connectors alongside user MCP servers in the same ai_mcp_configs table.
+  mcpServerId?: string;
 
   constructor(options: {
     tool: StructuredToolInterface<ToolType>;
     sourceId?: string;
     sourceType?: string;
+    mcpServerId?: string;
   }) {
     this.base = options.tool;
     this.sourceId = options.sourceId;
     this.sourceType = options.sourceType;
+    this.mcpServerId = options.mcpServerId;
   }
 
   get sanitizedName() {

packages/ai-proxy/src/server-remote-tool.ts

@@ -3,7 +3,11 @@ import type { StructuredToolInterface } from '@langchain/core/tools';
 import RemoteTool from './remote-tool';
 
 export default class ServerRemoteTool<ToolType = unknown> extends RemoteTool {
-  constructor(options: { tool: StructuredToolInterface<ToolType>; sourceId?: string }) {
+  constructor(options: {
+    tool: StructuredToolInterface<ToolType>;
+    sourceId?: string;
+    mcpServerId?: string;
+  }) {
     super({ ...options, sourceType: 'server' });
   }
 }

packages/ai-proxy/test/forest-integration-client.test.ts

@@ -1,6 +1,9 @@
 import ForestIntegrationClient from '../src/forest-integration-client';
+import getKolarTools from '../src/integrations/kolar/tools';
 import { validateKolarConfig } from '../src/integrations/kolar/utils';
+import getSnowflakeTools from '../src/integrations/snowflake/tools';
 import { validateSnowflakeConfig } from '../src/integrations/snowflake/utils';
+import getZendeskTools from '../src/integrations/zendesk/tools';
 import { validateZendeskConfig } from '../src/integrations/zendesk/utils';
 
 const mockZendeskTools = [{ name: 'zendesk_get_tickets' }, { name: 'zendesk_get_ticket' }];
@@ -37,6 +40,7 @@ describe('ForestIntegrationClient', () => {
     it('should load zendesk tools when integration is zendesk', async () => {
       const client = new ForestIntegrationClient([
         {
+          id: '1',
           integrationName: 'Zendesk',
           config: { subdomain: 'test', email: 'a@b.com', apiToken: 'tok' },
           isForestConnector: true,
@@ -52,7 +56,7 @@ describe('ForestIntegrationClient', () => {
       const logger = jest.fn();
       const client = new ForestIntegrationClient(
         // @ts-expect-error Testing unsupported integration
-        [{ integrationName: 'unknown', config: {} as any, isForestConnector: true }],
+        [{ id: '1', integrationName: 'unknown', config: {} as any, isForestConnector: true }],
         logger,
       );
 
@@ -64,6 +68,7 @@ describe('ForestIntegrationClient', () => {
     it('should load kolar tools when integration is Kolar', async () => {
       const client = new ForestIntegrationClient([
         {
+          id: '1',
           integrationName: 'Kolar',
           config: { apiKey: 'key' },
           isForestConnector: true,
@@ -78,6 +83,7 @@ describe('ForestIntegrationClient', () => {
     it('should load snowflake tools when integration is Snowflake', async () => {
       const client = new ForestIntegrationClient([
         {
+          id: '1',
           integrationName: 'Snowflake',
           config: {
             accountIdentifier: 'a',
@@ -101,11 +107,13 @@ describe('ForestIntegrationClient', () => {
     it('should load tools from multiple configs', async () => {
       const client = new ForestIntegrationClient([
         {
+          id: '1',
           integrationName: 'Zendesk',
           config: { subdomain: 'a', email: 'a@b.com', apiToken: 'tok' },
           isForestConnector: true,
         },
         {
+          id: '2',
           integrationName: 'Zendesk',
           config: { subdomain: 'b', email: 'c@d.com', apiToken: 'tok2' },
           isForestConnector: true,
@@ -122,7 +130,7 @@ describe('ForestIntegrationClient', () => {
     it('should call validateZendeskConfig for Zendesk integration', async () => {
       const zendeskConfig = { subdomain: 'test', email: 'a@b.com', apiToken: 'tok' };
       const client = new ForestIntegrationClient([
-        { integrationName: 'Zendesk', config: zendeskConfig, isForestConnector: true },
+        { id: '1', integrationName: 'Zendesk', config: zendeskConfig, isForestConnector: true },
       ]);
 
       await client.checkConnection();
@@ -133,6 +141,7 @@ describe('ForestIntegrationClient', () => {
     it('should return true on success', async () => {
       const client = new ForestIntegrationClient([
         {
+          id: '1',
           integrationName: 'Zendesk',
           config: { subdomain: 'test', email: 'a@b.com', apiToken: 'tok' },
           isForestConnector: true,
@@ -147,7 +156,7 @@ describe('ForestIntegrationClient', () => {
     it('should call validateKolarConfig for Kolar integration', async () => {
       const kolarConfig = { apiKey: 'key' };
       const client = new ForestIntegrationClient([
-        { integrationName: 'Kolar', config: kolarConfig, isForestConnector: true },
+        { id: '1', integrationName: 'Kolar', config: kolarConfig, isForestConnector: true },
       ]);
 
       await client.checkConnection();
@@ -161,7 +170,7 @@ describe('ForestIntegrationClient', () => {
         programmaticAccessToken: 'tok',
       };
       const client = new ForestIntegrationClient([
-        { integrationName: 'Snowflake', config: snowflakeConfig, isForestConnector: true },
+        { id: '1', integrationName: 'Snowflake', config: snowflakeConfig, isForestConnector: true },
       ]);
 
       await client.checkConnection();
@@ -172,12 +181,10 @@ describe('ForestIntegrationClient', () => {
     it('should throw for unsupported integration', async () => {
       const client = new ForestIntegrationClient([
         // @ts-expect-error Testing unsupported integration
-        { integrationName: 'Unknown', config: {}, isForestConnector: true },
+        { id: '1', integrationName: 'Unknown', config: {}, isForestConnector: true },
       ]);
 
-      await expect(client.checkConnection()).rejects.toThrow(
-        'Unsupported integration: Unknown',
-      );
+      await expect(client.checkConnection()).rejects.toThrow('Unsupported integration: Unknown');
     });
   });
 
@@ -188,4 +195,77 @@ describe('ForestIntegrationClient', () => {
       await expect(client.dispose()).resolves.toBeUndefined();
     });
   });
+
+  describe('ForestIntegrationConfig.id is threaded as mcpServerId into integration tool factories', () => {
+    it('passes the config id to getZendeskTools so produced tools can be matched by step.mcpServerId', async () => {
+      const client = new ForestIntegrationClient([
+        {
+          id: 'forest-zendesk-42',
+          integrationName: 'Zendesk',
+          config: { subdomain: 'test', email: 'a@b.com', apiToken: 'tok' },
+          isForestConnector: true,
+        },
+      ]);
+
+      await client.loadTools();
+
+      expect(getZendeskTools).toHaveBeenCalledWith(
+        expect.objectContaining({ subdomain: 'test' }),
+        'forest-zendesk-42',
+      );
+    });
+
+    it('passes the config id to getKolarTools', async () => {
+      const client = new ForestIntegrationClient([
+        {
+          id: 'forest-kolar-7',
+          integrationName: 'Kolar',
+          config: { apiKey: 'key' },
+          isForestConnector: true,
+        },
+      ]);
+
+      await client.loadTools();
+
+      expect(getKolarTools).toHaveBeenCalledWith(
+        expect.objectContaining({ apiKey: 'key' }),
+        'forest-kolar-7',
+      );
+    });
+
+    it('passes the config id to getSnowflakeTools', async () => {
+      const client = new ForestIntegrationClient([
+        {
+          id: 'forest-snowflake-99',
+          integrationName: 'Snowflake',
+          config: { accountIdentifier: 'a', programmaticAccessToken: 'tok' },
+          isForestConnector: true,
+        },
+      ]);
+
+      await client.loadTools();
+
+      expect(getSnowflakeTools).toHaveBeenCalledWith(
+        expect.objectContaining({ accountIdentifier: 'a' }),
+        'forest-snowflake-99',
+      );
+    });
+
+    it('passes undefined to the factory when the config entry has no id', async () => {
+      const client = new ForestIntegrationClient([
+        {
+          integrationName: 'Zendesk',
+          config: { subdomain: 'test', email: 'a@b.com', apiToken: 'tok' },
+          isForestConnector: true,
+        },
+      ]);
+
+      await client.loadTools();
+
+      expect(getZendeskTools).toHaveBeenCalledWith(
+        expect.objectContaining({ subdomain: 'test' }),
+        undefined,
+      );
+    });
+  });
 });

packages/ai-proxy/test/integrations/kolar/tools.test.ts

@@ -26,4 +26,23 @@ describe('getKolarTools', () => {
       'kolar_get_screening_result',
     ]);
   });
+
+  describe('mcpServerId propagation', () => {
+    it('sets RemoteTool.mcpServerId on every produced tool when mcpServerId is provided', () => {
+      const tools = getKolarTools(config, 'forest-kolar-7');
+
+      expect(tools).not.toHaveLength(0);
+      tools.forEach(tool => {
+        expect(tool.mcpServerId).toBe('forest-kolar-7');
+      });
+    });
+
+    it('leaves RemoteTool.mcpServerId undefined when no mcpServerId is provided', () => {
+      const tools = getKolarTools(config);
+
+      tools.forEach(tool => {
+        expect(tool.mcpServerId).toBeUndefined();
+      });
+    });
+  });
 });