From 6a990150b15399a255b95fe1c7e01e334ab41564 Mon Sep 17 00:00:00 2001 From: ryoppippi <1560508+ryoppippi@users.noreply.github.com> Date: Thu, 9 Oct 2025 14:01:23 +0100 Subject: [PATCH] feat(toolsets): add multi-account support for fetchTools Add support for fetching tools from multiple accounts via MCP: - Add `setAccounts()` method to configure default account IDs - Add `FetchToolsOptions.accountIds` to specify accounts per fetch - Options override `setAccounts()` when both provided - Each account gets separate MCP session with `x-account-id` header - Tools from all accounts are merged into single collection Tests: - Mock MCP server now respects `x-account-id` header - Test account-specific tool filtering - Test `setAccounts()` and `fetchTools()` override behavior - Verify tools from multiple accounts are correctly merged Related: unified-cloud-api MCP endpoint uses x-account-id header --- src/toolsets/stackone.ts | 67 ++++- src/toolsets/tests/stackone.mcp-fetch.spec.ts | 233 ++++++++++++++++-- 2 files changed, 280 insertions(+), 20 deletions(-) diff --git a/src/toolsets/stackone.ts b/src/toolsets/stackone.ts index 907b5a9..bf44271 100644 --- a/src/toolsets/stackone.ts +++ b/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 { + // 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 diff --git a/src/toolsets/tests/stackone.mcp-fetch.spec.ts b/src/toolsets/tests/stackone.mcp-fetch.spec.ts index 31609b8..6037f27 100644 --- a/src/toolsets/tests/stackone.mcp-fetch.spec.ts +++ b/src/toolsets/tests/stackone.mcp-fetch.spec.ts @@ -6,6 +6,7 @@ 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; @@ -13,27 +14,33 @@ type MockTool = { 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) { 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'); + }); +});