feat(toolsets): add multi-account support for fetchTools

src/toolsets/stackone.ts

@@ -1,5 +1,5 @@
 import { loadStackOneSpecs } from '../openapi/loader';
-import { StackOneTool, type Tools } from '../tool';
+import { StackOneTool, Tools } from '../tool';
 import type { ToolDefinition } from '../types';
 import { removeJsonSchemaProperty } from '../utils/schema';
 import { type BaseToolSetConfig, ToolSet, ToolSetConfigError } from './base';
@@ -14,6 +14,17 @@ export interface StackOneToolSetConfig extends BaseToolSetConfig {
   removedParams?: string[]; // List of parameters to remove from all tools
 }
 
+/**
+ * Options for filtering tools when fetching from MCP
+ */
+export interface FetchToolsOptions {
+  /**
+   * Filter tools by account IDs
+   * Only tools available on these accounts will be returned
+   */
+  accountIds?: string[];
+}
+
 /**
  * Configuration for workflow
  */
@@ -34,6 +45,7 @@ export class StackOneToolSet extends ToolSet {
    * Account ID for StackOne API
    */
   private accountId?: string;
+  private accountIds: string[] = [];
   private readonly _removedParams: string[];
 
   /**
@@ -103,6 +115,59 @@ export class StackOneToolSet extends ToolSet {
     return this.getTools(filterPattern, headers);
   }
 
+  /**
+   * Set account IDs for filtering tools
+   * @param accountIds Array of account IDs to filter tools by
+   * @returns This toolset instance for chaining
+   */
+  setAccounts(accountIds: string[]): this {
+    this.accountIds = accountIds;
+    return this;
+  }
+
+  /**
+   * Fetch tools from MCP with optional account ID filtering
+   * @param options Optional filtering options for account IDs
+   * @returns Collection of tools matching the filter criteria
+   *
+   * TODO: Add support for filtering by providers and actions
+   * - providers: Filter tools by provider names (e.g., ['hibob', 'bamboohr'])
+   * - actions: Filter tools by action patterns with glob support (e.g., ['*_list_employees'])
+   */
+  async fetchTools(options?: FetchToolsOptions): Promise<Tools> {
+    // Use account IDs from options, or fall back to instance state
+    const effectiveAccountIds = options?.accountIds || this.accountIds;
+
+    // If account IDs are specified, fetch tools for each account and merge
+    if (effectiveAccountIds.length > 0) {
+      const toolsPromises = effectiveAccountIds.map(async (accountId) => {
+        const headers = { 'x-account-id': accountId };
+        const mergedHeaders = { ...this.headers, ...headers };
+
+        // Create a temporary toolset instance with the account-specific headers
+        const tempHeaders = mergedHeaders;
+        const originalHeaders = this.headers;
+        this.headers = tempHeaders;
+
+        try {
+          const tools = await super.fetchTools();
+          return tools.toArray();
+        } finally {
+          // Restore original headers
+          this.headers = originalHeaders;
+        }
+      });
+
+      const toolArrays = await Promise.all(toolsPromises);
+      const allTools = toolArrays.flat();
+
+      return new Tools(allTools);
+    }
+
+    // No account filtering - fetch all tools
+    return await super.fetchTools();
+  }
+
   /**
    * Plan a workflow
    * @param config Configuration object containing workflow details

src/toolsets/tests/stackone.mcp-fetch.spec.ts

@@ -6,34 +6,41 @@ import { Hono } from 'hono';
 import { z } from 'zod';
 import { server as mswServer } from '../../../mocks/node';
 import { ToolSet } from '../base';
+import { StackOneToolSet } from '../stackone';
 
 type MockTool = {
   name: string;
   description?: string;
   shape: z.ZodRawShape;
 };
 
-async function createMockMcpServer(tools: MockTool[]) {
-  const mcp = new McpServer({ name: 'test-mcp', version: '1.0.0' });
-
-  for (const tool of tools) {
-    mcp.registerTool(
-      tool.name,
-      {
-        description: tool.description,
-        inputSchema: tool.shape,
-      },
-      async ({ params }) => ({
-        content: [],
-        structuredContent: params.arguments ?? {},
-        _meta: undefined,
-      })
-    );
-  }
-
+async function createMockMcpServer(accountTools: Record<string, MockTool[]>) {
   const app = new Hono();
+
   app.all('/mcp', async (c) => {
+    // Get account ID from header
+    const accountId = c.req.header('x-account-id') || 'default';
+    const tools = accountTools[accountId] || [];
+
+    // Create a new MCP server instance per account
+    const mcp = new McpServer({ name: 'test-mcp', version: '1.0.0' });
     const transport = new StreamableHTTPTransport();
+
+    for (const tool of tools) {
+      mcp.registerTool(
+        tool.name,
+        {
+          description: tool.description,
+          inputSchema: tool.shape,
+        },
+        async ({ params }) => ({
+          content: [],
+          structuredContent: params.arguments ?? {},
+          _meta: undefined,
+        })
+      );
+    }
+
     await mcp.connect(transport);
     return transport.handleRequest(c);
   });
@@ -66,7 +73,10 @@ describe('ToolSet.fetchTools (MCP + RPC integration)', () => {
     mswServer.close();
     restoreMsw = () => mswServer.listen({ onUnhandledRequest: 'warn' });
 
-    const server = await createMockMcpServer(mockTools);
+    const server = await createMockMcpServer({
+      default: mockTools,
+      'test-account': mockTools,
+    });
     origin = server.origin;
     closeServer = server.close;
   });
@@ -117,3 +127,188 @@ describe('ToolSet.fetchTools (MCP + RPC integration)', () => {
     expect(result).toEqual({ data: null });
   });
 });
+
+describe('StackOneToolSet account filtering', () => {
+  const acc1Tools = [
+    {
+      name: 'acc1_tool_1',
+      description: 'Account 1 Tool 1',
+      shape: { fields: z.string().optional() },
+    },
+    {
+      name: 'acc1_tool_2',
+      description: 'Account 1 Tool 2',
+      shape: { id: z.string() },
+    },
+  ] as const satisfies MockTool[];
+
+  const acc2Tools = [
+    {
+      name: 'acc2_tool_1',
+      description: 'Account 2 Tool 1',
+      shape: { fields: z.string().optional() },
+    },
+    {
+      name: 'acc2_tool_2',
+      description: 'Account 2 Tool 2',
+      shape: { id: z.string() },
+    },
+  ] as const satisfies MockTool[];
+
+  const acc3Tools = [
+    {
+      name: 'acc3_tool_1',
+      description: 'Account 3 Tool 1',
+      shape: { fields: z.string().optional() },
+    },
+  ] as const satisfies MockTool[];
+
+  const defaultTools = [
+    {
+      name: 'default_tool_1',
+      description: 'Default Tool 1',
+      shape: { fields: z.string().optional() },
+    },
+    {
+      name: 'default_tool_2',
+      description: 'Default Tool 2',
+      shape: { id: z.string() },
+    },
+  ] as const satisfies MockTool[];
+
+  let origin: string;
+  let closeServer: () => void;
+  let restoreMsw: (() => void) | undefined;
+
+  beforeAll(async () => {
+    mswServer.close();
+    restoreMsw = () => mswServer.listen({ onUnhandledRequest: 'warn' });
+
+    const server = await createMockMcpServer({
+      default: defaultTools,
+      acc1: acc1Tools,
+      acc2: acc2Tools,
+      acc3: acc3Tools,
+    });
+    origin = server.origin;
+    closeServer = server.close;
+  });
+
+  afterAll(() => {
+    closeServer();
+    restoreMsw?.();
+  });
+
+  it('supports setAccounts() for chaining', async () => {
+    const stackOneClient = {
+      actions: {
+        rpcAction: mock(async () => ({ actionsRpcResponse: { data: null } })),
+      },
+    } as unknown as StackOne;
+
+    const toolset = new StackOneToolSet({
+      baseUrl: origin,
+      apiKey: 'test-key',
+      stackOneClient,
+    });
+
+    // Test chaining
+    const result = toolset.setAccounts(['acc1', 'acc2']);
+    expect(result).toBe(toolset);
+  });
+
+  it('fetches tools without account filtering when no accountIds provided', async () => {
+    const stackOneClient = {
+      actions: {
+        rpcAction: mock(async () => ({ actionsRpcResponse: { data: null } })),
+      },
+    } as unknown as StackOne;
+
+    const toolset = new StackOneToolSet({
+      baseUrl: origin,
+      apiKey: 'test-key',
+      stackOneClient,
+    });
+
+    const tools = await toolset.fetchTools();
+    expect(tools.length).toBe(2);
+    const toolNames = tools.toArray().map((t) => t.name);
+    expect(toolNames).toContain('default_tool_1');
+    expect(toolNames).toContain('default_tool_2');
+  });
+
+  it('uses x-account-id header when fetching tools with accountIds', async () => {
+    const stackOneClient = {
+      actions: {
+        rpcAction: mock(async () => ({ actionsRpcResponse: { data: null } })),
+      },
+    } as unknown as StackOne;
+
+    const toolset = new StackOneToolSet({
+      baseUrl: origin,
+      apiKey: 'test-key',
+      stackOneClient,
+    });
+
+    // Fetch tools for acc1
+    const tools = await toolset.fetchTools({ accountIds: ['acc1'] });
+    expect(tools.length).toBe(2);
+    const toolNames = tools.toArray().map((t) => t.name);
+    expect(toolNames).toContain('acc1_tool_1');
+    expect(toolNames).toContain('acc1_tool_2');
+  });
+
+  it('uses setAccounts when no accountIds provided in fetchTools', async () => {
+    const stackOneClient = {
+      actions: {
+        rpcAction: mock(async () => ({ actionsRpcResponse: { data: null } })),
+      },
+    } as unknown as StackOne;
+
+    const toolset = new StackOneToolSet({
+      baseUrl: origin,
+      apiKey: 'test-key',
+      stackOneClient,
+    });
+
+    // Set accounts using setAccounts
+    toolset.setAccounts(['acc1', 'acc2']);
+
+    // Fetch without accountIds - should use setAccounts
+    const tools = await toolset.fetchTools();
+
+    // Should fetch tools for 2 accounts from setAccounts
+    // acc1 has 2 tools, acc2 has 2 tools, so total should be 4
+    expect(tools.length).toBe(4);
+    const toolNames = tools.toArray().map((t) => t.name);
+    expect(toolNames).toContain('acc1_tool_1');
+    expect(toolNames).toContain('acc1_tool_2');
+    expect(toolNames).toContain('acc2_tool_1');
+    expect(toolNames).toContain('acc2_tool_2');
+  });
+
+  it('overrides setAccounts when accountIds provided in fetchTools', async () => {
+    const stackOneClient = {
+      actions: {
+        rpcAction: mock(async () => ({ actionsRpcResponse: { data: null } })),
+      },
+    } as unknown as StackOne;
+
+    const toolset = new StackOneToolSet({
+      baseUrl: origin,
+      apiKey: 'test-key',
+      stackOneClient,
+    });
+
+    // Set accounts using setAccounts
+    toolset.setAccounts(['acc1', 'acc2']);
+
+    // Fetch with accountIds - should override setAccounts
+    const tools = await toolset.fetchTools({ accountIds: ['acc3'] });
+
+    // Should fetch tools only for acc3 (ignoring acc1, acc2)
+    expect(tools.length).toBe(1);
+    const toolNames = tools.toArray().map((t) => t.name);
+    expect(toolNames).toContain('acc3_tool_1');
+  });
+});