diff --git a/packages/_example/.env.example b/packages/_example/.env.example index 7c4ebaf6e7..85da10b99d 100644 --- a/packages/_example/.env.example +++ b/packages/_example/.env.example @@ -21,6 +21,8 @@ FOREST_AUTH_SECRET= EXECUTOR_AGENT_URL=http://localhost:3351 WORKFLOW_EXECUTOR_URL=http://localhost:3400 EXECUTOR_DATABASE_URL=postgresql://executor:password@localhost:5459/workflow_executor +# At-rest encryption key for secrets the executor stores (openssl rand -hex 32; same value across instances) +FOREST_EXECUTOR_ENCRYPTION_KEY= # when start:with-executor:multiple-instances command # EXECUTOR_AGENT_URL=http://host.docker.internal:3351 # WORKFLOW_EXECUTOR_URL=http://localhost:3400 diff --git a/packages/ai-proxy/src/ai-client.ts b/packages/ai-proxy/src/ai-client.ts index 87feb1f405..8ebbc5890e 100644 --- a/packages/ai-proxy/src/ai-client.ts +++ b/packages/ai-proxy/src/ai-client.ts @@ -1,3 +1,4 @@ +import type { McpServerLoadFailure } from './mcp-client'; import type { AiConfiguration } from './provider'; import type RemoteTool from './remote-tool'; import type { ToolProvider } from './tool-provider'; @@ -39,13 +40,31 @@ export class AiClient { } async loadRemoteTools(configs: Record): Promise { + return (await this.loadRemoteToolsWithFailures(configs)).tools; + } + + // Same load as loadRemoteTools, but also returns the classified per-server failures providers + // surface (only MCP providers do today). The default loadRemoteTools drops them, so existing + // consumers are unaffected. + async loadRemoteToolsWithFailures( + configs: Record, + ): Promise<{ tools: RemoteTool[]; failures: McpServerLoadFailure[] }> { await this.disposeToolProviders('Error closing previous remote tool connection'); const providers = createToolProviders(configs, this.logger); - const toolsByProvider = await Promise.all(providers.map(p => p.loadTools())); + const resultsByProvider = await Promise.all( + providers.map(async provider => + provider.loadToolsWithFailures + ? provider.loadToolsWithFailures() + : { tools: await provider.loadTools(), failures: [] }, + ), + ); this.toolProviders = providers; - return toolsByProvider.flat(); + return { + tools: resultsByProvider.flatMap(result => result.tools), + failures: resultsByProvider.flatMap(result => result.failures), + }; } async closeConnections(): Promise { diff --git a/packages/ai-proxy/src/index.ts b/packages/ai-proxy/src/index.ts index f0bf0a6e4b..197797caa1 100644 --- a/packages/ai-proxy/src/index.ts +++ b/packages/ai-proxy/src/index.ts @@ -20,6 +20,7 @@ export * from './remote-tools'; export { default as RemoteTool } from './remote-tool'; export * from './router'; export * from './mcp-client'; +export * from './mcp-auth-error'; export * from './oauth-token-injector'; export * from './errors'; export * from './tool-provider'; diff --git a/packages/ai-proxy/src/mcp-auth-error.ts b/packages/ai-proxy/src/mcp-auth-error.ts new file mode 100644 index 0000000000..187921a00b --- /dev/null +++ b/packages/ai-proxy/src/mcp-auth-error.ts @@ -0,0 +1,61 @@ +// Classifies errors surfaced while connecting to or calling an MCP server. Only 401 (the token was +// rejected) is a refreshable auth failure; 403 is a permission/scope problem a token refresh or +// re-consent cannot resolve, so it is left to surface as an ordinary failure. The MCP SDK / HTTP +// transport reports failures in several shapes (a numeric status field, or only a message string), +// so the checks walk the cause chain and inspect both structured status and the message text. +const AUTH_STATUSES = new Set([401]); +const AUTH_PATTERN = /\b401\b|unauthorized/i; +const CONNECTION_PATTERN = + /econnrefused|econnreset|etimedout|enotfound|eai_again|fetch failed|network|socket|timeout|connect/i; + +export type McpLoadFailureKind = 'auth' | 'connection' | 'unknown'; + +function statusOf(value: unknown): number | undefined { + const candidate = value as { code?: unknown; status?: unknown; statusCode?: unknown }; + + for (const field of [candidate?.code, candidate?.status, candidate?.statusCode]) { + if (typeof field === 'number') return field; + } + + return undefined; +} + +function messageOf(value: unknown): string { + if (value instanceof Error) return value.message; + if (typeof value === 'string') return value; + + return ''; +} + +function errorChain(error: unknown): unknown[] { + const links: unknown[] = []; + let current: unknown = error; + + while (current && links.length < 10 && !links.includes(current)) { + links.push(current); + current = (current as { cause?: unknown }).cause; + } + + return links; +} + +export function isMcpAuthError(error: unknown): boolean { + return errorChain(error).some(link => { + const status = statusOf(link); + // Only an HTTP-range status is authoritative (a 403 isn't refreshable even if its message says + // "unauthorized"). A JSON-RPC error code like -32603 isn't a status, so fall back to the message. + if (status !== undefined && status >= 100 && status <= 599) return AUTH_STATUSES.has(status); + + return AUTH_PATTERN.test(messageOf(link)); + }); +} + +export function classifyMcpLoadError(error: unknown): McpLoadFailureKind { + if (isMcpAuthError(error)) return 'auth'; + + const isConnectionFailure = errorChain(error).some(link => + CONNECTION_PATTERN.test(messageOf(link)), + ); + + return isConnectionFailure ? 'connection' : 'unknown'; +} diff --git a/packages/ai-proxy/src/mcp-client.ts b/packages/ai-proxy/src/mcp-client.ts index bb834c576e..0008a1bf8b 100644 --- a/packages/ai-proxy/src/mcp-client.ts +++ b/packages/ai-proxy/src/mcp-client.ts @@ -4,18 +4,28 @@ import type { Logger } from '@forestadmin/datasource-toolkit'; import { MultiServerMCPClient } from '@langchain/mcp-adapters'; import { McpConnectionError } from './errors'; +import { type McpLoadFailureKind, classifyMcpLoadError } from './mcp-auth-error'; import McpServerRemoteTool from './mcp-server-remote-tool'; export type McpServers = MultiServerMCPClient['config']['mcpServers']; export type McpServerConfig = MultiServerMCPClient['config']['mcpServers'][string] & { id?: string; + // Executor-side routing hint served by the orchestrator; stripped before reaching the SDK. + authType?: string; }; export type McpConfiguration = { configs: Record; } & Omit; +export interface McpServerLoadFailure { + server: string; + mcpServerId?: string; + kind: McpLoadFailureKind; + error: Error; +} + export default class McpClient implements ToolProvider { private readonly mcpClients: Record = {}; private readonly mcpServerIdsByName: Record = {}; @@ -26,8 +36,11 @@ export default class McpClient implements ToolProvider { // 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; + const { + id: mcpServerId, + authType, + ...rest + } = serverConfig as McpServerConfig & Record; this.mcpServerIdsByName[name] = mcpServerId; this.mcpClients[name] = new MultiServerMCPClient({ mcpServers: { [name]: rest as McpServerConfig }, @@ -36,9 +49,15 @@ export default class McpClient implements ToolProvider { }); } - async loadTools(): Promise { + // Exposes per-server failures classified by cause (auth vs connection) alongside the tools that + // did load, so a caller holding a per-user token can tell a revoked token (retry after refresh) + // from an unreachable server (genuine failure). loadTools() keeps its tools-only contract. + async loadToolsWithFailures(): Promise<{ + tools: McpServerRemoteTool[]; + failures: McpServerLoadFailure[]; + }> { const tools: McpServerRemoteTool[] = []; - const errors: Array<{ server: string; error: Error }> = []; + const failures: McpServerLoadFailure[] = []; await Promise.all( Object.entries(this.mcpClients).map(async ([name, client]) => { @@ -55,22 +74,31 @@ export default class McpClient implements ToolProvider { tools.push(...extendedTools); } catch (error) { this.logger?.('Error', `Error loading tools for ${name}`, error as Error); - errors.push({ server: name, error: error as Error }); + failures.push({ + server: name, + mcpServerId: this.mcpServerIdsByName[name], + kind: classifyMcpLoadError(error), + error: error as Error, + }); } }), ); // Surface partial failures to provide better feedback - if (errors.length > 0) { - const errorMessage = errors.map(e => `${e.server}: ${e.error.message}`).join('; '); + if (failures.length > 0) { + const errorMessage = failures.map(f => `${f.server}: ${f.error.message}`).join('; '); this.logger?.( 'Error', - `Failed to load tools from ${errors.length}/${Object.keys(this.mcpClients).length} ` + + `Failed to load tools from ${failures.length}/${Object.keys(this.mcpClients).length} ` + `MCP server(s): ${errorMessage}`, ); } - return tools; + return { tools, failures }; + } + + async loadTools(): Promise { + return (await this.loadToolsWithFailures()).tools; } async checkConnection(): Promise { diff --git a/packages/ai-proxy/src/tool-provider.ts b/packages/ai-proxy/src/tool-provider.ts index ed2a2408fe..a7d0eff4df 100644 --- a/packages/ai-proxy/src/tool-provider.ts +++ b/packages/ai-proxy/src/tool-provider.ts @@ -1,7 +1,10 @@ +import type { McpServerLoadFailure } from './mcp-client'; import type RemoteTool from './remote-tool'; export interface ToolProvider { loadTools(): Promise; + // Optional richer variant: providers that can classify per-server failures expose them here. + loadToolsWithFailures?(): Promise<{ tools: RemoteTool[]; failures: McpServerLoadFailure[] }>; checkConnection(): Promise; dispose(): Promise; } diff --git a/packages/ai-proxy/test/ai-client.test.ts b/packages/ai-proxy/test/ai-client.test.ts index d3919dfc3b..a56572fb2e 100644 --- a/packages/ai-proxy/test/ai-client.test.ts +++ b/packages/ai-proxy/test/ai-client.test.ts @@ -242,6 +242,46 @@ describe('loadRemoteTools', () => { }); }); +describe('loadRemoteToolsWithFailures', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('aggregates tools and classified failures from providers that expose them', async () => { + const mcpTool = { name: 'mcp-tool' }; + const failure = { + server: 'slack', + mcpServerId: 'srv-a', + kind: 'auth' as const, + error: new Error('401'), + }; + mockedCreateToolProviders.mockReturnValue([ + mockProvider({ + loadToolsWithFailures: jest + .fn() + .mockResolvedValue({ tools: [mcpTool], failures: [failure] }), + }), + ]); + + const result = await new AiClient({}).loadRemoteToolsWithFailures({} as never); + + expect(result.tools).toEqual([mcpTool]); + expect(result.failures).toEqual([failure]); + }); + + it('falls back to loadTools with no failures for providers that do not classify', async () => { + const integrationTool = { name: 'zendesk-tool' }; + mockedCreateToolProviders.mockReturnValue([ + mockProvider({ loadTools: jest.fn().mockResolvedValue([integrationTool]) }), + ]); + + const result = await new AiClient({}).loadRemoteToolsWithFailures({} as never); + + expect(result.tools).toEqual([integrationTool]); + expect(result.failures).toEqual([]); + }); +}); + describe('closeConnections', () => { beforeEach(() => { jest.clearAllMocks(); diff --git a/packages/ai-proxy/test/mcp-auth-error.test.ts b/packages/ai-proxy/test/mcp-auth-error.test.ts new file mode 100644 index 0000000000..f936a11140 --- /dev/null +++ b/packages/ai-proxy/test/mcp-auth-error.test.ts @@ -0,0 +1,64 @@ +import { classifyMcpLoadError, isMcpAuthError } from '../src/mcp-auth-error'; + +function withCause(message: string, cause: unknown): Error { + const error = new Error(message); + (error as { cause?: unknown }).cause = cause; + + return error; +} + +describe('isMcpAuthError', () => { + it('detects a 401 numeric status field', () => { + expect(isMcpAuthError({ code: 401 })).toBe(true); + }); + + it('detects 401 / unauthorized in the message', () => { + expect(isMcpAuthError(new Error('Request failed with status code 401'))).toBe(true); + expect(isMcpAuthError(new Error('Unauthorized'))).toBe(true); + }); + + it('walks the cause chain', () => { + expect( + isMcpAuthError(withCause('wrapper', Object.assign(new Error('inner'), { status: 401 }))), + ).toBe(true); + }); + + it('returns false for 403 (forbidden), non-auth errors, and nullish input', () => { + expect(isMcpAuthError(Object.assign(new Error('denied'), { status: 403 }))).toBe(false); + expect(isMcpAuthError(new Error('403 Forbidden'))).toBe(false); + expect(isMcpAuthError(new Error('ECONNREFUSED'))).toBe(false); + expect(isMcpAuthError(new Error('500 Internal Server Error'))).toBe(false); + expect(isMcpAuthError(undefined)).toBe(false); + }); + + it('lets an explicit status win over an "unauthorized" message (a 403 stays non-auth)', () => { + expect(isMcpAuthError(Object.assign(new Error('Unauthorized scope'), { status: 403 }))).toBe( + false, + ); + }); + + it('classifies a JSON-RPC McpError by message, ignoring its non-HTTP code (e.g. -32603)', () => { + expect(isMcpAuthError(Object.assign(new Error('Unauthorized'), { code: -32603 }))).toBe(true); + expect(isMcpAuthError(Object.assign(new Error('boom'), { code: -32603 }))).toBe(false); + }); +}); + +describe('classifyMcpLoadError', () => { + it("classifies a 401 as 'auth'", () => { + expect(classifyMcpLoadError(new Error('HTTP 401 Unauthorized'))).toBe('auth'); + expect(classifyMcpLoadError({ status: 401 })).toBe('auth'); + }); + + it("classifies network failures as 'connection'", () => { + expect(classifyMcpLoadError(new Error('connect ECONNREFUSED 127.0.0.1:3000'))).toBe( + 'connection', + ); + expect(classifyMcpLoadError(new Error('fetch failed'))).toBe('connection'); + expect(classifyMcpLoadError(new Error('socket hang up'))).toBe('connection'); + }); + + it("classifies a 403 (forbidden) and anything else as 'unknown'", () => { + expect(classifyMcpLoadError(new Error('HTTP 403 Forbidden'))).toBe('unknown'); + expect(classifyMcpLoadError(new Error('tool schema invalid'))).toBe('unknown'); + }); +}); diff --git a/packages/ai-proxy/test/mcp-client.test.ts b/packages/ai-proxy/test/mcp-client.test.ts index ad128c2fbd..22cd950a57 100644 --- a/packages/ai-proxy/test/mcp-client.test.ts +++ b/packages/ai-proxy/test/mcp-client.test.ts @@ -1,6 +1,7 @@ import type { McpConfiguration } from '../src'; import { tool } from '@langchain/core/tools'; +import { MultiServerMCPClient } from '@langchain/mcp-adapters'; import { McpConnectionError } from '../src'; import McpClient from '../src/mcp-client'; @@ -487,3 +488,96 @@ describe('McpClient', () => { }); }); }); + +// Additive auth-error classification: loadToolsWithFailures exposes per-server failures alongside +// the loaded tools so a token holder can tell a revoked token from an unreachable server. +describe('McpClient.loadToolsWithFailures', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + const makeTool = (name: string) => + tool(() => {}, { name, description: name, schema: undefined, responseFormat: 'content' }); + + const singleServer = (id?: string) => + ({ + configs: { + slack: { transport: 'stdio', command: 'npx', args: [], env: {}, ...(id ? { id } : {}) }, + }, + } as unknown as McpConfiguration); + + it('returns the loaded tools and no failures when every server loads', async () => { + getToolsMock.mockResolvedValue([makeTool('t1')]); + + const result = await new McpClient(singleServer()).loadToolsWithFailures(); + + expect(result.tools).toHaveLength(1); + expect(result.failures).toEqual([]); + }); + + it("classifies a 401 as an 'auth' failure tagged with server and mcpServerId, keeping other tools", async () => { + getToolsMock + .mockRejectedValueOnce(new Error('Request failed: 401 Unauthorized')) + .mockResolvedValueOnce([makeTool('t2')]); + const client = new McpClient({ + configs: { + slack: { transport: 'stdio', command: 'npx', args: [], env: {}, id: 'srv-a' }, + github: { transport: 'stdio', command: 'npx', args: [], env: {} }, + }, + } as unknown as McpConfiguration); + + const result = await client.loadToolsWithFailures(); + + expect(result.tools).toHaveLength(1); + expect(result.failures).toEqual([ + expect.objectContaining({ server: 'slack', mcpServerId: 'srv-a', kind: 'auth' }), + ]); + }); + + it("classifies a connection error as 'connection', not 'auth'", async () => { + getToolsMock.mockRejectedValue(new Error('connect ECONNREFUSED 127.0.0.1:3000')); + + const result = await new McpClient(singleServer()).loadToolsWithFailures(); + + expect(result.failures[0].kind).toBe('connection'); + }); + + it('logs the aggregated failure summary and does not throw on a per-server failure', async () => { + const logger = jest.fn(); + getToolsMock.mockRejectedValue(new Error('401')); + + const result = await new McpClient(singleServer(), logger).loadToolsWithFailures(); + + expect(result.tools).toEqual([]); + expect(logger).toHaveBeenCalledWith( + 'Error', + expect.stringContaining('Failed to load tools from 1/1'), + ); + }); + + it('loadTools delegates and returns only the tools', async () => { + getToolsMock.mockResolvedValue([makeTool('t3')]); + + const tools = await new McpClient(singleServer()).loadTools(); + + expect(tools).toHaveLength(1); + }); +}); + +// authType is an executor-side routing hint; like id it must be stripped in the constructor so it +// never reaches MultiServerMCPClient. +describe('McpClient constructor — authType stripping', () => { + it('strips authType and id, passing only the transport config to MultiServerMCPClient', () => { + (MultiServerMCPClient as jest.Mock).mockClear(); + + // eslint-disable-next-line no-new + new McpClient({ + configs: { + slack: { type: 'http', url: 'https://example.com/mcp', id: 'srv-a', authType: 'oauth2' }, + }, + } as unknown as McpConfiguration); + + const passedConfig = (MultiServerMCPClient as jest.Mock).mock.calls[0][0].mcpServers.slack; + expect(passedConfig).toEqual({ type: 'http', url: 'https://example.com/mcp' }); + }); +}); diff --git a/packages/workflow-executor/CLAUDE.md b/packages/workflow-executor/CLAUDE.md index abf827aa6e..72d2557be0 100644 --- a/packages/workflow-executor/CLAUDE.md +++ b/packages/workflow-executor/CLAUDE.md @@ -24,11 +24,12 @@ Front ◀──▶ Orchestrator ◀──pull/push──▶ Executor (this p - `runner.ts` — main entry: `Runner` (start/stop/triggerPoll, HTTP wiring, graceful drain, auto-chain). - `executors/` — one per step type + infra: `base-step-executor.ts` (template method, idempotency hook, error→outcome), `record-step-executor.ts` (shared record-step base), `trigger-record-action-step-executor.ts` (the `trigger-action` step), `step-executor-factory.ts`, `agent-with-log.ts` + `activity-log.ts` (audit-logged agent wrapper for mutating ops). Step types (`StepType`): `condition`, `read-record`, `update-record`, `trigger-action`, `load-related-record`, `mcp`, `guidance`. -- `ports/` — IO interfaces (all external IO is injected): `agent-port.ts` (datasource: getRecord/updateRecord/getRelatedData/getSingleRelatedData/resolvePolymorphicType/executeAction/getActionFormInfo/getActionForm/probe), `workflow-port.ts` (orchestrator), `run-store.ts` (per-run state). +- `ports/` — IO interfaces (all external IO is injected): `agent-port.ts` (datasource: getRecord/updateRecord/getRelatedData/getSingleRelatedData/resolvePolymorphicType/executeAction/getActionFormInfo/getActionForm/probe), `workflow-port.ts` (orchestrator), `run-store.ts` (per-run state), `mcp-oauth-credentials-store.ts` (`McpOAuthCredentialsStore` — at-rest OAuth-MCP creds; Database + InMemory impls). - `adapters/` — port impls: `agent-client-agent-port.ts` (via `@forestadmin/agent-client`), `forest-server-workflow-port.ts` (HTTP via `forestadmin-client` `ServerUtils`). -- `stores/` — `InMemoryStore` (tests + the `--in-memory` CLI mode, not for prod), `DatabaseStore` (Sequelize + umzug, `forest` schema), `build-run-store.ts` factories. +- `stores/` — `InMemoryStore` (tests + the `--in-memory` CLI mode, not for prod), `DatabaseStore` (Sequelize + umzug, `forest` schema), `database-mcp-oauth-credentials-store.ts` / `in-memory-mcp-oauth-credentials-store.ts` (OAuth-MCP creds store, migration `002`), `schema-migrations.ts` (shared `forest`-schema + advisory-lock migration runner), `build-run-store.ts` factories. +- `crypto/credential-encryption.ts` — at-rest encryption: HKDF (`FOREST_EXECUTOR_ENCRYPTION_KEY`) + AES-256-GCM, lazy key, fail-closed. - `types/validated/` — zod schemas + inferred types for everything crossing a trust boundary (`step-definition.ts`, `step-outcome.ts`, `collection.ts`, `execution.ts`). Non-validated runtime types live in `types/*.ts`. -- `errors.ts` — error hierarchy (see below). `http/executor-http-server.ts` — optional Koa server for the front. +- `errors.ts` — error hierarchy (see below). `http/executor-http-server.ts` — optional Koa server for the front (`GET /runs/:runId`, `POST /runs/:runId/trigger`, `POST`/`DELETE /mcp-oauth-credentials`); `http/mcp-oauth-credentials.ts` — `.strict()` deposit-body zod + `buildMcpOAuthCredentialInput` mapper. ## Step types & execution modes diff --git a/packages/workflow-executor/README.md b/packages/workflow-executor/README.md index 0420dfc3dd..f6b597ada3 100644 --- a/packages/workflow-executor/README.md +++ b/packages/workflow-executor/README.md @@ -152,3 +152,7 @@ FOREST_AUTH_SECRET="your-auth-secret" \ AGENT_URL="https://your-agent-url" \ npx @forestadmin/workflow-executor --in-memory ``` + +### OAuth2-protected MCP + +To exercise the executor's OAuth2-protected MCP path end-to-end, point it at [mcp-oauth-test-server](https://github.com/hercemer42/mcp-oauth-test-server) — a standalone OAuth2 + MCP server that can simulate refresh-token revocation/rotation, consent denial, and upstream 403s. diff --git a/packages/workflow-executor/example/.env.example b/packages/workflow-executor/example/.env.example index fb96106528..354b34b67a 100644 --- a/packages/workflow-executor/example/.env.example +++ b/packages/workflow-executor/example/.env.example @@ -2,6 +2,9 @@ FOREST_ENV_SECRET= FOREST_AUTH_SECRET= +# At-rest encryption key for secrets the executor stores (openssl rand -hex 32; same value across all instances) +FOREST_EXECUTOR_ENCRYPTION_KEY= + # Your locally running Forest Admin agent AGENT_URL=http://localhost:3351 diff --git a/packages/workflow-executor/src/adapters/ai-client-adapter.ts b/packages/workflow-executor/src/adapters/ai-client-adapter.ts index fc5a98df13..43369caf16 100644 --- a/packages/workflow-executor/src/adapters/ai-client-adapter.ts +++ b/packages/workflow-executor/src/adapters/ai-client-adapter.ts @@ -1,5 +1,11 @@ import type { AiModelPort, GetModelOptions } from '../ports/ai-model-port'; -import type { AiConfiguration, BaseChatModel, RemoteTool, ToolConfig } from '@forestadmin/ai-proxy'; +import type { + AiConfiguration, + BaseChatModel, + McpServerLoadFailure, + RemoteTool, + ToolConfig, +} from '@forestadmin/ai-proxy'; import { AiClient } from '@forestadmin/ai-proxy'; @@ -26,6 +32,14 @@ export default class AiClientAdapter implements AiModelPort { return this.callPort('loadRemoteTools', () => this.aiClient.loadRemoteTools(configs)); } + loadRemoteToolsWithFailures( + configs: Record, + ): Promise<{ tools: RemoteTool[]; failures: McpServerLoadFailure[] }> { + return this.callPort('loadRemoteToolsWithFailures', () => + this.aiClient.loadRemoteToolsWithFailures(configs), + ); + } + closeConnections(): Promise { return this.callPort('closeConnections', () => this.aiClient.closeConnections()); } diff --git a/packages/workflow-executor/src/adapters/always-error-ai-model-port.ts b/packages/workflow-executor/src/adapters/always-error-ai-model-port.ts index 24a7426c5c..a51e5ef061 100644 --- a/packages/workflow-executor/src/adapters/always-error-ai-model-port.ts +++ b/packages/workflow-executor/src/adapters/always-error-ai-model-port.ts @@ -1,5 +1,10 @@ import type { AiModelPort } from '../ports/ai-model-port'; -import type { BaseChatModel, RemoteTool, ToolConfig } from '@forestadmin/ai-proxy'; +import type { + BaseChatModel, + McpServerLoadFailure, + RemoteTool, + ToolConfig, +} from '@forestadmin/ai-proxy'; import { AiModelPortError } from '../errors'; @@ -22,6 +27,13 @@ export default class AlwaysErrorAiModelPort implements AiModelPort { return Promise.resolve([]); } + loadRemoteToolsWithFailures( + // eslint-disable-next-line @typescript-eslint/no-unused-vars + _configs: Record, + ): Promise<{ tools: RemoteTool[]; failures: McpServerLoadFailure[] }> { + return Promise.resolve({ tools: [], failures: [] }); + } + closeConnections(): Promise { return Promise.resolve(); } diff --git a/packages/workflow-executor/src/adapters/server-ai-adapter.ts b/packages/workflow-executor/src/adapters/server-ai-adapter.ts index fae591047b..2d05307d30 100644 --- a/packages/workflow-executor/src/adapters/server-ai-adapter.ts +++ b/packages/workflow-executor/src/adapters/server-ai-adapter.ts @@ -1,5 +1,11 @@ import type { AiModelPort, GetModelOptions } from '../ports/ai-model-port'; -import type { AiConfiguration, BaseChatModel, RemoteTool, ToolConfig } from '@forestadmin/ai-proxy'; +import type { + AiConfiguration, + BaseChatModel, + McpServerLoadFailure, + RemoteTool, + ToolConfig, +} from '@forestadmin/ai-proxy'; import { AiClient } from '@forestadmin/ai-proxy'; @@ -38,6 +44,14 @@ export default class ServerAiAdapter implements AiModelPort { return this.callPort('loadRemoteTools', () => this.aiClient.loadRemoteTools(configs)); } + loadRemoteToolsWithFailures( + configs: Record, + ): Promise<{ tools: RemoteTool[]; failures: McpServerLoadFailure[] }> { + return this.callPort('loadRemoteToolsWithFailures', () => + this.aiClient.loadRemoteToolsWithFailures(configs), + ); + } + closeConnections(): Promise { return this.callPort('closeConnections', () => this.aiClient.closeConnections()); } diff --git a/packages/workflow-executor/src/adapters/step-outcome-to-update-step-mapper.ts b/packages/workflow-executor/src/adapters/step-outcome-to-update-step-mapper.ts index 70d72627e7..3965001a14 100644 --- a/packages/workflow-executor/src/adapters/step-outcome-to-update-step-mapper.ts +++ b/packages/workflow-executor/src/adapters/step-outcome-to-update-step-mapper.ts @@ -36,6 +36,10 @@ export default function toUpdateStepRequest( context.approvalRequest = outcome.approvalRequest; } + if (outcome.type === 'mcp' && outcome.awaitingInputReason !== undefined) { + context.awaitingInputReason = outcome.awaitingInputReason; + } + const attributes: ServerStepHistoryUpdate = { done: outcome.status !== 'awaiting-input', context, diff --git a/packages/workflow-executor/src/build-workflow-executor.ts b/packages/workflow-executor/src/build-workflow-executor.ts index 64f8e332fb..0a717c8caa 100644 --- a/packages/workflow-executor/src/build-workflow-executor.ts +++ b/packages/workflow-executor/src/build-workflow-executor.ts @@ -13,6 +13,9 @@ import createConsoleLogger from './adapters/console-logger'; import ForestServerWorkflowPort from './adapters/forest-server-workflow-port'; import ForestadminClientActivityLogPortFactory from './adapters/forestadmin-client-activity-log-port-factory'; import ServerAiAdapter from './adapters/server-ai-adapter'; +import CredentialEncryption, { + isExecutorEncryptionKeyConfigured, +} from './crypto/credential-encryption'; import { DEFAULT_AI_INVOKE_TIMEOUT_S, DEFAULT_FOREST_SERVER_URL, @@ -22,9 +25,13 @@ import { DEFAULT_STEP_TIMEOUT_S, } from './defaults'; import ExecutorHttpServer from './http/executor-http-server'; +import OAuthTokenService from './oauth/token-service'; +import RemoteToolFetcher from './remote-tool-fetcher'; import Runner from './runner'; import SchemaCache from './schema-cache'; +import DatabaseMcpOAuthCredentialsStore from './stores/database-mcp-oauth-credentials-store'; import DatabaseStore from './stores/database-store'; +import InMemoryMcpOAuthCredentialsStore from './stores/in-memory-mcp-oauth-credentials-store'; import InMemoryStore from './stores/in-memory-store'; const FORCE_EXIT_DELAY_S = 5; @@ -74,6 +81,15 @@ function buildCommonDependencies(options: ExecutorOptions) { const forestServerUrl = options.forestServerUrl ?? DEFAULT_FOREST_SERVER_URL; const logger = options.logger ?? createConsoleLogger(options.loggerLevel ?? DEFAULT_LOGGER_LEVEL); + // Lazy key by design (OAuth-less executors boot without it), but surface a security-relevant + // heads-up so a missing key isn't discovered only when a deposit 503s. + if (!isExecutorEncryptionKeyConfigured()) { + logger( + 'Warn', + 'FOREST_EXECUTOR_ENCRYPTION_KEY is not set — OAuth-protected MCP servers are unavailable (credential deposits return 503) until it is configured.', + ); + } + const workflowPort = new ForestServerWorkflowPort({ envSecret: options.envSecret, forestServerUrl, @@ -188,7 +204,15 @@ function createWorkflowExecutor( async start() { await runner.start(); - await server.start(); + + // server.start() runs the OAuth-store migration; if it fails, don't leave the runner polling + // in the background — stop it so start() is all-or-nothing. + try { + await server.start(); + } catch (err) { + await runner.stop().catch(() => {}); + throw err; + } // Only own the host's signals when explicitly allowed (the standalone CLI). When embedded, // the host process must keep control of SIGTERM/SIGINT and its own exit. @@ -213,9 +237,28 @@ function createWorkflowExecutor( export function buildInMemoryExecutor(options: ExecutorOptions): WorkflowExecutor { const deps = buildCommonDependencies(options); + const mcpOAuthCredentialsStore = new InMemoryMcpOAuthCredentialsStore(); + const credentialEncryption = new CredentialEncryption(); + // Shares the store + encryption with the deposit endpoint so runtime reads and writes (rotation) + // go through the same instance the HTTP server exposes. In-memory is dev-only: credentials live + // only for the process lifetime, but oauth2 steps work end-to-end just like the database executor. + const mcpOAuthTokenService = new OAuthTokenService({ + store: mcpOAuthCredentialsStore, + encryption: credentialEncryption, + logger: deps.logger, + }); + + const remoteToolFetcher = new RemoteToolFetcher( + deps.workflowPort, + deps.aiModelPort, + deps.logger, + mcpOAuthTokenService, + ); + const runner = new Runner({ ...deps, runStore: new InMemoryStore(), + mcpOAuthTokenService, }); const server = new ExecutorHttpServer({ @@ -224,6 +267,10 @@ export function buildInMemoryExecutor(options: ExecutorOptions): WorkflowExecuto authSecret: options.authSecret, workflowPort: deps.workflowPort, logger: deps.logger, + mcpOAuthCredentialsStore, + credentialEncryption, + remoteToolFetcher, + oauthTokenService: mcpOAuthTokenService, }); return createWorkflowExecutor(runner, server, deps.logger, options.manageProcessSignals ?? true); @@ -241,9 +288,30 @@ export function buildDatabaseExecutor(options: DatabaseExecutorOptions): Workflo if (mergedOptions.logging === undefined) mergedOptions.logging = false; const sequelize = uri ? new Sequelize(uri, mergedOptions) : new Sequelize(mergedOptions); + const mcpOAuthCredentialsStore = new DatabaseMcpOAuthCredentialsStore({ + sequelize, + schema: mergedOptions.schema, + }); + const credentialEncryption = new CredentialEncryption(); + // Shares the store + encryption with the deposit endpoint so runtime reads and writes (rotation) + // go through the same instance the HTTP server migrates on start. + const mcpOAuthTokenService = new OAuthTokenService({ + store: mcpOAuthCredentialsStore, + encryption: credentialEncryption, + logger: deps.logger, + }); + + const remoteToolFetcher = new RemoteToolFetcher( + deps.workflowPort, + deps.aiModelPort, + deps.logger, + mcpOAuthTokenService, + ); + const runner = new Runner({ ...deps, runStore: new DatabaseStore({ sequelize, schema: mergedOptions.schema }), + mcpOAuthTokenService, }); const server = new ExecutorHttpServer({ @@ -252,6 +320,10 @@ export function buildDatabaseExecutor(options: DatabaseExecutorOptions): Workflo authSecret: options.authSecret, workflowPort: deps.workflowPort, logger: deps.logger, + mcpOAuthCredentialsStore, + credentialEncryption, + remoteToolFetcher, + oauthTokenService: mcpOAuthTokenService, }); return createWorkflowExecutor(runner, server, deps.logger, options.manageProcessSignals ?? true); diff --git a/packages/workflow-executor/src/crypto/credential-encryption.ts b/packages/workflow-executor/src/crypto/credential-encryption.ts new file mode 100644 index 0000000000..ddaacc2956 --- /dev/null +++ b/packages/workflow-executor/src/crypto/credential-encryption.ts @@ -0,0 +1,89 @@ +import { createCipheriv, createDecipheriv, hkdfSync, randomFillSync } from 'crypto'; + +import { ExecutorEncryptionKeyMissingError } from '../errors'; + +const ENV_KEY = 'FOREST_EXECUTOR_ENCRYPTION_KEY'; +// Fixed context label bound into the HKDF derivation — domain-separates this key from any other +// use of the same secret. Changing it would make every existing row undecryptable. +const HKDF_INFO = 'forest-executor:mcp-oauth-credentials'; +const HKDF_DIGEST = 'sha256'; +const KEY_BYTES = 32; // AES-256 +const IV_BYTES = 12; // GCM standard nonce length +const AUTH_TAG_BYTES = 16; +const ALGORITHM = 'aes-256-gcm'; + +// Boot-time check (no derivation, no throw): is the at-rest key configured? Lets startup emit a +// warning that OAuth-protected MCP servers are unavailable until it is set. There is deliberately no +// default key — a hardcoded one would be public in this open-source package and defeat the at-rest +// encryption (a DB dump would be decryptable with the shipped key). +export function isExecutorEncryptionKeyConfigured(): boolean { + return Boolean(process.env[ENV_KEY]); +} + +export interface EncryptedValue { + // Packed layout: iv | authTag | ciphertext — stored as a single BLOB column. + ciphertext: Buffer; +} + +// Concatenate byte arrays without going through Buffer.concat — keeps everything in the concrete +// Uint8Array domain the Node crypto types expect. +function concatBytes(parts: Uint8Array[]): Uint8Array { + const total = parts.reduce((length, part) => length + part.length, 0); + const out = new Uint8Array(total); + let offset = 0; + + for (const part of parts) { + out.set(part, offset); + offset += part.length; + } + + return out; +} + +// At-rest encryption for secrets the executor stores. The HKDF key (from +// FOREST_EXECUTOR_ENCRYPTION_KEY) is read lazily — an executor that stores no such secrets boots +// without it — and fails closed: a missing key throws rather than persisting or returning an +// unprotected value. +export default class CredentialEncryption { + encrypt(plaintext: string): EncryptedValue { + const iv = randomFillSync(new Uint8Array(IV_BYTES)); + const cipher = createCipheriv(ALGORITHM, this.deriveKey(), iv); + const encrypted = concatBytes([ + new Uint8Array(cipher.update(plaintext, 'utf8')), + new Uint8Array(cipher.final()), + ]); + const authTag = new Uint8Array(cipher.getAuthTag()); + + return { + ciphertext: Buffer.from(concatBytes([iv, authTag, encrypted])), + }; + } + + decrypt(value: Buffer): string { + const bytes = new Uint8Array(value); + const iv = bytes.subarray(0, IV_BYTES); + const authTag = bytes.subarray(IV_BYTES, IV_BYTES + AUTH_TAG_BYTES); + const encrypted = bytes.subarray(IV_BYTES + AUTH_TAG_BYTES); + + const decipher = createDecipheriv(ALGORITHM, this.deriveKey(), iv); + decipher.setAuthTag(authTag); + + const decrypted = concatBytes([ + new Uint8Array(decipher.update(encrypted)), + new Uint8Array(decipher.final()), + ]); + + return Buffer.from(decrypted).toString('utf8'); + } + + private deriveKey(): Uint8Array { + const secret = process.env[ENV_KEY]; + + if (!secret) throw new ExecutorEncryptionKeyMissingError(); + + // Empty salt is intentional: the fixed HKDF_INFO label gives domain separation and the + // single high-entropy secret needs no salt. Wrap hkdfSync's ArrayBuffer as a concrete + // Uint8Array to satisfy CipherKey (Buffer's ArrayBufferLike backing does not). + return new Uint8Array(hkdfSync(HKDF_DIGEST, secret, new Uint8Array(0), HKDF_INFO, KEY_BYTES)); + } +} diff --git a/packages/workflow-executor/src/defaults.ts b/packages/workflow-executor/src/defaults.ts index 9a1221cfb2..3ad7cd7669 100644 --- a/packages/workflow-executor/src/defaults.ts +++ b/packages/workflow-executor/src/defaults.ts @@ -9,3 +9,6 @@ export const DEFAULT_STOP_TIMEOUT_S = 30; export const DEFAULT_MAX_CHAIN_DEPTH = 50; export const DEFAULT_SCHEMA_CACHE_TTL_S = 10 * 60; export const DEFAULT_LOGGER_LEVEL: LoggerLevel = 'Info'; +// Refresh an OAuth access token this many seconds before it actually expires, so a token never +// goes stale mid-request between the skew check and the downstream MCP call. +export const DEFAULT_OAUTH_EXPIRY_SKEW_S = 60; diff --git a/packages/workflow-executor/src/errors.ts b/packages/workflow-executor/src/errors.ts index 4f5307879c..7656c9f166 100644 --- a/packages/workflow-executor/src/errors.ts +++ b/packages/workflow-executor/src/errors.ts @@ -1,6 +1,7 @@ /* eslint-disable max-classes-per-file */ import type { MalformedRunInfo } from './ports/workflow-port'; import type { RecordId } from './types/validated/collection'; +import type { AwaitingInputReason } from './types/validated/step-outcome'; import type { z } from 'zod'; export function causeMessage(error: unknown): string | undefined { @@ -361,6 +362,67 @@ export class ConfigurationError extends Error { } } +// Carries the typed pause reason so the step executor can emit an awaiting-input outcome (and the +// orchestrator/front can prompt the user to reconnect) instead of a generic error. +export class OAuthReauthRequiredError extends WorkflowExecutorError { + readonly awaitingInputReason: AwaitingInputReason; + + constructor( + mcpServerId: string, + awaitingInputReason: AwaitingInputReason = 'needs-oauth-reauth', + ) { + super( + `OAuth re-authentication required for mcpServerId="${mcpServerId}"`, + 'This tool needs to be reconnected before it can run. Please re-authorize it and try again.', + ); + this.awaitingInputReason = awaitingInputReason; + } +} + +// Transient refresh failure (network error, 5xx, malformed response). Surfaces as a retryable step +// error rather than a re-auth pause — re-authenticating would not fix a temporarily unreachable endpoint. +export class OAuthRefreshError extends WorkflowExecutorError { + constructor(message: string, cause?: unknown) { + super( + `OAuth token refresh failed: ${message}`, + 'Could not refresh the connection to this tool. Please try again.', + ); + if (cause !== undefined) this.cause = cause; + } +} + +// The token endpoint rejected the refresh token (RFC 6749 invalid_grant). Internal control-flow +// signal: the token service catches it to drive the re-read + single-retry, then converts a genuine +// failure into OAuthReauthRequiredError. +export class OAuthInvalidGrantError extends WorkflowExecutorError { + constructor(detail?: string) { + super(`OAuth refresh token rejected${detail ? `: ${detail}` : ''}`); + } +} + +// The token endpoint is where the executor POSTs the refresh grant (with client credentials), so an +// unconstrained value is an SSRF vector. A rejected one fails closed (terminal) before any network +// call, rather than letting an authenticated caller aim the executor at an internal address. +export class InvalidTokenEndpointError extends WorkflowExecutorError { + constructor(reason: string) { + super( + `Invalid OAuth token endpoint: ${reason}`, + 'This tool is misconfigured (invalid token endpoint). Contact your administrator.', + ); + } +} + +// Boundary error — the deposit endpoint maps it to a typed HTTP response so the frontend can tell +// an operator to provision the key, not a generic or re-consent failure. +export class ExecutorEncryptionKeyMissingError extends Error { + readonly code = 'executor_encryption_key_missing'; + + constructor() { + super('FOREST_EXECUTOR_ENCRYPTION_KEY is not set'); + this.name = 'ExecutorEncryptionKeyMissingError'; + } +} + // Run lifecycle/access errors raised by the Runner. Each extends a domain category, so toHttpError // maps them by category (404/409/403) — no per-error HTTP binding. export class RunNotFoundError extends NotFoundError { @@ -433,7 +495,7 @@ export class SourceRecordMissingError extends WorkflowExecutorError { // Boundary error — surfaces from Runner.start() and is caught at the CLI/HTTP layer, not by step executors. export class AgentProbeError extends Error { - // Manual `cause` assignment: Error accepts it natively since Node 16.9 but our TS target is ES2020. + // Manual `cause` assignment: our ES2020 TS target doesn't type the native Error `cause` option. readonly cause?: unknown; constructor(message: string, options?: { cause?: unknown }) { diff --git a/packages/workflow-executor/src/executors/mcp-step-executor.ts b/packages/workflow-executor/src/executors/mcp-step-executor.ts index 2447c8a837..27c9b8f105 100644 --- a/packages/workflow-executor/src/executors/mcp-step-executor.ts +++ b/packages/workflow-executor/src/executors/mcp-step-executor.ts @@ -1,16 +1,22 @@ import type { ExecutionContext, StepExecutionResult } from '../types/execution-context'; import type { McpStepExecutionData, McpToolCall } from '../types/step-execution-data'; import type { McpStepDefinition } from '../types/validated/step-definition'; -import type { RecordStepStatus } from '../types/validated/step-outcome'; +import type { AwaitingInputReason, RecordStepStatus } from '../types/validated/step-outcome'; import type { RemoteTool } from '@forestadmin/ai-proxy'; -import { DynamicStructuredTool, HumanMessage, SystemMessage } from '@forestadmin/ai-proxy'; +import { + DynamicStructuredTool, + HumanMessage, + SystemMessage, + isMcpAuthError, +} from '@forestadmin/ai-proxy'; import { z } from 'zod'; import { McpToolInvocationError, McpToolNotFoundError, NoMcpToolsError, + OAuthReauthRequiredError, StepStateError, } from '../errors'; import BaseStepExecutor from './base-step-executor'; @@ -28,14 +34,18 @@ export default class McpStepExecutor extends BaseStepExecutor private readonly mcpServerName?: string; + private readonly reloadWithFreshAuth?: () => Promise; + constructor( context: ExecutionContext, remoteTools: readonly RemoteTool[], mcpServerName?: string, + reloadWithFreshAuth?: () => Promise, ) { super(context); this.remoteTools = remoteTools; this.mcpServerName = mcpServerName; + this.reloadWithFreshAuth = reloadWithFreshAuth; } protected override getExtraLogContext(): Record { @@ -48,6 +58,7 @@ export default class McpStepExecutor extends BaseStepExecutor protected buildOutcomeResult(outcome: { status: RecordStepStatus; error?: string; + awaitingInputReason?: AwaitingInputReason; }): StepExecutionResult { return { stepOutcome: { @@ -74,7 +85,41 @@ export default class McpStepExecutor extends BaseStepExecutor } protected async doExecute(): Promise { - // Branch A -- Re-entry after pending execution found in RunStore + try { + return await this.runStep(); + } catch (error) { + // An unrefreshable OAuth credential pauses the step for re-authentication rather than failing + // it. Clear the write-ahead marker so the resumed step is not rejected as interrupted. + if (error instanceof OAuthReauthRequiredError) { + await this.clearReauthPauseState(); + + return this.buildOutcomeResult({ + status: 'awaiting-input', + awaitingInputReason: error.awaitingInputReason, + }); + } + + throw error; + } + } + + // Keep a confirmation-flow record's approved pendingData (clear only the marker) so resume replays + // it; delete a pendingData-less record, which would otherwise mis-route resume into confirmation. + private async clearReauthPauseState(): Promise { + const existing = await this.findPendingExecution('mcp'); + if (!existing) return; + + if (existing.pendingData) { + await this.context.runStore.saveStepExecution(this.context.runId, { + ...existing, + idempotencyPhase: undefined, + }); + } else { + await this.context.runStore.deleteStepExecution(this.context.runId, this.context.stepIndex); + } + } + + private async runStep(): Promise { const pending = await this.patchAndReloadPendingData( this.context.incomingPendingData, ); @@ -85,7 +130,6 @@ export default class McpStepExecutor extends BaseStepExecutor ); } - // Branches B & C -- First call const tools = this.requireTools(); const { toolName, args } = await this.selectTool(tools); const selectedTool = tools.find(t => t.base.name === toolName); @@ -93,11 +137,9 @@ export default class McpStepExecutor extends BaseStepExecutor const target: McpToolCall = { name: toolName, sourceId: selectedTool.sourceId, input: args }; if (this.context.stepDefinition.executionType === StepExecutionMode.FullyAutomated) { - // Branch B -- direct execution return this.executeToolAndPersist(target); } - // Branch C -- Awaiting confirmation await this.context.runStore.saveStepExecution(this.context.runId, { type: 'mcp', stepIndex: this.context.stepIndex, @@ -124,13 +166,7 @@ export default class McpStepExecutor extends BaseStepExecutor recordId: this.context.baseRecordRef.recordId, }, { - operation: async () => { - try { - return await tool.base.invoke(target.input); - } catch (cause) { - throw new McpToolInvocationError(target.name, cause); - } - }, + operation: () => this.invokeWithReauthRetry(tool, target), beforeCall: () => this.context.runStore.saveStepExecution(this.context.runId, { ...existingExecution, @@ -195,6 +231,55 @@ export default class McpStepExecutor extends BaseStepExecutor return this.buildOutcomeResult({ status: 'success' }); } + // No-op for bearer/none steps (no reloadWithFreshAuth). For an OAuth2 step, a 401 on the call means + // the token was rejected after listing tools succeeded: force one refresh, rebuild the tool, retry + // once. A second 401 pauses the step for re-authentication. + private async invokeWithReauthRetry(tool: RemoteTool, target: McpToolCall): Promise { + try { + return await tool.base.invoke(target.input); + } catch (cause) { + if (!this.reloadWithFreshAuth || !isMcpAuthError(cause)) { + throw new McpToolInvocationError(target.name, cause); + } + + let refreshedTools: RemoteTool[]; + + try { + refreshedTools = await this.reloadWithFreshAuth(); + } catch (refreshError) { + // A non-auth refresh failure means nothing ran (the first call was a rejected 401), so clear + // the write-ahead marker to keep the step retryable; OAuthReauthRequiredError still pauses. + if (!(refreshError instanceof OAuthReauthRequiredError)) { + await this.clearReauthPauseState(); + } + + throw refreshError; + } + + const refreshedTool = refreshedTools.find( + t => t.base.name === target.name && t.sourceId === target.sourceId, + ); + + if (!refreshedTool) { + // The 401 first call never ran and the reload yielded no tool (empty on a connection + // failure), so clear the marker to keep the step retryable rather than wedged. + await this.clearReauthPauseState(); + + throw new McpToolNotFoundError(target.name); + } + + try { + return await refreshedTool.base.invoke(target.input); + } catch (retryCause) { + if (isMcpAuthError(retryCause)) { + throw new OAuthReauthRequiredError(this.context.stepDefinition.mcpServerId); + } + + throw new McpToolInvocationError(target.name, retryCause); + } + } + } + private async formatToolResult(tool: McpToolCall, toolResult: unknown): Promise { if (toolResult === null || toolResult === undefined) return null; diff --git a/packages/workflow-executor/src/executors/step-executor-factory.ts b/packages/workflow-executor/src/executors/step-executor-factory.ts index cc01e32dec..65e09ca151 100644 --- a/packages/workflow-executor/src/executors/step-executor-factory.ts +++ b/packages/workflow-executor/src/executors/step-executor-factory.ts @@ -23,6 +23,7 @@ import type { } from '../types/validated/step-definition'; import { + OAuthReauthRequiredError, StepStateError, WorkflowExecutorError, causeMessage, @@ -57,7 +58,7 @@ export default class StepExecutorFactory { step: AvailableStepExecution, contextConfig: StepContextConfig, activityLogPort: ActivityLogPort, - fetchRemoteTools: (mcpServerId: string) => Promise, + fetchRemoteTools: (mcpServerId: string, userId: number) => Promise, incomingPendingData?: unknown, forestServerToken?: string, ): Promise { @@ -90,11 +91,12 @@ export default class StepExecutorFactory { case StepType.Mcp: { const mcpContext = context as ExecutionContext; - const { tools, mcpServerName } = await fetchRemoteTools( + const { tools, mcpServerName, reloadWithFreshAuth } = await fetchRemoteTools( mcpContext.stepDefinition.mcpServerId, + step.user.id, ); - return new McpStepExecutor(mcpContext, tools, mcpServerName); + return new McpStepExecutor(mcpContext, tools, mcpServerName, reloadWithFreshAuth); } case StepType.Guidance: @@ -105,6 +107,29 @@ export default class StepExecutorFactory { ); } } catch (error) { + // Re-auth required while loading tools is an expected pause, not a failure: emit awaiting-input + // with the typed reason so the user can reconnect, rather than erroring the step. + if (error instanceof OAuthReauthRequiredError) { + contextConfig.logger('Info', 'MCP step paused for OAuth re-authentication', { + runId: step.runId, + stepId: step.stepId, + stepIndex: step.stepIndex, + awaitingInputReason: error.awaitingInputReason, + }); + + return { + execute: async (): Promise => ({ + stepOutcome: { + type: 'mcp', + stepId: step.stepId, + stepIndex: step.stepIndex, + status: 'awaiting-input', + awaitingInputReason: error.awaitingInputReason, + }, + }), + }; + } + contextConfig.logger('Error', 'Step execution failed unexpectedly', { runId: step.runId, stepId: step.stepId, diff --git a/packages/workflow-executor/src/http/executor-http-server.ts b/packages/workflow-executor/src/http/executor-http-server.ts index b164f86a4a..094bc4472e 100644 --- a/packages/workflow-executor/src/http/executor-http-server.ts +++ b/packages/workflow-executor/src/http/executor-http-server.ts @@ -1,5 +1,9 @@ +import type CredentialEncryption from '../crypto/credential-encryption'; +import type OAuthTokenService from '../oauth/token-service'; import type { Logger } from '../ports/logger-port'; +import type { McpOAuthCredentialsStore } from '../ports/mcp-oauth-credentials-store'; import type { WorkflowPort } from '../ports/workflow-port'; +import type RemoteToolFetcher from '../remote-tool-fetcher'; import type Runner from '../runner'; import type { Server } from 'http'; @@ -11,14 +15,24 @@ import koaJwt from 'koa-jwt'; import { type BearerClaims, BearerClaimsSchema } from './bearer-claims'; import { + BadRequestHttpError, ForbiddenHttpError, + NotFoundHttpError, ServiceUnavailableHttpError, UnauthorizedHttpError, toHttpError, } from './http-errors'; +import { + buildMcpOAuthCredentialInput, + depositCredentialsBodySchema, +} from './mcp-oauth-credentials'; import serializeStepForWire from './step-serializer'; import createConsoleLogger from '../adapters/console-logger'; -import { extractErrorMessage } from '../errors'; +import { + ExecutorEncryptionKeyMissingError, + OAuthReauthRequiredError, + extractErrorMessage, +} from '../errors'; // eslint-disable-next-line @typescript-eslint/no-var-requires, import/no-dynamic-require, global-require const { version } = require('../../package.json') as { version: string }; @@ -29,17 +43,25 @@ export interface ExecutorHttpServerOptions { authSecret: string; workflowPort: WorkflowPort; logger?: Logger; + mcpOAuthCredentialsStore: McpOAuthCredentialsStore; + credentialEncryption: CredentialEncryption; + remoteToolFetcher: RemoteToolFetcher; + // The runtime always provides this (build-workflow-executor); optional so tests that don't + // exercise credential deletion don't all have to construct one. + oauthTokenService?: OAuthTokenService; } export default class ExecutorHttpServer { private readonly app: Koa; private readonly options: ExecutorHttpServerOptions; private readonly logger: Logger; + private readonly mcpOAuthCredentialsStore: McpOAuthCredentialsStore; private server: Server | null = null; constructor(options: ExecutorHttpServerOptions) { this.options = options; this.logger = options.logger ?? createConsoleLogger(); + this.mcpOAuthCredentialsStore = options.mcpOAuthCredentialsStore; this.app = new Koa(); this.app.use(async (ctx, next) => { @@ -142,11 +164,30 @@ export default class ExecutorHttpServer { ); router.post('/runs/:runId/trigger', this.handleTrigger.bind(this)); + const { + mcpOAuthCredentialsStore: credentialsStore, + credentialEncryption, + remoteToolFetcher, + } = this.options; + + // Design-time tool listing for the MCP-server details page: resolve the caller's vault + // credential, refresh, and list the oauth2 server's tools — no workflow run involved. + router.get('/list-mcp-tools', ctx => this.handleListMcpTools(ctx, remoteToolFetcher)); + + router.post('/mcp-oauth-credentials', ctx => + this.handleDepositCredentials(ctx, credentialsStore, credentialEncryption), + ); + router.delete('/mcp-oauth-credentials/:mcpServerId', ctx => + this.handleDeleteCredentials(ctx, credentialsStore), + ); + this.app.use(router.routes()); this.app.use(router.allowedMethods()); } async start(): Promise { + await this.mcpOAuthCredentialsStore.init(this.logger); + return new Promise((resolve, reject) => { this.server = http.createServer(this.app.callback()); this.server.once('error', reject); @@ -218,4 +259,110 @@ export default class ExecutorHttpServer { ctx.status = 200; ctx.body = { triggered: true }; } + + private async handleDepositCredentials( + ctx: Koa.Context, + store: McpOAuthCredentialsStore, + encryption: CredentialEncryption, + ): Promise { + const userId = (ctx.state.user as BearerClaims).id; + const parsed = depositCredentialsBodySchema.safeParse(ctx.request.body ?? {}); + + if (!parsed.success) { + const details = parsed.error.issues + .map(issue => `${issue.path.join('.') || 'body'}: ${issue.message}`) + .join('; '); + + throw new BadRequestHttpError(`Invalid request body — ${details}`); + } + + try { + await store.upsert(buildMcpOAuthCredentialInput({ body: parsed.data, userId, encryption })); + } catch (err) { + // The frontend must tell this missing-key config error apart from a generic failure (to route + // the user to an admin rather than retry), so it returns a typed { code } the middleware won't. + if (err instanceof ExecutorEncryptionKeyMissingError) { + ctx.status = 503; + ctx.body = { code: err.code }; + + return; + } + + throw err; + } + + // Evict any cached access token so a reconnect/re-deposit takes effect immediately, instead of + // serving the token minted from the previous credential until it expires (mirrors delete). + this.options.oauthTokenService?.evict(userId, parsed.data.mcpServerId); + ctx.status = 200; + ctx.body = { stored: true }; + } + + private async handleDeleteCredentials( + ctx: Koa.Context, + store: McpOAuthCredentialsStore, + ): Promise { + const userId = (ctx.state.user as BearerClaims).id; + + await store.delete(userId, ctx.params.mcpServerId); + // Evict any in-process cached access token so the disconnect is immediate, not deferred until + // the cached token expires (the row is gone, but the runtime would otherwise still serve it). + this.options.oauthTokenService?.evict(userId, ctx.params.mcpServerId); + ctx.status = 204; + } + + // Design-time tool listing: resolve the caller's vault credential for the target oauth2 server, + // refresh + inject the Bearer, and return the server's tool definitions. user_id comes from the + // validated JWT, so a caller only ever lists with their own stored credential. A missing/dead + // credential surfaces as a typed needs-oauth-reauth response (not a generic error or empty list) + // so the details page can prompt a reconnect — mirroring the run path's awaiting-input reason. + private async handleListMcpTools(ctx: Koa.Context, fetcher: RemoteToolFetcher): Promise { + const userId = (ctx.state.user as BearerClaims).id; + const { mcpServerId } = ctx.query; + + if (typeof mcpServerId !== 'string' || mcpServerId.length === 0) { + throw new BadRequestHttpError('Missing required query parameter "mcpServerId"'); + } + + try { + const { tools, mcpServerName, loadFailed } = await fetcher.fetch(mcpServerId, userId); + + // A dead or misrouted server otherwise renders as an (indistinguishable) empty tool list, so + // surface it — mirroring the run path, which errors rather than executing against no tools. + if (!mcpServerName) { + throw new NotFoundHttpError(`No MCP server is configured for id "${mcpServerId}"`); + } + + if (loadFailed) { + throw new ServiceUnavailableHttpError( + 'The MCP server could not be reached to list its tools', + ); + } + + ctx.status = 200; + ctx.body = { + tools: tools.map(tool => ({ name: tool.base.name, description: tool.base.description })), + }; + } catch (err) { + // Set the body directly so the error middleware (which would map this to a generic 400 and + // drop the typed reason) doesn't touch it — the frontend needs the reason to prompt reconnect. + if (err instanceof OAuthReauthRequiredError) { + ctx.status = 409; + ctx.body = { awaitingInputReason: err.awaitingInputReason, mcpServerId }; + + return; + } + + // Key-missing is an operator misconfig, not re-consent-resolvable: return the same typed 503 + // the deposit endpoint uses so the details page shows the admin message, not a generic error. + if (err instanceof ExecutorEncryptionKeyMissingError) { + ctx.status = 503; + ctx.body = { code: err.code }; + + return; + } + + throw err; + } + } } diff --git a/packages/workflow-executor/src/http/mcp-oauth-credentials.ts b/packages/workflow-executor/src/http/mcp-oauth-credentials.ts new file mode 100644 index 0000000000..ae31812d89 --- /dev/null +++ b/packages/workflow-executor/src/http/mcp-oauth-credentials.ts @@ -0,0 +1,85 @@ +import type CredentialEncryption from '../crypto/credential-encryption'; +import type { McpOAuthCredentialInput } from '../ports/mcp-oauth-credentials-store'; + +import { z } from 'zod'; + +import assertSafeTokenEndpoint from '../oauth/token-endpoint-url'; + +// String lengths mirror the DB column limits, so oversized values are rejected here at the boundary +// rather than at insert. .strict() matters for security: it blocks a body-supplied user id, since +// identity comes only from the JWT. +export const depositCredentialsBodySchema = z + .object({ + mcpServerId: z.string().min(1).max(255), + refreshToken: z.string().min(1), + clientId: z.string().min(1).max(255).optional(), + clientSecret: z.string().min(1).optional(), + clientSecretExpiresAt: z + .string() + .refine(value => !Number.isNaN(Date.parse(value)), { message: 'must be a parseable date' }) + .optional(), + tokenEndpoint: z + .string() + .min(1) + .max(2048) + .superRefine((value, ctx) => { + // The executor POSTs the refresh grant here, so reject SSRF-prone endpoints at deposit + // rather than discovering them at refresh time. + try { + assertSafeTokenEndpoint(value); + } catch (error) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: error instanceof Error ? error.message : 'invalid token endpoint', + }); + } + }), + // Restrict to the methods the refresh grant implements, so an unsupported/typo'd one is rejected + // here rather than silently mis-authenticating at refresh time ('none' = public client). + tokenEndpointAuthMethod: z + .enum(['client_secret_basic', 'client_secret_post', 'none']) + .optional(), + scopes: z.string().max(2048).optional(), + }) + .strict() + .superRefine((body, ctx) => { + // RFC 6749 requires client_id alongside client_secret; a secret with no id can never + // authenticate, so reject it here instead of persisting an unusable credential. + if (body.clientSecret && !body.clientId) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + path: ['clientId'], + message: 'clientId is required when clientSecret is provided', + }); + } + }); + +export type DepositCredentialsBody = z.infer; + +// Translates a validated deposit body into the at-rest record: encrypts the refresh token (and +// client secret when present) and maps optional fields to their nullable columns. encrypt() throws +// ExecutorEncryptionKeyMissingError when the key is unset; the caller maps that to a 503. +export function buildMcpOAuthCredentialInput({ + body, + userId, + encryption, +}: { + body: DepositCredentialsBody; + userId: number; + encryption: CredentialEncryption; +}): McpOAuthCredentialInput { + const refreshToken = encryption.encrypt(body.refreshToken); + const clientSecret = body.clientSecret ? encryption.encrypt(body.clientSecret) : null; + + return { + userId, + mcpServerId: body.mcpServerId, + refreshTokenEnc: refreshToken.ciphertext, + clientId: body.clientId ?? null, + clientSecretEnc: clientSecret?.ciphertext ?? null, + clientSecretExpiresAt: body.clientSecretExpiresAt ? new Date(body.clientSecretExpiresAt) : null, + tokenEndpoint: body.tokenEndpoint, + tokenEndpointAuthMethod: body.tokenEndpointAuthMethod ?? null, + scopes: body.scopes ?? null, + }; +} diff --git a/packages/workflow-executor/src/index.ts b/packages/workflow-executor/src/index.ts index 7fa349a083..e18ed25d2b 100644 --- a/packages/workflow-executor/src/index.ts +++ b/packages/workflow-executor/src/index.ts @@ -100,6 +100,7 @@ export { AiModelPortError, AgentProbeError, ConfigurationError, + ExecutorEncryptionKeyMissingError, InvalidPreRecordedArgsError, UnsupportedStepTypeError, UnsupportedActionFormError, @@ -126,6 +127,16 @@ export { default as SchemaResolver } from './schema-resolver'; export { default as InMemoryStore } from './stores/in-memory-store'; export { default as DatabaseStore } from './stores/database-store'; export type { DatabaseStoreOptions } from './stores/database-store'; +export { default as DatabaseMcpOAuthCredentialsStore } from './stores/database-mcp-oauth-credentials-store'; +export type { DatabaseMcpOAuthCredentialsStoreOptions } from './stores/database-mcp-oauth-credentials-store'; +export { default as InMemoryMcpOAuthCredentialsStore } from './stores/in-memory-mcp-oauth-credentials-store'; +export type { + McpOAuthCredentialsStore, + McpOAuthCredentialInput, + StoredMcpOAuthCredential, +} from './ports/mcp-oauth-credentials-store'; +export { default as CredentialEncryption } from './crypto/credential-encryption'; +export type { EncryptedValue } from './crypto/credential-encryption'; export { buildDatabaseRunStore, buildInMemoryRunStore } from './stores/build-run-store'; export { buildInMemoryExecutor, buildDatabaseExecutor } from './build-workflow-executor'; export { runCli } from './cli-core'; diff --git a/packages/workflow-executor/src/oauth/keyed-mutex.ts b/packages/workflow-executor/src/oauth/keyed-mutex.ts new file mode 100644 index 0000000000..76dc6769c0 --- /dev/null +++ b/packages/workflow-executor/src/oauth/keyed-mutex.ts @@ -0,0 +1,24 @@ +// Serializes async work per key: callers sharing a key run one at a time (FIFO), while different +// keys proceed in parallel. Used to collapse concurrent token refreshes for the same (user, server) +// into a single in-flight refresh within one process. +export default class KeyedMutex { + private readonly tails = new Map>(); + + async runExclusive(key: string, task: () => Promise): Promise { + const previous = this.tails.get(key) ?? Promise.resolve(); + const result = previous.then(() => task()); + // Store a non-rejecting tail so the next caller chains regardless of this task's outcome. + const tail = result.then( + () => undefined, + () => undefined, + ); + this.tails.set(key, tail); + + try { + return await result; + } finally { + // Drop the entry once this was the last queued task for the key, so the map can't grow. + if (this.tails.get(key) === tail) this.tails.delete(key); + } + } +} diff --git a/packages/workflow-executor/src/oauth/refresh-grant.ts b/packages/workflow-executor/src/oauth/refresh-grant.ts new file mode 100644 index 0000000000..6cef868d78 --- /dev/null +++ b/packages/workflow-executor/src/oauth/refresh-grant.ts @@ -0,0 +1,151 @@ +import { OAuthInvalidGrantError, OAuthRefreshError } from '../errors'; +import assertSafeTokenEndpoint from './token-endpoint-url'; + +export interface RefreshGrantParams { + tokenEndpoint: string; + refreshToken: string; + clientId?: string | null; + clientSecret?: string | null; + tokenEndpointAuthMethod?: string | null; + scopes?: string | null; +} + +export interface RefreshGrantResult { + accessToken: string; + expiresInS?: number; + // Present only when the authorization server rotates the refresh token. + refreshToken?: string; +} + +interface TokenEndpointResponse { + access_token?: unknown; + expires_in?: unknown; + refresh_token?: unknown; + error?: unknown; + error_description?: unknown; +} + +// RFC 6749 §2.3.1 form-url-encodes the client id/secret before base64 (space → '+', unlike +// encodeURIComponent's '%20'); identical for the base64url/hex credentials providers actually issue. +function formUrlEncode(value: string): string { + const params = new URLSearchParams(); + params.set('v', value); + + return params.toString().slice('v='.length); +} + +// Some providers return expires_in as a numeric string; coerce it so the token is still cached +// (an unusable value leaves it uncached, forcing a full refresh — and rotation — on every call). +function coerceExpiresInS(value: unknown): number | undefined { + if (typeof value === 'number') return value; + if (typeof value === 'string' && /^\d+$/.test(value)) return Number(value); + + return undefined; +} + +function buildRequest(params: RefreshGrantParams): { + headers: Record; + body: URLSearchParams; +} { + const headers: Record = { + 'content-type': 'application/x-www-form-urlencoded', + accept: 'application/json', + }; + const body = new URLSearchParams({ + grant_type: 'refresh_token', + refresh_token: params.refreshToken, + }); + + if (params.scopes) body.set('scope', params.scopes); + + const { clientId, clientSecret, tokenEndpointAuthMethod } = params; + + if (clientSecret) { + if (tokenEndpointAuthMethod === 'client_secret_post') { + if (clientId) body.set('client_id', clientId); + body.set('client_secret', clientSecret); + } else { + const credentials = `${formUrlEncode(clientId ?? '')}:${formUrlEncode(clientSecret)}`; + headers.authorization = `Basic ${Buffer.from(credentials).toString('base64')}`; + } + } else if (clientId) { + body.set('client_id', clientId); + } + + return { headers, body }; +} + +// Runs the OAuth2 refresh-token grant (RFC 6749 §6) against the stored token endpoint. Distinguishes +// a non-retryable invalid_grant (revoked / rotated) from a transient failure so the caller can decide +// between forcing re-auth and retrying. +export default async function refreshAccessToken( + params: RefreshGrantParams, +): Promise { + // Defense in depth: deposit validation rejects SSRF-prone endpoints, but a row predating that + // validation could still carry one — re-check before the outbound POST. Throws (terminal) rather + // than reaching the network. + assertSafeTokenEndpoint(params.tokenEndpoint); + + const { headers, body } = buildRequest(params); + + let response: Awaited>; + + try { + // Never follow redirects: a 3xx to an internal host would re-send the grant body (refresh + // token, plus the client secret in client_secret_post) there and parse its reply as a token — + // bypassing the endpoint validation above. Any redirect is treated as a failure. + response = await fetch(params.tokenEndpoint, { + method: 'POST', + headers, + body, + redirect: 'manual', + }); + } catch (cause) { + throw new OAuthRefreshError('the token endpoint could not be reached', cause); + } + + if (response.type === 'opaqueredirect' || (response.status >= 300 && response.status < 400)) { + throw new OAuthRefreshError('the token endpoint attempted a redirect'); + } + + let payload: TokenEndpointResponse = {}; + + try { + const parsed: unknown = await response.json(); + // Ignore a non-object body (e.g. literal null), keeping the {} default so the status checks + // below still surface a typed OAuthRefreshError. + if (parsed && typeof parsed === 'object') payload = parsed as TokenEndpointResponse; + } catch { + // Non-JSON body — handled by the status checks below. + } + + if (!response.ok) { + // invalid_client / unauthorized_client mean the client credential itself is dead (e.g. an expired + // secret) — non-retryable like invalid_grant, so route to re-consent instead of a doomed retry. + if ( + payload.error === 'invalid_grant' || + payload.error === 'invalid_client' || + payload.error === 'unauthorized_client' + ) { + throw new OAuthInvalidGrantError( + typeof payload.error_description === 'string' ? payload.error_description : undefined, + ); + } + + throw new OAuthRefreshError( + typeof payload.error === 'string' + ? payload.error + : `the token endpoint returned ${response.status}`, + ); + } + + if (typeof payload.access_token !== 'string' || !payload.access_token) { + throw new OAuthRefreshError('the token endpoint response had no access_token'); + } + + return { + accessToken: payload.access_token, + expiresInS: coerceExpiresInS(payload.expires_in), + refreshToken: typeof payload.refresh_token === 'string' ? payload.refresh_token : undefined, + }; +} diff --git a/packages/workflow-executor/src/oauth/token-endpoint-url.ts b/packages/workflow-executor/src/oauth/token-endpoint-url.ts new file mode 100644 index 0000000000..4e413f5a82 --- /dev/null +++ b/packages/workflow-executor/src/oauth/token-endpoint-url.ts @@ -0,0 +1,117 @@ +import net from 'net'; + +import { InvalidTokenEndpointError } from '../errors'; + +// The token endpoint comes from the deposit body, and the executor POSTs the refresh grant (with +// client credentials) to it — so an unconstrained value is an SSRF vector. The guard rejects hosts +// that are never a legitimate endpoint but are prime SSRF targets (loopback, link-local / +// cloud-metadata) and requires TLS, while deliberately allowing private (RFC1918) hosts: a private +// OAuth provider reached over the internal network is a supported deployment, so blocking it would +// break a real use case. http and loopback are relaxed only when NODE_ENV explicitly opts in +// (development/test) — see isRelaxedEnv. Only IP literals are inspected — resolving a hostname here +// would be racy (it can resolve to a different address at fetch time), not safer. + +// Fail closed: the strict rules (https + no loopback) apply everywhere UNLESS NODE_ENV explicitly +// opts into the local-dev relaxation. An unset or misspelled NODE_ENV stays strict, so a +// misconfigured deploy can't silently allow loopback targets or cleartext http. +function isRelaxedEnv(): boolean { + return process.env.NODE_ENV === 'development' || process.env.NODE_ENV === 'test'; +} + +function unbracket(hostname: string): string { + return hostname.startsWith('[') && hostname.endsWith(']') ? hostname.slice(1, -1) : hostname; +} + +type HostClass = { loopback: boolean; linkLocal: boolean; unspecified: boolean }; + +function classifyIpv4(host: string): HostClass { + const [a, b] = host.split('.').map(Number); + + return { + loopback: a === 127, // 127.0.0.0/8 + linkLocal: a === 169 && b === 254, // 169.254.0.0/16 + unspecified: a === 0, // 0.0.0.0/8 — "this host"; connects to loopback on Linux + }; +} + +// Returns the embedded IPv4 of an IPv4-mapped IPv6 address (otherwise null). The WHATWG URL parser +// normalizes `::ffff:127.0.0.1` to the hex form `::ffff:7f00:1`, which would otherwise dodge the +// IPv4 loopback/link-local checks and still reach the executor's loopback / metadata interface. +function embeddedIpv4(host: string): string | null { + const tail = host.toLowerCase().match(/^::ffff:(.+)$/)?.[1]; + if (!tail) return null; + + if (net.isIPv4(tail)) return tail; + + const hex = tail.match(/^([0-9a-f]{1,4}):([0-9a-f]{1,4})$/); + if (!hex) return null; + + const hi = parseInt(hex[1], 16); // high 16 bits → first two octets + const lo = parseInt(hex[2], 16); // low 16 bits → last two octets + + return `${Math.floor(hi / 256)}.${hi % 256}.${Math.floor(lo / 256)}.${lo % 256}`; +} + +function classify(host: string): HostClass { + const kind = net.isIP(host); + + if (kind === 4) return classifyIpv4(host); + + if (kind === 6) { + const mapped = embeddedIpv4(host); + if (mapped) return classifyIpv4(mapped); + + const normalized = host.toLowerCase().split('%')[0]; // drop any zone id + const firstHextet = normalized.split(':')[0]; + + return { + loopback: normalized === '::1' || normalized === '0:0:0:0:0:0:0:1', + linkLocal: /^fe[89ab]/.test(firstHextet), // fe80::/10 + unspecified: normalized === '::' || normalized === '0:0:0:0:0:0:0:0', + }; + } + + // Not an IP literal. Reject the reserved loopback names — bare `localhost`, the `.localhost` TLD + // (RFC 6761), and their trailing-dot FQDN forms — since they resolve to loopback. Other hostnames + // are left to the scheme rule; their DNS is not resolved here (a hostname pointed at an internal + // IP is the filtering-agent's job, not this string check). + return { + loopback: /^localhost\.?$|\.localhost\.?$/.test(host.toLowerCase()), + linkLocal: false, + unspecified: false, + }; +} + +export default function assertSafeTokenEndpoint(raw: string): void { + let url: URL; + + try { + url = new URL(raw); + } catch { + throw new InvalidTokenEndpointError('it is not a valid absolute URL'); + } + + if (url.protocol !== 'https:' && url.protocol !== 'http:') { + throw new InvalidTokenEndpointError('it must use http or https'); + } + + const { loopback, linkLocal, unspecified } = classify(unbracket(url.hostname)); + + if (unspecified) { + throw new InvalidTokenEndpointError('it points at the unspecified address (0.0.0.0 / ::)'); + } + + if (linkLocal) { + throw new InvalidTokenEndpointError('it points at a link-local / metadata address'); + } + + if (!isRelaxedEnv()) { + if (url.protocol !== 'https:') { + throw new InvalidTokenEndpointError('it must use https'); + } + + if (loopback) { + throw new InvalidTokenEndpointError('it points at the executor host (loopback)'); + } + } +} diff --git a/packages/workflow-executor/src/oauth/token-service.ts b/packages/workflow-executor/src/oauth/token-service.ts new file mode 100644 index 0000000000..ab6628f74a --- /dev/null +++ b/packages/workflow-executor/src/oauth/token-service.ts @@ -0,0 +1,230 @@ +import type { RefreshGrantParams, RefreshGrantResult } from './refresh-grant'; +import type CredentialEncryption from '../crypto/credential-encryption'; +import type { Logger } from '../ports/logger-port'; +import type { + McpOAuthCredentialsStore, + StoredMcpOAuthCredential, +} from '../ports/mcp-oauth-credentials-store'; + +import { DEFAULT_OAUTH_EXPIRY_SKEW_S } from '../defaults'; +import { + ExecutorEncryptionKeyMissingError, + OAuthInvalidGrantError, + OAuthReauthRequiredError, +} from '../errors'; +import KeyedMutex from './keyed-mutex'; +import defaultRefreshAccessToken from './refresh-grant'; + +interface CachedToken { + accessToken: string; + expiresAtMs: number; +} + +export interface OAuthTokenServiceOptions { + store: McpOAuthCredentialsStore; + encryption: CredentialEncryption; + logger?: Logger; + expirySkewS?: number; + refreshAccessToken?: (params: RefreshGrantParams) => Promise; + now?: () => number; +} + +// Acquires an MCP OAuth access token for a (user, server): serves a cached token until it nears +// expiry, otherwise runs the refresh-token grant under a per-key mutex (one in-flight refresh per +// user+server in this process) and recovers from a concurrent refresh-token rotation via a single +// re-read + retry. A missing credential or a genuinely rejected refresh raises +// OAuthReauthRequiredError so the step pauses for re-authentication. +export default class OAuthTokenService { + private readonly store: McpOAuthCredentialsStore; + private readonly encryption: CredentialEncryption; + private readonly logger?: Logger; + private readonly expirySkewMs: number; + private readonly refreshAccessToken: (params: RefreshGrantParams) => Promise; + private readonly now: () => number; + private readonly cache = new Map(); + private readonly mutex = new KeyedMutex(); + // Bumped on evict() so a refresh already in flight for a key can detect a disconnect that landed + // mid-refresh and skip repopulating the cache for the now-deleted credential. + private readonly evictionEpoch = new Map(); + + constructor(options: OAuthTokenServiceOptions) { + this.store = options.store; + this.encryption = options.encryption; + this.logger = options.logger; + this.expirySkewMs = (options.expirySkewS ?? DEFAULT_OAUTH_EXPIRY_SKEW_S) * 1000; + this.refreshAccessToken = options.refreshAccessToken ?? defaultRefreshAccessToken; + this.now = options.now ?? Date.now; + } + + async getAccessToken( + userId: number, + mcpServerId: string, + options?: { forceRefresh?: boolean }, + ): Promise { + const key = `${userId}:${mcpServerId}`; + const forceRefresh = options?.forceRefresh ?? false; + + if (!forceRefresh) { + const cached = this.readCache(key); + if (cached) return cached; + } + + return this.mutex.runExclusive(key, async () => { + // Re-check inside the lock: a concurrent caller may have just refreshed for this key. + if (!forceRefresh) { + const cached = this.readCache(key); + if (cached) return cached; + } + + return this.refreshAndCache(userId, mcpServerId, key); + }); + } + + // Drop the cached access token for a (user, server). Called on credential delete so a disconnect + // takes effect immediately, instead of the executor serving the cached token until it expires. + evict(userId: number, mcpServerId: string): void { + const key = `${userId}:${mcpServerId}`; + this.cache.delete(key); + this.evictionEpoch.set(key, (this.evictionEpoch.get(key) ?? 0) + 1); + } + + private readCache(key: string): string | undefined { + const entry = this.cache.get(key); + if (!entry) return undefined; + + if (this.now() >= entry.expiresAtMs - this.expirySkewMs) { + // Drop expired entries on read so the cache can't grow without bound across many keys. + this.cache.delete(key); + + return undefined; + } + + return entry.accessToken; + } + + private async refreshAndCache(userId: number, mcpServerId: string, key: string): Promise { + const epochAtStart = this.evictionEpoch.get(key) ?? 0; + const credential = await this.store.get(userId, mcpServerId); + if (!credential) throw new OAuthReauthRequiredError(mcpServerId); + + const { result, credential: grantedCredential } = await this.runGrantWithRotationRetry( + credential, + userId, + mcpServerId, + ); + + const evictedDuringRefresh = (this.evictionEpoch.get(key) ?? 0) !== epochAtStart; + + // Skip caching if the credential was disconnected mid-refresh (would serve a token for a deleted + // credential) or the grant carried no expiry (never servable); the in-flight caller still gets it. + if (!evictedDuringRefresh && result.expiresInS !== undefined) { + this.cache.set(key, { + accessToken: result.accessToken, + expiresAtMs: this.now() + result.expiresInS * 1000, + }); + } else { + this.cache.delete(key); + } + + if (result.refreshToken) { + await this.persistRotatedRefreshToken(grantedCredential, result.refreshToken); + } + + return result.accessToken; + } + + // On invalid_grant a peer instance likely rotated the refresh token out from under us: re-read the + // row and retry once with the current token; a second invalid_grant (or an unchanged token) is a + // genuine revocation and forces re-auth. Returns the credential whose token produced the grant so + // the caller writes the rotated token back onto that (current) row, not the stale one. + private async runGrantWithRotationRetry( + credential: StoredMcpOAuthCredential, + userId: number, + mcpServerId: string, + ): Promise<{ result: RefreshGrantResult; credential: StoredMcpOAuthCredential }> { + try { + const result = await this.refreshAccessToken(this.toGrantParams(credential)); + + return { result, credential }; + } catch (error) { + if (!(error instanceof OAuthInvalidGrantError)) throw error; + + // A peer rotated the refresh token iff the stored value changed since we read it. + const latest = await this.store.get(userId, mcpServerId); + const wasRotated = + latest !== null && + latest.refreshTokenEnc.toString('base64') !== credential.refreshTokenEnc.toString('base64'); + + if (!latest || !wasRotated) { + throw new OAuthReauthRequiredError(mcpServerId); + } + + try { + const result = await this.refreshAccessToken(this.toGrantParams(latest)); + + return { result, credential: latest }; + } catch (retryError) { + if (retryError instanceof OAuthInvalidGrantError) { + throw new OAuthReauthRequiredError(mcpServerId); + } + + throw retryError; + } + } + } + + // Decrypt happens here. A decrypt failure with the key PRESENT (auth-tag mismatch — the row was + // encrypted under a since-rotated/hard-swapped key, or is corrupt) is recoverable: re-consent + // re-deposits under the current key, so surface it as needs-oauth-reauth. A missing key + // (ExecutorEncryptionKeyMissingError) is an operator misconfig, not re-consent-resolvable, so it + // propagates as a terminal error (a re-deposit would just 503 at the deposit endpoint). + private toGrantParams(credential: StoredMcpOAuthCredential): RefreshGrantParams { + try { + return { + tokenEndpoint: credential.tokenEndpoint, + refreshToken: this.encryption.decrypt(credential.refreshTokenEnc), + clientId: credential.clientId, + clientSecret: credential.clientSecretEnc + ? this.encryption.decrypt(credential.clientSecretEnc) + : null, + tokenEndpointAuthMethod: credential.tokenEndpointAuthMethod, + scopes: credential.scopes, + }; + } catch (error) { + if (error instanceof ExecutorEncryptionKeyMissingError) throw error; + + throw new OAuthReauthRequiredError(credential.mcpServerId); + } + } + + private async persistRotatedRefreshToken( + credential: StoredMcpOAuthCredential, + refreshToken: string, + ): Promise { + try { + const encrypted = this.encryption.encrypt(refreshToken); + + // Id-scoped so a disconnect (or disconnect + re-authorize) after the grant read leaves an + // absent or different row — the rotated-token write-back then can't resurrect or clobber it. + await this.store.updateIfPresent(credential.id, { + userId: credential.userId, + mcpServerId: credential.mcpServerId, + refreshTokenEnc: encrypted.ciphertext, + clientId: credential.clientId, + clientSecretEnc: credential.clientSecretEnc, + clientSecretExpiresAt: credential.clientSecretExpiresAt, + tokenEndpoint: credential.tokenEndpoint, + tokenEndpointAuthMethod: credential.tokenEndpointAuthMethod, + scopes: credential.scopes, + }); + } catch (error) { + // A failed write-back is not fatal: the access token just obtained is valid for this call. + // The rotated refresh token is lost, so a later refresh forces re-authentication — the safe + // fallback, and better than failing the current operation over a transient write error. + this.logger?.('Error', 'Failed to persist rotated MCP OAuth refresh token', { + mcpServerId: credential.mcpServerId, + error: error instanceof Error ? error.message : String(error), + }); + } + } +} diff --git a/packages/workflow-executor/src/ports/ai-model-port.ts b/packages/workflow-executor/src/ports/ai-model-port.ts index 5241a7a32c..30d0f18abf 100644 --- a/packages/workflow-executor/src/ports/ai-model-port.ts +++ b/packages/workflow-executor/src/ports/ai-model-port.ts @@ -1,4 +1,9 @@ -import type { BaseChatModel, RemoteTool, ToolConfig } from '@forestadmin/ai-proxy'; +import type { + BaseChatModel, + McpServerLoadFailure, + RemoteTool, + ToolConfig, +} from '@forestadmin/ai-proxy'; export interface GetModelOptions { aiConfigName?: string; @@ -8,5 +13,10 @@ export interface GetModelOptions { export interface AiModelPort { getModel(options?: GetModelOptions): BaseChatModel; loadRemoteTools(configs: Record): Promise; + // Loads tools and exposes per-server failures classified by cause (auth vs connection), so the + // OAuth path can tell a revoked token from an unreachable server. Default consumers use loadRemoteTools. + loadRemoteToolsWithFailures( + configs: Record, + ): Promise<{ tools: RemoteTool[]; failures: McpServerLoadFailure[] }>; closeConnections(): Promise; } diff --git a/packages/workflow-executor/src/ports/mcp-oauth-credentials-store.ts b/packages/workflow-executor/src/ports/mcp-oauth-credentials-store.ts new file mode 100644 index 0000000000..ef745b9c1c --- /dev/null +++ b/packages/workflow-executor/src/ports/mcp-oauth-credentials-store.ts @@ -0,0 +1,31 @@ +import type { Logger } from './logger-port'; + +export interface McpOAuthCredentialInput { + userId: number; + mcpServerId: string; + refreshTokenEnc: Buffer; + clientId?: string | null; + clientSecretEnc?: Buffer | null; + clientSecretExpiresAt?: Date | null; + tokenEndpoint: string; + tokenEndpointAuthMethod?: string | null; + scopes?: string | null; +} + +export interface StoredMcpOAuthCredential extends McpOAuthCredentialInput { + id: number; +} + +// Persists OAuth MCP credentials, one row per (userId, mcpServerId). Backed by a Sequelize store +// (real executors) or an in-memory one (--in-memory / dev). Holds opaque encrypted bytes — +// encryption happens upstream in buildMcpOAuthCredentialInput — so implementations do no crypto. +export interface McpOAuthCredentialsStore { + init(logger?: Logger): Promise; + close(logger?: Logger): Promise; + get(userId: number, mcpServerId: string): Promise; + upsert(credential: McpOAuthCredentialInput): Promise; + // Update-only by the row id the caller read: never inserts, and no-ops if that row is gone or was + // re-created with a new id — so a stale write-back can't resurrect or clobber a credential. + updateIfPresent(id: number, credential: McpOAuthCredentialInput): Promise; + delete(userId: number, mcpServerId: string): Promise; +} diff --git a/packages/workflow-executor/src/ports/run-store.ts b/packages/workflow-executor/src/ports/run-store.ts index de5a2da1ab..7e30facbfb 100644 --- a/packages/workflow-executor/src/ports/run-store.ts +++ b/packages/workflow-executor/src/ports/run-store.ts @@ -6,4 +6,5 @@ export interface RunStore { close(logger?: Logger): Promise; getStepExecutions(runId: string): Promise; saveStepExecution(runId: string, stepExecution: StepExecutionData): Promise; + deleteStepExecution(runId: string, stepIndex: number): Promise; } diff --git a/packages/workflow-executor/src/remote-tool-fetcher.ts b/packages/workflow-executor/src/remote-tool-fetcher.ts index 1b79c44c24..21c854a690 100644 --- a/packages/workflow-executor/src/remote-tool-fetcher.ts +++ b/packages/workflow-executor/src/remote-tool-fetcher.ts @@ -1,8 +1,15 @@ +import type OAuthTokenService from './oauth/token-service'; import type { AiModelPort } from './ports/ai-model-port'; import type { Logger } from './ports/logger-port'; import type { WorkflowPort } from './ports/workflow-port'; import type { RemoteTool, ToolConfig } from '@forestadmin/ai-proxy'; +import { injectOauthTokens } from '@forestadmin/ai-proxy'; + +import { OAuthReauthRequiredError } from './errors'; + +const OAUTH2_AUTH_TYPE = 'oauth2'; + // Match by config.id, not by Record key: server names can collide across configs. export function scopeConfigsToServer( configs: Record, @@ -11,23 +18,47 @@ export function scopeConfigsToServer( return Object.fromEntries(Object.entries(configs).filter(([, cfg]) => cfg.id === mcpServerId)); } +function readAuthType(config: ToolConfig | undefined): string | undefined { + return (config as { authType?: string } | undefined)?.authType; +} + export interface FetchRemoteToolsResult { tools: RemoteTool[]; mcpServerName?: string; + // True when the server was found but its tools failed to load, so a caller can tell a genuinely + // empty server from an unreachable one. + loadFailed?: boolean; + // Present only for OAuth2 servers: re-mints the token (forced refresh) and reloads the tools, so + // the executor can retry once after an upstream 401 on a tool call. Throws OAuthReauthRequiredError + // when the credential can no longer be refreshed. + reloadWithFreshAuth?: () => Promise; } export default class RemoteToolFetcher { private readonly workflowPort: WorkflowPort; private readonly aiModelPort: AiModelPort; private readonly logger: Logger; + private readonly oauthTokenService: OAuthTokenService; + + constructor( + workflowPort: WorkflowPort, + aiModelPort: AiModelPort, + logger: Logger, + oauthTokenService: OAuthTokenService, + ) { + // Fail loudly here instead of a cryptic undefined dereference at the first OAuth fetch — + // Runner/RunnerConfig are public API, so a JS caller can bypass the compile-time check. + if (!oauthTokenService) { + throw new Error('RemoteToolFetcher requires an OAuth token service (mcpOAuthTokenService)'); + } - constructor(workflowPort: WorkflowPort, aiModelPort: AiModelPort, logger: Logger) { this.workflowPort = workflowPort; this.aiModelPort = aiModelPort; this.logger = logger; + this.oauthTokenService = oauthTokenService; } - async fetch(mcpServerId: string): Promise { + async fetch(mcpServerId: string, userId: number): Promise { const configs = await this.workflowPort.getMcpServerConfigs(); const scoped = scopeConfigsToServer(configs, mcpServerId); const [mcpServerName] = Object.keys(scoped); @@ -36,11 +67,68 @@ export default class RemoteToolFetcher { if (Object.keys(scoped).length === 0) return { tools: [], mcpServerName }; + if (readAuthType(scoped[mcpServerName]) === OAUTH2_AUTH_TYPE) { + return this.fetchOAuthTools(scoped, mcpServerName, mcpServerId, userId); + } + const tools = await this.aiModelPort.loadRemoteTools(scoped); + const loadFailed = this.errorOnPartialLoadFailure(scoped, tools, mcpServerId, mcpServerName); + + return { tools, mcpServerName, loadFailed }; + } - this.errorOnPartialLoadFailure(scoped, tools, mcpServerId, mcpServerName); + // OAuth2 path: acquire a per-user access token, inject it as a Bearer header, then list tools. + // Tool listing is the first authenticated call, so an auth failure here forces one token refresh + // and a single retry; a still-failing load is a genuine re-auth case. + private async fetchOAuthTools( + scoped: Record, + mcpServerName: string, + mcpServerId: string, + userId: number, + ): Promise { + const tokenService = this.oauthTokenService; + + const attemptLoad = async ( + forceRefresh: boolean, + ): Promise<{ tools: RemoteTool[]; hasAuthFailure: boolean }> => { + const token = await tokenService.getAccessToken(userId, mcpServerId, { forceRefresh }); + const bearer = `Bearer ${token}`; + // All scoped configs share this mcpServerId, so inject the token for every one — not just the + // first key — or extra same-id servers load unauthenticated and report spurious auth failures. + const injected = + injectOauthTokens({ + configs: scoped, + tokensByMcpServerName: Object.fromEntries( + Object.keys(scoped).map(name => [name, bearer]), + ), + }) ?? scoped; + const { tools, failures } = await this.aiModelPort.loadRemoteToolsWithFailures(injected); + + return { tools, hasAuthFailure: failures.some(failure => failure.kind === 'auth') }; + }; + + const reloadWithFreshAuth = async (): Promise => { + const attempt = await attemptLoad(true); + if (attempt.hasAuthFailure) throw new OAuthReauthRequiredError(mcpServerId); + this.errorOnPartialLoadFailure(scoped, attempt.tools, mcpServerId, mcpServerName); + + return attempt.tools; + }; + + const initial = await attemptLoad(false); + + if (initial.hasAuthFailure) { + return { tools: await reloadWithFreshAuth(), mcpServerName, reloadWithFreshAuth }; + } + + const loadFailed = this.errorOnPartialLoadFailure( + scoped, + initial.tools, + mcpServerId, + mcpServerName, + ); - return { tools, mcpServerName }; + return { tools: initial.tools, mcpServerName, reloadWithFreshAuth, loadFailed }; } // Distinguish "no configs at all" (deployment misconfig) from "configs exist but none match" @@ -75,18 +163,20 @@ export default class RemoteToolFetcher { tools: RemoteTool[], mcpServerId: string, mcpServerName: string | undefined, - ): void { + ): boolean { const loadedMcpServerIds = new Set(tools.map(t => t.mcpServerId)); const failedConfigNames = Object.entries(scoped) .filter(([, cfg]) => !loadedMcpServerIds.has(cfg.id)) .map(([name]) => name); - if (failedConfigNames.length === 0) return; + if (failedConfigNames.length === 0) return false; this.logger('Error', 'MCP servers failed to load tools', { requestedMcpServerId: mcpServerId, mcpServerName, failedConfigNames, }); + + return true; } } diff --git a/packages/workflow-executor/src/runner.ts b/packages/workflow-executor/src/runner.ts index 5a0ba81966..02be945160 100644 --- a/packages/workflow-executor/src/runner.ts +++ b/packages/workflow-executor/src/runner.ts @@ -1,4 +1,5 @@ import type { StepContextConfig } from './executors/step-executor-factory'; +import type OAuthTokenService from './oauth/token-service'; import type { ActivityLogPortFactory } from './ports/activity-log-port'; import type { AgentPort } from './ports/agent-port'; import type { AiModelPort } from './ports/ai-model-port'; @@ -50,6 +51,9 @@ export interface RunnerConfig { // Max number of ADDITIONAL steps auto-chained via /update-step response before yielding to the // next poll cycle (counted after the initial step). 0 disables chaining entirely. Default 50. maxChainDepth?: number; + // Per-user OAuth access-token service for oauth2 MCP steps. Wired by both the in-memory and + // database executors, sharing the credential store the HTTP deposit endpoint writes to. + mcpOAuthTokenService: OAuthTokenService; } // eslint-disable-next-line @typescript-eslint/no-var-requires, import/no-dynamic-require, global-require @@ -70,6 +74,7 @@ export default class Runner { config.workflowPort, config.aiModelPort, this.logger, + config.mcpOAuthTokenService, ); } @@ -313,7 +318,7 @@ export default class Runner { currentStep, this.contextConfig, this.config.activityLogPortFactory.forRun(currentToken), - mcpServerId => this.remoteToolFetcher.fetch(mcpServerId), + (mcpServerId, userId) => this.remoteToolFetcher.fetch(mcpServerId, userId), currentIncomingData, currentToken, ); diff --git a/packages/workflow-executor/src/stores/database-mcp-oauth-credentials-store.ts b/packages/workflow-executor/src/stores/database-mcp-oauth-credentials-store.ts new file mode 100644 index 0000000000..613aa2be52 --- /dev/null +++ b/packages/workflow-executor/src/stores/database-mcp-oauth-credentials-store.ts @@ -0,0 +1,256 @@ +import type { Logger } from '../ports/logger-port'; +import type { + McpOAuthCredentialInput, + McpOAuthCredentialsStore, + StoredMcpOAuthCredential, +} from '../ports/mcp-oauth-credentials-store'; +import type { QueryInterface, Sequelize } from 'sequelize'; + +import { DataTypes } from 'sequelize'; +import { SequelizeStorage, Umzug } from 'umzug'; + +import { extractErrorMessage } from '../errors'; +import { + resolveSchema, + runMigrations, + tableId as toTableId, + tableReference as toTableReference, +} from './schema-migrations'; + +const TABLE_NAME = 'ai_mcp_oauth_credentials'; + +export interface DatabaseMcpOAuthCredentialsStoreOptions { + sequelize: Sequelize; + schema?: string; +} + +interface CredentialRow { + id: number; + user_id: number; + mcp_server_id: string; + refresh_token_enc: Buffer; + client_id: string | null; + client_secret_enc: Buffer | null; + client_secret_expires_at: string | Date | null; + token_endpoint: string; + token_endpoint_auth_method: string | null; + scopes: string | null; +} + +export default class DatabaseMcpOAuthCredentialsStore implements McpOAuthCredentialsStore { + private readonly sequelize: Sequelize; + + private readonly configuredSchema?: string; + + constructor(options: DatabaseMcpOAuthCredentialsStoreOptions) { + this.sequelize = options.sequelize; + this.configuredSchema = options.schema; + } + + private get schema(): string | undefined { + return resolveSchema(this.sequelize, this.configuredSchema); + } + + private get tableReference(): string { + return toTableReference(this.schema, TABLE_NAME); + } + + async init(logger?: Logger): Promise { + const { schema } = this; + const tableId = toTableId(schema, TABLE_NAME); + + const umzug = new Umzug({ + migrations: [ + { + name: '002_create_mcp_oauth_credentials', + up: async ({ context }: { context: QueryInterface }) => { + // Atomic (table + index) and idempotent so a half-applied or already-applied run can't + // crash-loop boot. + await context.sequelize.transaction(async transaction => { + if (await context.tableExists(tableId, { transaction })) return; + + await context.createTable( + tableId, + { + id: { type: DataTypes.INTEGER, primaryKey: true, autoIncrement: true }, + userId: { type: DataTypes.INTEGER, allowNull: false, field: 'user_id' }, + mcpServerId: { + type: DataTypes.STRING(255), + allowNull: false, + field: 'mcp_server_id', + }, + refreshTokenEnc: { + type: DataTypes.BLOB, + allowNull: false, + field: 'refresh_token_enc', + }, + clientId: { type: DataTypes.STRING(255), allowNull: true, field: 'client_id' }, + clientSecretEnc: { + type: DataTypes.BLOB, + allowNull: true, + field: 'client_secret_enc', + }, + clientSecretExpiresAt: { + type: DataTypes.DATE, + allowNull: true, + field: 'client_secret_expires_at', + }, + tokenEndpoint: { + type: DataTypes.STRING(2048), + allowNull: false, + field: 'token_endpoint', + }, + tokenEndpointAuthMethod: { + type: DataTypes.STRING(64), + allowNull: true, + field: 'token_endpoint_auth_method', + }, + scopes: { type: DataTypes.STRING(2048), allowNull: true, field: 'scopes' }, + createdAt: { + type: DataTypes.DATE, + allowNull: false, + defaultValue: DataTypes.NOW, + field: 'created_at', + }, + updatedAt: { + type: DataTypes.DATE, + allowNull: false, + defaultValue: DataTypes.NOW, + field: 'updated_at', + }, + }, + { transaction }, + ); + + await context.addIndex(tableId, ['user_id', 'mcp_server_id'], { + unique: true, + name: 'idx_user_id_mcp_server_id', + transaction, + }); + }); + }, + down: async ({ context }: { context: QueryInterface }) => { + await context.dropTable(tableId); + }, + }, + ], + context: this.sequelize.getQueryInterface(), + storage: new SequelizeStorage({ + sequelize: this.sequelize, + ...(schema ? { schema } : {}), + }), + logger: undefined, + }); + + await runMigrations({ + sequelize: this.sequelize, + umzug, + schema, + logger, + failMessage: 'MCP OAuth credentials migration failed', + }); + } + + async get(userId: number, mcpServerId: string): Promise { + const [rows] = await this.sequelize.query( + `SELECT * FROM ${this.tableReference} WHERE user_id = :userId AND mcp_server_id = :mcpServerId`, + { replacements: { userId, mcpServerId } }, + ); + + const row = (rows as CredentialRow[])[0]; + + return row ? DatabaseMcpOAuthCredentialsStore.toCredential(row) : null; + } + + async upsert(credential: McpOAuthCredentialInput): Promise { + await this.sequelize.transaction(async transaction => { + const now = new Date(); + const replacements = { + userId: credential.userId, + mcpServerId: credential.mcpServerId, + refreshTokenEnc: credential.refreshTokenEnc, + clientId: credential.clientId ?? null, + clientSecretEnc: credential.clientSecretEnc ?? null, + clientSecretExpiresAt: credential.clientSecretExpiresAt ?? null, + tokenEndpoint: credential.tokenEndpoint, + tokenEndpointAuthMethod: credential.tokenEndpointAuthMethod ?? null, + scopes: credential.scopes ?? null, + now, + }; + + // Delete + insert in transaction: dialect-agnostic upsert (avoids ON CONFLICT / ON DUPLICATE). + await this.sequelize.query( + `DELETE FROM ${this.tableReference} WHERE user_id = :userId AND mcp_server_id = :mcpServerId`, + { replacements, transaction }, + ); + await this.sequelize.query( + `INSERT INTO ${this.tableReference} ` + + '(user_id, mcp_server_id, refresh_token_enc, client_id, client_secret_enc, ' + + 'client_secret_expires_at, token_endpoint, token_endpoint_auth_method, scopes, ' + + 'created_at, updated_at) VALUES ' + + '(:userId, :mcpServerId, :refreshTokenEnc, :clientId, :clientSecretEnc, ' + + ':clientSecretExpiresAt, :tokenEndpoint, :tokenEndpointAuthMethod, :scopes, ' + + ':now, :now)', + { replacements, transaction }, + ); + }); + } + + async updateIfPresent(id: number, credential: McpOAuthCredentialInput): Promise { + // Single atomic UPDATE … WHERE id — affects zero rows if that row was deleted or re-created. + await this.sequelize.query( + `UPDATE ${this.tableReference} SET ` + + 'refresh_token_enc = :refreshTokenEnc, client_id = :clientId, ' + + 'client_secret_enc = :clientSecretEnc, client_secret_expires_at = :clientSecretExpiresAt, ' + + 'token_endpoint = :tokenEndpoint, token_endpoint_auth_method = :tokenEndpointAuthMethod, ' + + 'scopes = :scopes, updated_at = :now ' + + 'WHERE id = :id', + { + replacements: { + id, + refreshTokenEnc: credential.refreshTokenEnc, + clientId: credential.clientId ?? null, + clientSecretEnc: credential.clientSecretEnc ?? null, + clientSecretExpiresAt: credential.clientSecretExpiresAt ?? null, + tokenEndpoint: credential.tokenEndpoint, + tokenEndpointAuthMethod: credential.tokenEndpointAuthMethod ?? null, + scopes: credential.scopes ?? null, + now: new Date(), + }, + }, + ); + } + + async delete(userId: number, mcpServerId: string): Promise { + await this.sequelize.query( + `DELETE FROM ${this.tableReference} WHERE user_id = :userId AND mcp_server_id = :mcpServerId`, + { replacements: { userId, mcpServerId } }, + ); + } + + async close(logger?: Logger): Promise { + try { + await this.sequelize.close(); + } catch (error) { + logger?.('Error', 'Failed to close database connection', { + error: extractErrorMessage(error), + }); + } + } + + private static toCredential(row: CredentialRow): StoredMcpOAuthCredential { + return { + id: Number(row.id), + userId: Number(row.user_id), + mcpServerId: row.mcp_server_id, + refreshTokenEnc: row.refresh_token_enc, + clientId: row.client_id ?? null, + clientSecretEnc: row.client_secret_enc ?? null, + clientSecretExpiresAt: + row.client_secret_expires_at == null ? null : new Date(row.client_secret_expires_at), + tokenEndpoint: row.token_endpoint, + tokenEndpointAuthMethod: row.token_endpoint_auth_method ?? null, + scopes: row.scopes ?? null, + }; + } +} diff --git a/packages/workflow-executor/src/stores/database-store.ts b/packages/workflow-executor/src/stores/database-store.ts index 7d099dca4e..8a90ef234c 100644 --- a/packages/workflow-executor/src/stores/database-store.ts +++ b/packages/workflow-executor/src/stores/database-store.ts @@ -1,18 +1,20 @@ import type { Logger } from '../ports/logger-port'; import type { RunStore } from '../ports/run-store'; import type { StepExecutionData } from '../types/step-execution-data'; -import type { QueryInterface, Sequelize, Transaction } from 'sequelize'; +import type { QueryInterface, Sequelize } from 'sequelize'; import { DataTypes } from 'sequelize'; import { SequelizeStorage, Umzug } from 'umzug'; import { RunStorePortError, WorkflowExecutorError, extractErrorMessage } from '../errors'; +import { + resolveSchema, + runMigrations, + tableId as toTableId, + tableReference as toTableReference, +} from './schema-migrations'; const TABLE_NAME = 'workflow_step_executions'; -const DEFAULT_SCHEMA = 'forest'; - -// Must stay constant across releases, or an old and a new deploy could migrate concurrently. -const MIGRATION_ADVISORY_LOCK_KEY = 6_438_071_259_157; export interface DatabaseStoreOptions { sequelize: Sequelize; @@ -30,17 +32,15 @@ export default class DatabaseStore implements RunStore { } private get schema(): string | undefined { - if (this.sequelize.getDialect() === 'sqlite') return undefined; - - return this.configuredSchema || DEFAULT_SCHEMA; + return resolveSchema(this.sequelize, this.configuredSchema); } private get tableId(): string | { tableName: string; schema: string } { - return this.schema ? { tableName: TABLE_NAME, schema: this.schema } : TABLE_NAME; + return toTableId(this.schema, TABLE_NAME); } private get tableReference(): string { - return this.schema ? `"${this.schema}"."${TABLE_NAME}"` : `"${TABLE_NAME}"`; + return toTableReference(this.schema, TABLE_NAME); } async init(logger?: Logger): Promise { @@ -115,60 +115,15 @@ export default class DatabaseStore implements RunStore { logger: undefined, }); - return this.callPort('init', async () => { - try { - if (this.sequelize.getDialect() !== 'postgres') { - await umzug.up(); - - return; - } - - // The migration lock holds one pool connection while umzug opens a second. - const { pool } = this.sequelize.connectionManager as unknown as { - pool?: { maxSize?: number }; - }; - const poolMax = pool?.maxSize ?? 1; - - if (poolMax < 2) { - throw new Error( - 'workflow-executor requires pool.max >= 2 on Postgres: the migration lock holds one connection while migrations run on another', - ); - } - - // Schema in its own committed transaction so umzug (on other connections) sees it. - if (schema) { - await this.withMigrationLock(transaction => - this.sequelize.query(`CREATE SCHEMA IF NOT EXISTS "${schema}"`, { transaction }), - ); - } - - await this.withMigrationLock(() => umzug.up()); - } catch (error) { - logger?.('Error', 'Database migration failed', { - error: extractErrorMessage(error), - }); - throw error; - } - }); - } - - // Serializes booting instances via a transaction-scoped advisory lock: auto-releases at commit - // and is pooler-safe (RDS Proxy / PgBouncer), unlike a session lock which would leak there. - private async withMigrationLock( - run: (transaction: Transaction) => Promise, - ): Promise { - await this.sequelize.transaction(async transaction => { - // Stop a client idle-in-transaction timeout from killing this idle txn mid-migration, - // which would drop the lock. - await this.sequelize.query('SET LOCAL idle_in_transaction_session_timeout = 0', { - transaction, - }); - await this.sequelize.query('SELECT pg_advisory_xact_lock($1)', { - bind: [MIGRATION_ADVISORY_LOCK_KEY], - transaction, - }); - await run(transaction); - }); + return this.callPort('init', () => + runMigrations({ + sequelize: this.sequelize, + umzug, + schema, + logger, + failMessage: 'Database migration failed', + }), + ); } async getStepExecutions(runId: string): Promise { @@ -204,6 +159,15 @@ export default class DatabaseStore implements RunStore { }); } + async deleteStepExecution(runId: string, stepIndex: number): Promise { + return this.callPort('deleteStepExecution', async () => { + await this.sequelize.query( + `DELETE FROM ${this.tableReference} WHERE run_id = :runId AND step_index = :stepIndex`, + { replacements: { runId, stepIndex } }, + ); + }); + } + async close(logger?: Logger): Promise { return this.callPort('close', async () => { try { diff --git a/packages/workflow-executor/src/stores/in-memory-mcp-oauth-credentials-store.ts b/packages/workflow-executor/src/stores/in-memory-mcp-oauth-credentials-store.ts new file mode 100644 index 0000000000..f1e8ff9607 --- /dev/null +++ b/packages/workflow-executor/src/stores/in-memory-mcp-oauth-credentials-store.ts @@ -0,0 +1,51 @@ +import type { + McpOAuthCredentialInput, + McpOAuthCredentialsStore, + StoredMcpOAuthCredential, +} from '../ports/mcp-oauth-credentials-store'; + +// In-memory MCP OAuth credentials store for --in-memory / dev. Same throwaway semantics as +// InMemoryStore (the run store): state is lost on restart. It holds the already-encrypted Buffers +// produced by buildMcpOAuthCredentialInput — there is no crypto here, just a Map keyed by +// (userId, mcpServerId), mirroring the DB store's one-row-per-key contract. +export default class InMemoryMcpOAuthCredentialsStore implements McpOAuthCredentialsStore { + private readonly data = new Map(); + private nextId = 1; + + async init(): Promise { + // No-op: nothing to migrate for an in-memory store. + } + + async close(): Promise { + // No-op: nothing to close. + } + + async get(userId: number, mcpServerId: string): Promise { + return this.data.get(InMemoryMcpOAuthCredentialsStore.key(userId, mcpServerId)) ?? null; + } + + async upsert(credential: McpOAuthCredentialInput): Promise { + // Overwrite in place — one row per (userId, mcpServerId). A fresh id each time mirrors the DB + // store's delete-then-insert; nothing relies on id stability. + const key = InMemoryMcpOAuthCredentialsStore.key(credential.userId, credential.mcpServerId); + this.data.set(key, { ...credential, id: this.nextId }); + this.nextId += 1; + } + + async updateIfPresent(id: number, credential: McpOAuthCredentialInput): Promise { + const key = InMemoryMcpOAuthCredentialsStore.key(credential.userId, credential.mcpServerId); + const existing = this.data.get(key); + // Only update the exact row the caller read; a re-created row has a new id, so skip it. + if (!existing || existing.id !== id) return; + + this.data.set(key, { ...credential, id }); + } + + async delete(userId: number, mcpServerId: string): Promise { + this.data.delete(InMemoryMcpOAuthCredentialsStore.key(userId, mcpServerId)); + } + + private static key(userId: number, mcpServerId: string): string { + return `${userId}:${mcpServerId}`; + } +} diff --git a/packages/workflow-executor/src/stores/in-memory-store.ts b/packages/workflow-executor/src/stores/in-memory-store.ts index 90117d70aa..6a20242155 100644 --- a/packages/workflow-executor/src/stores/in-memory-store.ts +++ b/packages/workflow-executor/src/stores/in-memory-store.ts @@ -41,6 +41,12 @@ export default class InMemoryStore implements RunStore { }); } + async deleteStepExecution(runId: string, stepIndex: number): Promise { + return this.callPort('deleteStepExecution', async () => { + this.data.get(runId)?.delete(stepIndex); + }); + } + private async callPort(operation: string, fn: () => Promise): Promise { try { return await fn(); diff --git a/packages/workflow-executor/src/stores/schema-migrations.ts b/packages/workflow-executor/src/stores/schema-migrations.ts new file mode 100644 index 0000000000..7389483c1c --- /dev/null +++ b/packages/workflow-executor/src/stores/schema-migrations.ts @@ -0,0 +1,99 @@ +import type { Logger } from '../ports/logger-port'; +import type { Sequelize, Transaction } from 'sequelize'; + +import { extractErrorMessage } from '../errors'; + +// The schema both executor tables live under, so a shared customer database stays safe: the executor +// never touches `public`, and its umzug `SequelizeMeta` registries can't collide with the +// agent/server's own. Skipped on SQLite (test suite), which has no schema support. +export const DEFAULT_SCHEMA = 'forest'; + +// Must stay constant across releases, or an old and a new deploy could migrate concurrently. Shared +// by every executor migration runner so all of them serialize against one another on cold start. +const MIGRATION_ADVISORY_LOCK_KEY = 6_438_071_259_157; + +export function resolveSchema(sequelize: Sequelize, configured?: string): string | undefined { + if (sequelize.getDialect() === 'sqlite') return undefined; + + return configured || DEFAULT_SCHEMA; +} + +export function tableId( + schema: string | undefined, + tableName: string, +): string | { tableName: string; schema: string } { + return schema ? { tableName, schema } : tableName; +} + +export function tableReference(schema: string | undefined, tableName: string): string { + return schema ? `"${schema}"."${tableName}"` : `"${tableName}"`; +} + +// Serializes booting instances via a transaction-scoped advisory lock: auto-releases at commit and +// is pooler-safe (RDS Proxy / PgBouncer), unlike a session lock which would leak there. +async function withMigrationLock( + sequelize: Sequelize, + run: (transaction: Transaction) => Promise, +): Promise { + await sequelize.transaction(async transaction => { + // Stop a client idle-in-transaction timeout from killing this idle txn mid-migration, which + // would drop the lock. + await sequelize.query('SET LOCAL idle_in_transaction_session_timeout = 0', { transaction }); + await sequelize.query('SELECT pg_advisory_xact_lock($1)', { + bind: [MIGRATION_ADVISORY_LOCK_KEY], + transaction, + }); + await run(transaction); + }); +} + +// Runs an umzug migration set under the executor's shared-database safety rules: on Postgres the +// `forest` schema is created and migrations run behind the advisory lock (one writer across booting +// replicas); on other dialects (SQLite) it just runs. Logs `failMessage` and rethrows on failure. +export async function runMigrations({ + sequelize, + umzug, + schema, + logger, + failMessage, +}: { + sequelize: Sequelize; + umzug: { up: () => Promise }; + schema: string | undefined; + logger?: Logger; + failMessage: string; +}): Promise { + try { + if (sequelize.getDialect() !== 'postgres') { + await umzug.up(); + + return; + } + + // The migration lock holds one pool connection while umzug opens a second. + const { pool } = sequelize.connectionManager as unknown as { + pool?: { maxSize?: number; write?: { maxSize?: number }; read?: { maxSize?: number } }; + }; + // With replication configured, `pool` splits into write/read sub-pools and `maxSize` is + // undefined; the migration lock runs on the write pool, so size the check against it. + const poolMax = pool?.maxSize ?? pool?.write?.maxSize ?? pool?.read?.maxSize ?? 1; + + if (poolMax < 2) { + throw new Error( + 'workflow-executor requires pool.max >= 2 on Postgres: the migration lock holds one connection while migrations run on another', + ); + } + + // Schema in its own committed transaction so umzug (on other connections) sees it. + if (schema) { + await withMigrationLock(sequelize, transaction => + sequelize.query(`CREATE SCHEMA IF NOT EXISTS "${schema}"`, { transaction }), + ); + } + + await withMigrationLock(sequelize, () => umzug.up()); + } catch (error) { + logger?.('Error', failMessage, { error: extractErrorMessage(error) }); + throw error; + } +} diff --git a/packages/workflow-executor/src/types/validated/step-outcome.ts b/packages/workflow-executor/src/types/validated/step-outcome.ts index 58b4337aee..35d086d2a0 100644 --- a/packages/workflow-executor/src/types/validated/step-outcome.ts +++ b/packages/workflow-executor/src/types/validated/step-outcome.ts @@ -9,6 +9,11 @@ export type BaseStepStatus = z.infer; export const RecordStepStatusSchema = z.enum(['success', 'error', 'awaiting-input']); export type RecordStepStatus = z.infer; +// Typed reason for an awaiting-input pause the user must resolve out of band: the Forest server +// validates it and the front reads it to prompt the matching action. +export const AwaitingInputReasonSchema = z.enum(['needs-oauth-reauth']); +export type AwaitingInputReason = z.infer; + export type StepStatus = BaseStepStatus | RecordStepStatus; /** @@ -49,6 +54,8 @@ export const McpStepOutcomeSchema = z ...baseOutcomeFields, type: z.literal('mcp'), status: RecordStepStatusSchema, + /** Present when status is 'awaiting-input' because the step paused for re-authentication. */ + awaitingInputReason: AwaitingInputReasonSchema.optional(), }) .strict(); export type McpStepOutcome = z.infer; diff --git a/packages/workflow-executor/test/adapters/ai-client-adapter.test.ts b/packages/workflow-executor/test/adapters/ai-client-adapter.test.ts index 2c87b82951..5dee091eac 100644 --- a/packages/workflow-executor/test/adapters/ai-client-adapter.test.ts +++ b/packages/workflow-executor/test/adapters/ai-client-adapter.test.ts @@ -2,12 +2,14 @@ import AiClientAdapter from '../../src/adapters/ai-client-adapter'; const mockGetModel = jest.fn().mockReturnValue({ invoke: jest.fn() }); const mockLoadRemoteTools = jest.fn().mockResolvedValue([]); +const mockLoadRemoteToolsWithFailures = jest.fn().mockResolvedValue({ tools: [], failures: [] }); const mockCloseConnections = jest.fn().mockResolvedValue(undefined); jest.mock('@forestadmin/ai-proxy', () => ({ AiClient: jest.fn().mockImplementation(() => ({ getModel: mockGetModel, loadRemoteTools: mockLoadRemoteTools, + loadRemoteToolsWithFailures: mockLoadRemoteToolsWithFailures, closeConnections: mockCloseConnections, })), })); @@ -44,6 +46,15 @@ describe('AiClientAdapter', () => { expect(mockLoadRemoteTools).toHaveBeenCalledWith(configs); }); + it('delegates loadRemoteToolsWithFailures to AiClient', async () => { + const adapter = new AiClientAdapter([]); + const configs = {}; + + await adapter.loadRemoteToolsWithFailures(configs); + + expect(mockLoadRemoteToolsWithFailures).toHaveBeenCalledWith(configs); + }); + it('delegates closeConnections to AiClient', async () => { const adapter = new AiClientAdapter([]); diff --git a/packages/workflow-executor/test/adapters/always-error-ai-model-port.test.ts b/packages/workflow-executor/test/adapters/always-error-ai-model-port.test.ts index 91489f26f6..8a2cab8722 100644 --- a/packages/workflow-executor/test/adapters/always-error-ai-model-port.test.ts +++ b/packages/workflow-executor/test/adapters/always-error-ai-model-port.test.ts @@ -32,6 +32,15 @@ describe('AlwaysErrorAiModelPort', () => { }); }); + describe('loadRemoteToolsWithFailures', () => { + it('returns no tools and no failures', async () => { + await expect(port.loadRemoteToolsWithFailures({} as never)).resolves.toEqual({ + tools: [], + failures: [], + }); + }); + }); + describe('closeConnections', () => { it('resolves without error', async () => { await expect(port.closeConnections()).resolves.toBeUndefined(); diff --git a/packages/workflow-executor/test/adapters/server-ai-adapter.test.ts b/packages/workflow-executor/test/adapters/server-ai-adapter.test.ts index cd765e16ce..42b917973e 100644 --- a/packages/workflow-executor/test/adapters/server-ai-adapter.test.ts +++ b/packages/workflow-executor/test/adapters/server-ai-adapter.test.ts @@ -2,6 +2,7 @@ import ServerAiAdapter from '../../src/adapters/server-ai-adapter'; const mockGetModel = jest.fn().mockReturnValue({ id: 'fake-model' }); const mockLoadRemoteTools = jest.fn().mockResolvedValue([]); +const mockLoadRemoteToolsWithFailures = jest.fn().mockResolvedValue({ tools: [], failures: [] }); const mockCloseConnections = jest.fn().mockResolvedValue(undefined); const mockAiClientConstructor = jest.fn(); @@ -12,6 +13,7 @@ jest.mock('@forestadmin/ai-proxy', () => ({ return { getModel: mockGetModel, loadRemoteTools: mockLoadRemoteTools, + loadRemoteToolsWithFailures: mockLoadRemoteToolsWithFailures, closeConnections: mockCloseConnections, }; }), @@ -132,6 +134,17 @@ describe('ServerAiAdapter', () => { }); }); + describe('loadRemoteToolsWithFailures', () => { + it('delegates to internal AiClient with the given configs', async () => { + const configs = {}; + + const result = await adapter.loadRemoteToolsWithFailures(configs); + + expect(mockLoadRemoteToolsWithFailures).toHaveBeenCalledWith(configs); + expect(result).toEqual({ tools: [], failures: [] }); + }); + }); + describe('closeConnections', () => { it('delegates to internal AiClient', async () => { await adapter.closeConnections(); diff --git a/packages/workflow-executor/test/adapters/step-outcome-to-update-step-mapper.test.ts b/packages/workflow-executor/test/adapters/step-outcome-to-update-step-mapper.test.ts index ce9819bd42..620cf111d0 100644 --- a/packages/workflow-executor/test/adapters/step-outcome-to-update-step-mapper.test.ts +++ b/packages/workflow-executor/test/adapters/step-outcome-to-update-step-mapper.test.ts @@ -209,4 +209,50 @@ describe('toUpdateStepRequest', () => { expect(body.stepUpdate.stepIndex).toBe(7); }); + + describe('awaitingInputReason propagation', () => { + it('puts awaitingInputReason in the update-step context (done=false) when an mcp step pauses for reauth', () => { + const outcome: StepOutcome = { + type: 'mcp', + stepId: 'step-1', + stepIndex: 4, + status: 'awaiting-input', + awaitingInputReason: 'needs-oauth-reauth', + }; + + const body = toUpdateStepRequest('42', outcome); + + expect(body.stepUpdate.attributes).toEqual({ + done: false, + context: { status: 'awaiting-input', awaitingInputReason: 'needs-oauth-reauth' }, + }); + expect(body.executionStatus).toEqual({ type: 'awaiting-input' }); + }); + + it('omits awaitingInputReason for an awaiting-input pause without a reason', () => { + const outcome: StepOutcome = { + type: 'mcp', + stepId: 'step-1', + stepIndex: 0, + status: 'awaiting-input', + }; + + const body = toUpdateStepRequest('42', outcome); + + expect(body.stepUpdate.attributes.context).toEqual({ status: 'awaiting-input' }); + }); + + it('omits awaitingInputReason for a success outcome', () => { + const outcome: StepOutcome = { + type: 'mcp', + stepId: 'step-1', + stepIndex: 0, + status: 'success', + }; + + const body = toUpdateStepRequest('42', outcome); + + expect(body.stepUpdate.attributes.context).toEqual({ status: 'success' }); + }); + }); }); diff --git a/packages/workflow-executor/test/build-workflow-executor.test.ts b/packages/workflow-executor/test/build-workflow-executor.test.ts index 30c0e1b72d..7e8dc01011 100644 --- a/packages/workflow-executor/test/build-workflow-executor.test.ts +++ b/packages/workflow-executor/test/build-workflow-executor.test.ts @@ -1,9 +1,14 @@ import ForestServerWorkflowPort from '../src/adapters/forest-server-workflow-port'; import { buildDatabaseExecutor, buildInMemoryExecutor } from '../src/build-workflow-executor'; +import CredentialEncryption from '../src/crypto/credential-encryption'; import { DEFAULT_SCHEMA_CACHE_TTL_S } from '../src/defaults'; +import ExecutorHttpServer from '../src/http/executor-http-server'; +import OAuthTokenService from '../src/oauth/token-service'; import Runner from '../src/runner'; import SchemaCache from '../src/schema-cache'; +import DatabaseMcpOAuthCredentialsStore from '../src/stores/database-mcp-oauth-credentials-store'; import DatabaseStore from '../src/stores/database-store'; +import InMemoryMcpOAuthCredentialsStore from '../src/stores/in-memory-mcp-oauth-credentials-store'; import InMemoryStore from '../src/stores/in-memory-store'; jest.mock('../src/runner'); @@ -56,6 +61,25 @@ describe('buildInMemoryExecutor', () => { ); }); + it('wires the in-memory OAuth credentials store and encryption into the HTTP server', () => { + buildInMemoryExecutor(BASE_OPTIONS); + + expect(ExecutorHttpServer).toHaveBeenCalledWith( + expect.objectContaining({ + mcpOAuthCredentialsStore: expect.any(InMemoryMcpOAuthCredentialsStore), + credentialEncryption: expect.any(CredentialEncryption), + }), + ); + }); + + it('wires an OAuth token service into the in-memory runner', () => { + buildInMemoryExecutor(BASE_OPTIONS); + + expect(MockedRunner).toHaveBeenCalledWith( + expect.objectContaining({ mcpOAuthTokenService: expect.any(OAuthTokenService) }), + ); + }); + it('creates ForestServerWorkflowPort with default forestServerUrl', () => { buildInMemoryExecutor(BASE_OPTIONS); @@ -258,6 +282,17 @@ describe('buildDatabaseExecutor', () => { ); }); + it('wires the database OAuth credentials store and encryption into the HTTP server', () => { + buildDatabaseExecutor(DB_OPTIONS); + + expect(ExecutorHttpServer).toHaveBeenCalledWith( + expect.objectContaining({ + mcpOAuthCredentialsStore: expect.any(DatabaseMcpOAuthCredentialsStore), + credentialEncryption: expect.any(CredentialEncryption), + }), + ); + }); + it('creates Sequelize with uri and passes remaining options through', () => { buildDatabaseExecutor(DB_OPTIONS); @@ -440,6 +475,23 @@ describe('WorkflowExecutor lifecycle', () => { expect(MockedRunner.prototype.stop).toHaveBeenCalled(); }); + it('start() stops the runner when server.start fails (all-or-nothing boot)', async () => { + // eslint-disable-next-line @typescript-eslint/no-var-requires, global-require + const { default: MockedHttpServer } = require('../src/http/executor-http-server'); + const originalStart = MockedHttpServer.prototype.start; + MockedHttpServer.prototype.start = jest.fn().mockRejectedValue(new Error('migration failed')); + + try { + const exec = buildInMemoryExecutor(BASE_OPTIONS); + + await expect(exec.start()).rejects.toThrow('migration failed'); + // the migration source is in server.start(); the runner must not be left polling + expect(MockedRunner.prototype.stop).toHaveBeenCalled(); + } finally { + MockedHttpServer.prototype.start = originalStart; + } + }); + it('state getter returns runner state', () => { expect(executor.state).toBe('running'); }); @@ -498,3 +550,37 @@ describe('WorkflowExecutor lifecycle', () => { } }); }); + +describe('FOREST_EXECUTOR_ENCRYPTION_KEY startup warning', () => { + const ENV_KEY = 'FOREST_EXECUTOR_ENCRYPTION_KEY'; + const original = process.env[ENV_KEY]; + + afterEach(() => { + if (original === undefined) delete process.env[ENV_KEY]; + else process.env[ENV_KEY] = original; + }); + + it('warns at startup when the encryption key is unset (no default key — fail closed)', () => { + delete process.env[ENV_KEY]; + const logger = jest.fn(); + + buildInMemoryExecutor({ ...BASE_OPTIONS, logger }); + + expect(logger).toHaveBeenCalledWith( + 'Warn', + expect.stringContaining('FOREST_EXECUTOR_ENCRYPTION_KEY'), + ); + }); + + it('does not warn when the encryption key is set', () => { + process.env[ENV_KEY] = 'a'.repeat(64); + const logger = jest.fn(); + + buildInMemoryExecutor({ ...BASE_OPTIONS, logger }); + + expect(logger).not.toHaveBeenCalledWith( + 'Warn', + expect.stringContaining('FOREST_EXECUTOR_ENCRYPTION_KEY'), + ); + }); +}); diff --git a/packages/workflow-executor/test/crypto/credential-encryption.test.ts b/packages/workflow-executor/test/crypto/credential-encryption.test.ts new file mode 100644 index 0000000000..795dfe83bf --- /dev/null +++ b/packages/workflow-executor/test/crypto/credential-encryption.test.ts @@ -0,0 +1,196 @@ +/** + * Spec for the at-rest credential encryption helper. + * + * Behaviour: + * - Key is derived in-process via HKDF (`crypto.hkdfSync`, fixed context label) from a dedicated + * `FOREST_EXECUTOR_ENCRYPTION_KEY` env var — separate from `FOREST_AUTH_SECRET`. + * - The key is read LAZILY (never required at construction / boot). + * - AES-GCM is used (authenticated encryption — tampering must be detected on decrypt). + * - Fail closed: a missing key (or a failed decrypt) must throw, never return plaintext/garbage. + * + * Key rotation is a hard swap (re-consent per (user, server)), not version-aware multi-key decrypt, + * so there is no enc-key-version concept; `decrypt` takes only the packed ciphertext. + */ +import CredentialEncryption from '../../src/crypto/credential-encryption'; +import { ExecutorEncryptionKeyMissingError } from '../../src/errors'; + +const ENV_KEY = 'FOREST_EXECUTOR_ENCRYPTION_KEY'; +// 32-byte key as 64 hex chars (mirrors the envSecret format validated elsewhere). +const TEST_KEY = 'a'.repeat(64); +const OTHER_KEY = 'b'.repeat(64); + +describe('CredentialEncryption', () => { + const original = process.env[ENV_KEY]; + + beforeEach(() => { + process.env[ENV_KEY] = TEST_KEY; + }); + + afterEach(() => { + if (original === undefined) delete process.env[ENV_KEY]; + else process.env[ENV_KEY] = original; + }); + + describe('round-trip', () => { + it('decrypts back to the exact plaintext that was encrypted', () => { + const enc = new CredentialEncryption(); + const plaintext = 'refresh-token-abc123'; + + const { ciphertext } = enc.encrypt(plaintext); + + expect(enc.decrypt(ciphertext)).toBe(plaintext); + }); + + it('round-trips multi-byte unicode without corruption', () => { + const enc = new CredentialEncryption(); + const plaintext = 'tökén-🔐-Ω-secret'; + + const { ciphertext } = enc.encrypt(plaintext); + + expect(enc.decrypt(ciphertext)).toBe(plaintext); + }); + + it('round-trips an empty string (boundary: zero-length plaintext)', () => { + const enc = new CredentialEncryption(); + + const { ciphertext } = enc.encrypt(''); + + expect(enc.decrypt(ciphertext)).toBe(''); + }); + }); + + describe('output shape', () => { + it('returns ciphertext as a Buffer (blob-storable)', () => { + const enc = new CredentialEncryption(); + + const { ciphertext } = enc.encrypt('secret'); + + expect(Buffer.isBuffer(ciphertext)).toBe(true); + }); + + it('does not leak the plaintext into the ciphertext bytes', () => { + const enc = new CredentialEncryption(); + const plaintext = 'super-secret-refresh-token'; + + const { ciphertext } = enc.encrypt(plaintext); + + expect(ciphertext.toString('utf8')).not.toContain(plaintext); + expect(ciphertext.toString('latin1')).not.toContain(plaintext); + }); + }); + + describe('non-determinism (random IV per encryption)', () => { + it('produces different ciphertext for the same plaintext on repeated calls', () => { + const enc = new CredentialEncryption(); + + const a = enc.encrypt('same-plaintext'); + const b = enc.encrypt('same-plaintext'); + + expect(a.ciphertext.toString('hex')).not.toBe(b.ciphertext.toString('hex')); + }); + + it('still decrypts both independently to the same plaintext', () => { + const enc = new CredentialEncryption(); + + const a = enc.encrypt('same-plaintext'); + const b = enc.encrypt('same-plaintext'); + + expect(enc.decrypt(a.ciphertext)).toBe('same-plaintext'); + expect(enc.decrypt(b.ciphertext)).toBe('same-plaintext'); + }); + }); + + describe('authenticity (AES-GCM) — fail closed on tampering', () => { + it('throws when a ciphertext byte is flipped', () => { + const enc = new CredentialEncryption(); + const { ciphertext } = enc.encrypt('secret'); + + const tampered = Buffer.from(ciphertext.toString('hex'), 'hex'); + const last = tampered.length - 1; + tampered[last] = (tampered[last] + 1) % 256; + + expect(() => enc.decrypt(tampered)).toThrow(); + }); + + it('throws when a byte in the IV region is flipped', () => { + const enc = new CredentialEncryption(); + const { ciphertext } = enc.encrypt('secret'); + + // Packed layout is iv | authTag | ciphertext; the IV is the first 12 bytes. + const tampered = Buffer.from(ciphertext.toString('hex'), 'hex'); + tampered[0] = (tampered[0] + 1) % 256; + + expect(() => enc.decrypt(tampered)).toThrow(); + }); + + it('throws when a byte in the auth-tag region is flipped', () => { + const enc = new CredentialEncryption(); + const { ciphertext } = enc.encrypt('secret'); + + // The 16-byte auth tag immediately follows the 12-byte IV. + const tampered = Buffer.from(ciphertext.toString('hex'), 'hex'); + tampered[12] = (tampered[12] + 1) % 256; + + expect(() => enc.decrypt(tampered)).toThrow(); + }); + + it('throws when the ciphertext is truncated', () => { + const enc = new CredentialEncryption(); + const { ciphertext } = enc.encrypt('secret'); + + const truncated = ciphertext.subarray(0, ciphertext.length - 1); + + expect(() => enc.decrypt(truncated)).toThrow(); + }); + + it('throws when decrypting under a different key (cross-key, fail closed)', () => { + const enc = new CredentialEncryption(); + const { ciphertext } = enc.encrypt('secret'); + + // Rotate the host key out from under the same payload. + process.env[ENV_KEY] = OTHER_KEY; + const other = new CredentialEncryption(); + + expect(() => other.decrypt(ciphertext)).toThrow(); + }); + }); + + describe('key derivation', () => { + it('derives the same key across instances (empty-salt HKDF is deterministic)', () => { + const writer = new CredentialEncryption(); + const { ciphertext } = writer.encrypt('cross-instance-secret'); + + // A separate instance reading the same env key must decrypt the payload — this is what lets a + // restarted or horizontally-scaled executor read rows written by another instance, and pins + // the empty-salt / fixed-info derivation as stable rather than per-instance. + const reader = new CredentialEncryption(); + + expect(reader.decrypt(ciphertext)).toBe('cross-instance-secret'); + }); + }); + + describe('lazy key reading', () => { + it('does not throw at construction when the key is unset', () => { + delete process.env[ENV_KEY]; + + expect(() => new CredentialEncryption()).not.toThrow(); + }); + + it('throws ExecutorEncryptionKeyMissingError on encrypt when the key is unset', () => { + delete process.env[ENV_KEY]; + const enc = new CredentialEncryption(); + + expect(() => enc.encrypt('secret')).toThrow(ExecutorEncryptionKeyMissingError); + }); + + it('throws ExecutorEncryptionKeyMissingError on decrypt when the key is unset', () => { + const enc = new CredentialEncryption(); + const { ciphertext } = enc.encrypt('secret'); + + delete process.env[ENV_KEY]; + const cold = new CredentialEncryption(); + + expect(() => cold.decrypt(ciphertext)).toThrow(ExecutorEncryptionKeyMissingError); + }); + }); +}); diff --git a/packages/workflow-executor/test/errors.test.ts b/packages/workflow-executor/test/errors.test.ts index be03b9562d..17ffbf92e5 100644 --- a/packages/workflow-executor/test/errors.test.ts +++ b/packages/workflow-executor/test/errors.test.ts @@ -2,6 +2,9 @@ import { AiModelPortError, InvalidPendingDataError, NoMcpToolsError, + OAuthInvalidGrantError, + OAuthReauthRequiredError, + OAuthRefreshError, PendingDataNotFoundError, extractErrorMessage, } from '../src/errors'; @@ -143,3 +146,34 @@ describe('InvalidPendingDataError', () => { expect(err.userMessage).toBe('The request body is invalid.'); }); }); + +describe('OAuthReauthRequiredError', () => { + it("defaults awaitingInputReason to 'needs-oauth-reauth' and names the server in the technical message", () => { + const err = new OAuthReauthRequiredError('srv-1'); + + expect(err.awaitingInputReason).toBe('needs-oauth-reauth'); + expect(err.message).toMatch(/srv-1/); + }); + + it('keeps the user-facing message free of the internal server id', () => { + const err = new OAuthReauthRequiredError('srv-1'); + + expect(err.userMessage).not.toMatch(/srv-1/); + }); +}); + +describe('OAuthRefreshError', () => { + it('prefixes the technical message and stores the cause', () => { + const cause = new Error('5xx'); + const err = new OAuthRefreshError('endpoint unreachable', cause); + + expect(err.message).toMatch(/endpoint unreachable/); + expect(err.cause).toBe(cause); + }); +}); + +describe('OAuthInvalidGrantError', () => { + it('includes the detail when provided', () => { + expect(new OAuthInvalidGrantError('token expired').message).toMatch(/token expired/); + }); +}); diff --git a/packages/workflow-executor/test/executors/base-step-executor.test.ts b/packages/workflow-executor/test/executors/base-step-executor.test.ts index 63085ca5f1..b2b17fd97e 100644 --- a/packages/workflow-executor/test/executors/base-step-executor.test.ts +++ b/packages/workflow-executor/test/executors/base-step-executor.test.ts @@ -98,6 +98,7 @@ function makeMockRunStore(stepExecutions: StepExecutionData[] = []): RunStore { close: jest.fn().mockResolvedValue(undefined), getStepExecutions: jest.fn().mockResolvedValue(stepExecutions), saveStepExecution: jest.fn().mockResolvedValue(undefined), + deleteStepExecution: jest.fn().mockResolvedValue(undefined), }; } diff --git a/packages/workflow-executor/test/executors/condition-step-executor.test.ts b/packages/workflow-executor/test/executors/condition-step-executor.test.ts index d0c422596e..0242da473b 100644 --- a/packages/workflow-executor/test/executors/condition-step-executor.test.ts +++ b/packages/workflow-executor/test/executors/condition-step-executor.test.ts @@ -31,6 +31,7 @@ function makeMockRunStore(overrides: Partial = {}): RunStore { close: jest.fn().mockResolvedValue(undefined), getStepExecutions: jest.fn().mockResolvedValue([]), saveStepExecution: jest.fn().mockResolvedValue(undefined), + deleteStepExecution: jest.fn().mockResolvedValue(undefined), ...overrides, }; } diff --git a/packages/workflow-executor/test/executors/guidance-step-executor.test.ts b/packages/workflow-executor/test/executors/guidance-step-executor.test.ts index 7abc8cebca..308be4058e 100644 --- a/packages/workflow-executor/test/executors/guidance-step-executor.test.ts +++ b/packages/workflow-executor/test/executors/guidance-step-executor.test.ts @@ -20,6 +20,7 @@ function makeMockRunStore(overrides: Partial = {}): RunStore { close: jest.fn().mockResolvedValue(undefined), getStepExecutions: jest.fn().mockResolvedValue([]), saveStepExecution: jest.fn().mockResolvedValue(undefined), + deleteStepExecution: jest.fn().mockResolvedValue(undefined), ...overrides, }; } diff --git a/packages/workflow-executor/test/executors/load-related-record-step-executor.test.ts b/packages/workflow-executor/test/executors/load-related-record-step-executor.test.ts index 229242d07b..9e24e9b5bd 100644 --- a/packages/workflow-executor/test/executors/load-related-record-step-executor.test.ts +++ b/packages/workflow-executor/test/executors/load-related-record-step-executor.test.ts @@ -130,6 +130,7 @@ function makeMockRunStore(overrides: Partial = {}): RunStore { close: jest.fn().mockResolvedValue(undefined), getStepExecutions: jest.fn().mockResolvedValue([]), saveStepExecution: jest.fn().mockResolvedValue(undefined), + deleteStepExecution: jest.fn().mockResolvedValue(undefined), ...overrides, }; } diff --git a/packages/workflow-executor/test/executors/mcp-step-executor.test.ts b/packages/workflow-executor/test/executors/mcp-step-executor.test.ts index 7f88825796..cc03afe1ff 100644 --- a/packages/workflow-executor/test/executors/mcp-step-executor.test.ts +++ b/packages/workflow-executor/test/executors/mcp-step-executor.test.ts @@ -8,12 +8,13 @@ import type { McpStepDefinition } from '../../src/types/validated/step-definitio import RemoteTool from '@forestadmin/ai-proxy/src/remote-tool'; -import { RunStorePortError, StepStateError } from '../../src/errors'; +import { OAuthReauthRequiredError, RunStorePortError, StepStateError } from '../../src/errors'; import ActivityLog from '../../src/executors/activity-log'; import AgentWithLog from '../../src/executors/agent-with-log'; import McpStepExecutor from '../../src/executors/mcp-step-executor'; import SchemaCache from '../../src/schema-cache'; import SchemaResolver from '../../src/schema-resolver'; +import InMemoryStore from '../../src/stores/in-memory-store'; import { StepExecutionMode, StepType } from '../../src/types/validated/step-definition'; // --------------------------------------------------------------------------- @@ -58,6 +59,7 @@ function makeMockRunStore(overrides: Partial = {}): RunStore { close: jest.fn().mockResolvedValue(undefined), getStepExecutions: jest.fn().mockResolvedValue([]), saveStepExecution: jest.fn().mockResolvedValue(undefined), + deleteStepExecution: jest.fn().mockResolvedValue(undefined), ...overrides, }; } @@ -1137,3 +1139,336 @@ describe('McpStepExecutor', () => { }); }); }); + +// The executor's OAuth responsibility is the tool-call 401 retry and the re-auth pause mapping. +// authType identification, the credential lookup, list-tools retry and Bearer injection live in +// RemoteToolFetcher / StepExecutorFactory and are covered by their own suites. +describe('McpStepExecutor — OAuth2 tool-call re-authentication', () => { + const authError = () => new Error('Request failed with status 401'); + + function makeAutomated(invoke: jest.Mock) { + const tool = new MockRemoteTool({ + name: 'send_notification', + sourceId: 'mcp-server-1', + invoke, + }); + const { model } = makeMockModel('send_notification', { message: 'Hello' }); + const context = makeContext({ + model, + stepDefinition: makeStep({ executionType: StepExecutionMode.FullyAutomated }), + }); + + return { tool, context }; + } + + it('force-refreshes, rebuilds the tool, and retries once when the call returns 401', async () => { + const { tool, context } = makeAutomated(jest.fn().mockRejectedValue(authError())); + const freshInvoke = jest.fn().mockResolvedValue('ok-after-refresh'); + const freshTool = new MockRemoteTool({ + name: 'send_notification', + sourceId: 'mcp-server-1', + invoke: freshInvoke, + }); + const reloadWithFreshAuth = jest.fn().mockResolvedValue([freshTool]); + + const result = await new McpStepExecutor(context, [tool], 'srv', reloadWithFreshAuth).execute(); + + expect(result.stepOutcome.status).toBe('success'); + expect(reloadWithFreshAuth).toHaveBeenCalledTimes(1); + expect(freshInvoke).toHaveBeenCalledWith({ message: 'Hello' }); + }); + + it("pauses with awaiting-input/'needs-oauth-reauth' when the credential can no longer be refreshed", async () => { + const { tool, context } = makeAutomated(jest.fn().mockRejectedValue(authError())); + const reloadWithFreshAuth = jest.fn().mockRejectedValue(new OAuthReauthRequiredError('srv')); + + const result = await new McpStepExecutor(context, [tool], 'srv', reloadWithFreshAuth).execute(); + + expect(result.stepOutcome).toMatchObject({ + status: 'awaiting-input', + awaitingInputReason: 'needs-oauth-reauth', + }); + }); + + it("pauses with 'needs-oauth-reauth' when the retried call still returns 401", async () => { + const { tool, context } = makeAutomated(jest.fn().mockRejectedValue(authError())); + const freshTool = new MockRemoteTool({ + name: 'send_notification', + sourceId: 'mcp-server-1', + invoke: jest.fn().mockRejectedValue(authError()), + }); + const reloadWithFreshAuth = jest.fn().mockResolvedValue([freshTool]); + + const result = await new McpStepExecutor(context, [tool], 'srv', reloadWithFreshAuth).execute(); + + expect(result.stepOutcome).toMatchObject({ + status: 'awaiting-input', + awaitingInputReason: 'needs-oauth-reauth', + }); + expect(reloadWithFreshAuth).toHaveBeenCalledTimes(1); + }); + + it('does not refresh for a non-auth tool error — surfaces it as a step error', async () => { + const { tool, context } = makeAutomated(jest.fn().mockRejectedValue(new Error('boom'))); + const reloadWithFreshAuth = jest.fn(); + + const result = await new McpStepExecutor(context, [tool], 'srv', reloadWithFreshAuth).execute(); + + expect(result.stepOutcome.status).toBe('error'); + expect(reloadWithFreshAuth).not.toHaveBeenCalled(); + }); + + it('does not refresh on a 401 for a bearer/none step (no reload hook)', async () => { + const { tool, context } = makeAutomated(jest.fn().mockRejectedValue(authError())); + + const result = await new McpStepExecutor(context, [tool], 'srv').execute(); + + expect(result.stepOutcome.status).toBe('error'); + }); + + it('surfaces a non-auth error from the retried call as a step error (not a reauth pause)', async () => { + const { tool, context } = makeAutomated(jest.fn().mockRejectedValue(authError())); + const freshTool = new MockRemoteTool({ + name: 'send_notification', + sourceId: 'mcp-server-1', + invoke: jest.fn().mockRejectedValue(new Error('downstream 500')), + }); + const reloadWithFreshAuth = jest.fn().mockResolvedValue([freshTool]); + + const result = await new McpStepExecutor(context, [tool], 'srv', reloadWithFreshAuth).execute(); + + expect(result.stepOutcome.status).toBe('error'); + }); + + it('errors when the refreshed tool set no longer contains the selected tool', async () => { + const { tool, context } = makeAutomated(jest.fn().mockRejectedValue(authError())); + const reloadWithFreshAuth = jest.fn().mockResolvedValue([]); + + const result = await new McpStepExecutor(context, [tool], 'srv', reloadWithFreshAuth).execute(); + + expect(result.stepOutcome.status).toBe('error'); + }); +}); + +// On a re-auth pause the executor clears the 'executing' write-ahead marker so the resumed step is +// not rejected as interrupted; the failed tool call still surfaces as a failed audit-log entry. +describe('McpStepExecutor — re-auth pause hardening', () => { + const authError = () => new Error('Request failed with status 401'); + + // A FullyAutomated step whose tool call 401s and whose refresh cannot recover → re-auth pause. + function pauseFor(runStore: ReturnType | InMemoryStore) { + const tool = new MockRemoteTool({ + name: 'send_notification', + sourceId: 'mcp-server-1', + invoke: jest.fn().mockRejectedValue(authError()), + }); + const reloadWithFreshAuth = jest.fn().mockRejectedValue(new OAuthReauthRequiredError('srv')); + const context = makeContext({ + runStore: runStore as RunStore, + stepDefinition: makeStep({ executionType: StepExecutionMode.FullyAutomated }), + }); + + return new McpStepExecutor(context, [tool], 'srv', reloadWithFreshAuth); + } + + describe('idempotency phase cleared on the re-auth pause path', () => { + it('does not leave the step execution marked executing after pausing for re-auth', async () => { + // GIVEN a real store so the persisted write-ahead marker is observable. + const store = new InMemoryStore(); + + // WHEN the step pauses for re-authentication. + const result = await pauseFor(store).execute(); + + // THEN it pauses, and the 'executing' marker written by beforeCall must not survive — else + // a resume would read a stale marker. + expect(result.stepOutcome.status).toBe('awaiting-input'); + const persisted = await store.getStepExecutions('run-1'); + const mcpExecution = persisted.find(execution => execution.stepIndex === 0) as + | McpStepExecutionData + | undefined; + expect(mcpExecution?.idempotencyPhase).not.toBe('executing'); + }); + + it('resumes to success after a re-auth pause instead of failing as interrupted', async () => { + // GIVEN a step that paused for re-auth, persisting its state in a shared store. + const store = new InMemoryStore(); + const pause = await pauseFor(store).execute(); + expect(pause.stepOutcome.status).toBe('awaiting-input'); + + // WHEN the user reconnects and the step is re-dispatched against the same store, with the + // tool now succeeding. + const reconnectedTool = new MockRemoteTool({ + name: 'send_notification', + sourceId: 'mcp-server-1', + invoke: jest.fn().mockResolvedValue('ok-after-reconnect'), + }); + const resumeContext = makeContext({ + runStore: store, + stepDefinition: makeStep({ executionType: StepExecutionMode.FullyAutomated }), + }); + + const resumed = await new McpStepExecutor(resumeContext, [reconnectedTool], 'srv').execute(); + + // THEN checkIdempotency must not throw StepStateError on the stale 'executing' marker; the + // step runs to success. + expect(resumed.stepOutcome.status).toBe('success'); + }); + + it('preserves the approved pendingData on a confirmation-flow re-auth pause so resume replays it', async () => { + // GIVEN a confirmation-flow step with a user-approved tool call already persisted. + const store = new InMemoryStore(); + await store.saveStepExecution('run-1', { + type: 'mcp', + stepIndex: 0, + pendingData: { + name: 'send_notification', + sourceId: 'mcp-server-1', + input: { message: 'Hello' }, + }, + userConfirmation: { userConfirmed: true }, + } as McpStepExecutionData); + const tool = new MockRemoteTool({ + name: 'send_notification', + sourceId: 'mcp-server-1', + invoke: jest.fn().mockRejectedValue(authError()), + }); + const reloadWithFreshAuth = jest.fn().mockRejectedValue(new OAuthReauthRequiredError('srv')); + const context = makeContext({ runStore: store }); + + // WHEN the approved call 401s and cannot re-auth. + const result = await new McpStepExecutor( + context, + [tool], + 'srv', + reloadWithFreshAuth, + ).execute(); + + // THEN the marker is cleared but the approved call is kept, so a resume replays it rather than + // re-selecting a tool. + expect(result.stepOutcome.status).toBe('awaiting-input'); + const persisted = (await store.getStepExecutions('run-1')).find(e => e.stepIndex === 0) as + | McpStepExecutionData + | undefined; + expect(persisted?.idempotencyPhase).not.toBe('executing'); + expect(persisted?.pendingData).toEqual({ + name: 'send_notification', + sourceId: 'mcp-server-1', + input: { message: 'Hello' }, + }); + }); + + it('surfaces a store error from the re-auth cleanup as a step error, not a stuck pause', async () => { + // GIVEN cleanup that fails — a left-behind 'executing' marker would make the pause + // non-resumable, so the failure must surface as an error rather than a stuck awaiting-input. + const store = new InMemoryStore(); + jest + .spyOn(store, 'deleteStepExecution') + .mockRejectedValue(new RunStorePortError('deleteStepExecution', new Error('store down'))); + + // WHEN a FullyAutomated step pauses for re-auth (no pendingData → cleanup goes via delete). + const result = await pauseFor(store).execute(); + + // THEN the store failure propagates to a step error instead of a non-resumable pause. + expect(result.stepOutcome.status).toBe('error'); + expect(result.stepOutcome.error).toBe('The step state could not be accessed. Please retry.'); + }); + + it('clears the executing marker when the refresh fails with a non-auth error, leaving the step retryable', async () => { + // GIVEN a 401 that triggers a refresh which then fails with a non-auth error (e.g. the token + // endpoint is unreachable): nothing executed, so the step must stay retryable, not wedged. + const store = new InMemoryStore(); + const tool = new MockRemoteTool({ + name: 'send_notification', + sourceId: 'mcp-server-1', + invoke: jest.fn().mockRejectedValue(authError()), + }); + const reloadWithFreshAuth = jest + .fn() + .mockRejectedValue(new Error('token endpoint unreachable')); + const context = makeContext({ + runStore: store, + stepDefinition: makeStep({ executionType: StepExecutionMode.FullyAutomated }), + }); + + // WHEN the refresh path fails with a non-auth error. + const result = await new McpStepExecutor( + context, + [tool], + 'srv', + reloadWithFreshAuth, + ).execute(); + + // THEN it surfaces as a step error (not a re-auth pause), and the 'executing' marker is gone so + // a retry is not rejected as interrupted. + expect(result.stepOutcome.status).toBe('error'); + const persisted = (await store.getStepExecutions('run-1')).find(e => e.stepIndex === 0) as + | McpStepExecutionData + | undefined; + expect(persisted?.idempotencyPhase).not.toBe('executing'); + }); + + it('clears the executing marker when the reload yields no matching tool, leaving the step retryable', async () => { + // The reload succeeds but returns no tool (empty on a connection failure); the 401 first call + // never ran, so the step must stay retryable rather than wedged as interrupted. + const store = new InMemoryStore(); + const tool = new MockRemoteTool({ + name: 'send_notification', + sourceId: 'mcp-server-1', + invoke: jest.fn().mockRejectedValue(authError()), + }); + const reloadWithFreshAuth = jest.fn().mockResolvedValue([]); + const context = makeContext({ + runStore: store, + stepDefinition: makeStep({ executionType: StepExecutionMode.FullyAutomated }), + }); + + const result = await new McpStepExecutor( + context, + [tool], + 'srv', + reloadWithFreshAuth, + ).execute(); + + expect(result.stepOutcome.status).toBe('error'); + const persisted = (await store.getStepExecutions('run-1')).find(e => e.stepIndex === 0) as + | McpStepExecutionData + | undefined; + expect(persisted?.idempotencyPhase).not.toBe('executing'); + }); + }); + + describe('re-auth pause surfaces the tool-call failure in the audit log', () => { + it('marks the activity-log entry failed while the step pauses for re-auth', async () => { + // GIVEN an activity-log port we can inspect. + const activityLogPort = { + createPending: jest.fn().mockResolvedValue({ id: 'log-1', index: '0' }), + markSucceeded: jest.fn().mockResolvedValue(undefined), + markFailed: jest.fn().mockResolvedValue(undefined), + }; + const tool = new MockRemoteTool({ + name: 'send_notification', + sourceId: 'mcp-server-1', + invoke: jest.fn().mockRejectedValue(authError()), + }); + const reloadWithFreshAuth = jest.fn().mockRejectedValue(new OAuthReauthRequiredError('srv')); + const context = makeContext({ + activityLogPort, + stepDefinition: makeStep({ executionType: StepExecutionMode.FullyAutomated }), + }); + + // WHEN the step pauses for re-authentication. + const result = await new McpStepExecutor( + context, + [tool], + 'srv', + reloadWithFreshAuth, + ).execute(); + + // THEN the failed tool call is audited as failed, while the step itself pauses (not errors). + expect(result.stepOutcome.status).toBe('awaiting-input'); + expect(activityLogPort.createPending).toHaveBeenCalledTimes(1); + expect(activityLogPort.markFailed).toHaveBeenCalledTimes(1); + expect(activityLogPort.markSucceeded).not.toHaveBeenCalled(); + }); + }); +}); diff --git a/packages/workflow-executor/test/executors/read-record-step-executor.test.ts b/packages/workflow-executor/test/executors/read-record-step-executor.test.ts index 62da87bfa4..743bbab67e 100644 --- a/packages/workflow-executor/test/executors/read-record-step-executor.test.ts +++ b/packages/workflow-executor/test/executors/read-record-step-executor.test.ts @@ -72,6 +72,7 @@ function makeMockRunStore(overrides: Partial = {}): RunStore { close: jest.fn().mockResolvedValue(undefined), getStepExecutions: jest.fn().mockResolvedValue([]), saveStepExecution: jest.fn().mockResolvedValue(undefined), + deleteStepExecution: jest.fn().mockResolvedValue(undefined), ...overrides, }; } diff --git a/packages/workflow-executor/test/executors/step-executor-factory.test.ts b/packages/workflow-executor/test/executors/step-executor-factory.test.ts new file mode 100644 index 0000000000..88e2019957 --- /dev/null +++ b/packages/workflow-executor/test/executors/step-executor-factory.test.ts @@ -0,0 +1,102 @@ +import type { StepContextConfig } from '../../src/executors/step-executor-factory'; +import type { ActivityLogPort } from '../../src/ports/activity-log-port'; +import type { AvailableStepExecution } from '../../src/types/execution-context'; + +import { OAuthReauthRequiredError } from '../../src/errors'; +import StepExecutorFactory from '../../src/executors/step-executor-factory'; +import SchemaCache from '../../src/schema-cache'; +import { StepExecutionMode, StepType } from '../../src/types/validated/step-definition'; + +const activityLogPort = { + createPending: jest.fn().mockResolvedValue({ id: 'log-1', index: '0' }), + markSucceeded: jest.fn().mockResolvedValue(undefined), + markFailed: jest.fn().mockResolvedValue(undefined), +} as unknown as ActivityLogPort; + +function makeStep(): AvailableStepExecution { + return { + runId: 'run-1', + stepId: 'step-1', + stepIndex: 0, + collectionId: 'col-1', + baseRecordRef: { collectionName: 'customers', recordId: [1], stepIndex: 0 }, + stepDefinition: { + type: StepType.Mcp, + executionType: StepExecutionMode.FullyAutomated, + mcpServerId: 'srv-1', + }, + previousSteps: [], + user: { + id: 7, + email: 'a@b.com', + firstName: 'A', + lastName: 'B', + team: 't', + renderingId: 1, + role: 'admin', + permissionLevel: 'admin', + tags: {}, + }, + } as unknown as AvailableStepExecution; +} + +function makeContextConfig(): StepContextConfig { + return { + aiModelPort: { getModel: jest.fn().mockReturnValue({}) }, + agentPort: {}, + workflowPort: {}, + runStore: {}, + schemaCache: new SchemaCache(), + logger: jest.fn(), + } as unknown as StepContextConfig; +} + +describe('StepExecutorFactory.create — MCP OAuth re-auth', () => { + it('maps OAuthReauthRequiredError from tool loading to an awaiting-input outcome with the typed reason', async () => { + const fetchRemoteTools = jest.fn().mockRejectedValue(new OAuthReauthRequiredError('srv-1')); + + const executor = await StepExecutorFactory.create( + makeStep(), + makeContextConfig(), + activityLogPort, + fetchRemoteTools, + ); + const result = await executor.execute(); + + expect(result.stepOutcome).toEqual({ + type: 'mcp', + stepId: 'step-1', + stepIndex: 0, + status: 'awaiting-input', + awaitingInputReason: 'needs-oauth-reauth', + }); + }); + + it('maps a generic tool-loading failure to an error outcome (no awaitingInputReason)', async () => { + const fetchRemoteTools = jest.fn().mockRejectedValue(new Error('kaboom')); + + const executor = await StepExecutorFactory.create( + makeStep(), + makeContextConfig(), + activityLogPort, + fetchRemoteTools, + ); + const result = await executor.execute(); + + expect(result.stepOutcome.status).toBe('error'); + expect(result.stepOutcome).not.toHaveProperty('awaitingInputReason'); + }); + + it('passes the step user id to the tool fetcher', async () => { + const fetchRemoteTools = jest.fn().mockResolvedValue({ tools: [], mcpServerName: 'srv' }); + + await StepExecutorFactory.create( + makeStep(), + makeContextConfig(), + activityLogPort, + fetchRemoteTools, + ); + + expect(fetchRemoteTools).toHaveBeenCalledWith('srv-1', 7); + }); +}); diff --git a/packages/workflow-executor/test/executors/trigger-record-action-step-executor.test.ts b/packages/workflow-executor/test/executors/trigger-record-action-step-executor.test.ts index 830bce2107..ca27fa1dae 100644 --- a/packages/workflow-executor/test/executors/trigger-record-action-step-executor.test.ts +++ b/packages/workflow-executor/test/executors/trigger-record-action-step-executor.test.ts @@ -84,6 +84,7 @@ function makeMockRunStore(overrides: Partial = {}): RunStore { close: jest.fn().mockResolvedValue(undefined), getStepExecutions: jest.fn().mockResolvedValue([]), saveStepExecution: jest.fn().mockResolvedValue(undefined), + deleteStepExecution: jest.fn().mockResolvedValue(undefined), ...overrides, }; } diff --git a/packages/workflow-executor/test/executors/update-record-step-executor.test.ts b/packages/workflow-executor/test/executors/update-record-step-executor.test.ts index e8b358c8c1..701eb71f92 100644 --- a/packages/workflow-executor/test/executors/update-record-step-executor.test.ts +++ b/packages/workflow-executor/test/executors/update-record-step-executor.test.ts @@ -77,6 +77,7 @@ function makeMockRunStore(overrides: Partial = {}): RunStore { close: jest.fn().mockResolvedValue(undefined), getStepExecutions: jest.fn().mockResolvedValue([]), saveStepExecution: jest.fn().mockResolvedValue(undefined), + deleteStepExecution: jest.fn().mockResolvedValue(undefined), ...overrides, }; } diff --git a/packages/workflow-executor/test/http/executor-http-server.test.ts b/packages/workflow-executor/test/http/executor-http-server.test.ts index 7403dd0413..f5a9767fbf 100644 --- a/packages/workflow-executor/test/http/executor-http-server.test.ts +++ b/packages/workflow-executor/test/http/executor-http-server.test.ts @@ -1,10 +1,13 @@ import type { Logger } from '../../src/ports/logger-port'; +import type { McpOAuthCredentialsStore } from '../../src/ports/mcp-oauth-credentials-store'; import type { WorkflowPort } from '../../src/ports/workflow-port'; +import type RemoteToolFetcher from '../../src/remote-tool-fetcher'; import type Runner from '../../src/runner'; import jsonwebtoken from 'jsonwebtoken'; import request from 'supertest'; +import CredentialEncryption from '../../src/crypto/credential-encryption'; import { MalformedRunError, RunAlreadyInFlightError, @@ -14,6 +17,7 @@ import { WorkflowPortError, } from '../../src/errors'; import ExecutorHttpServer from '../../src/http/executor-http-server'; +import InMemoryMcpOAuthCredentialsStore from '../../src/stores/in-memory-mcp-oauth-credentials-store'; const AUTH_SECRET = 'test-auth-secret'; @@ -45,11 +49,20 @@ function createMockWorkflowPort(overrides: Partial = {}): Workflow } as unknown as WorkflowPort; } +function createMockFetcher(): RemoteToolFetcher { + return { + fetch: jest.fn().mockResolvedValue({ tools: [], mcpServerName: undefined }), + } as unknown as RemoteToolFetcher; +} + function createServer( overrides: { runner?: Runner; workflowPort?: WorkflowPort; logger?: jest.MockedFunction; + mcpOAuthCredentialsStore?: McpOAuthCredentialsStore; + credentialEncryption?: CredentialEncryption; + remoteToolFetcher?: RemoteToolFetcher; } = {}, ) { return new ExecutorHttpServer({ @@ -58,10 +71,39 @@ function createServer( authSecret: AUTH_SECRET, workflowPort: overrides.workflowPort ?? createMockWorkflowPort(), logger: overrides.logger, + mcpOAuthCredentialsStore: + overrides.mcpOAuthCredentialsStore ?? new InMemoryMcpOAuthCredentialsStore(), + credentialEncryption: overrides.credentialEncryption ?? new CredentialEncryption(), + remoteToolFetcher: overrides.remoteToolFetcher ?? createMockFetcher(), }); } describe('ExecutorHttpServer', () => { + describe('start()', () => { + it('initializes the MCP OAuth credentials store with the logger', async () => { + const init = jest.fn().mockResolvedValue(undefined); + const logger: jest.MockedFunction = jest.fn(); + const server = new ExecutorHttpServer({ + port: 0, + runner: createMockRunner(), + authSecret: AUTH_SECRET, + workflowPort: createMockWorkflowPort(), + logger, + mcpOAuthCredentialsStore: { init } as unknown as McpOAuthCredentialsStore, + credentialEncryption: {} as unknown as CredentialEncryption, + remoteToolFetcher: createMockFetcher(), + }); + + await server.start(); + + try { + expect(init).toHaveBeenCalledWith(logger); + } finally { + await server.stop(); + } + }); + }); + describe('JWT authentication', () => { it('should return 401 when no token is provided', async () => { const server = createServer(); diff --git a/packages/workflow-executor/test/http/mcp-oauth-credentials-route.test.ts b/packages/workflow-executor/test/http/mcp-oauth-credentials-route.test.ts new file mode 100644 index 0000000000..22b648dc6f --- /dev/null +++ b/packages/workflow-executor/test/http/mcp-oauth-credentials-route.test.ts @@ -0,0 +1,561 @@ +/** + * Spec for the OAuth credential deposit endpoint. + * + * Behaviour: + * - POST /mcp-oauth-credentials and DELETE deposit/disconnect, on the SAME HTTP server as /trigger. + * - Authenticated by the existing koaJwt middleware (forest_session_token, FOREST_AUTH_SECRET). + * - user_id is taken from the validated token, NEVER from the request body. + * - The executor encrypts the refresh token (+ client secret) and upserts one row per (user, server). + * - When FOREST_EXECUTOR_ENCRYPTION_KEY is unset, encryption fails closed and the endpoint returns + * a distinct, typed `executor_encryption_key_missing` (HTTP 503) — never a generic failure. + * - Dormant for now: nothing reads the table at runtime yet; only this deposit path writes it. + * + * Endpoint contract: + * - ExecutorHttpServer options gain: `mcpOAuthCredentialsStore` (upsert/get/delete) and + * `credentialEncryption` (encrypt/decrypt). Both injected like `runner` / `workflowPort`. + * - POST body (camelCase JSON): { mcpServerId, refreshToken, clientId?, clientSecret?, + * clientSecretExpiresAt?, tokenEndpoint?, tokenEndpointAuthMethod?, scopes? }. + * - DELETE path: /mcp-oauth-credentials/:mcpServerId. + * - Typed key-missing response: HTTP 503 with body { code: 'executor_encryption_key_missing' }. + */ +import jsonwebtoken from 'jsonwebtoken'; +import request from 'supertest'; + +import { ExecutorEncryptionKeyMissingError, OAuthReauthRequiredError } from '../../src/errors'; +import ExecutorHttpServer from '../../src/http/executor-http-server'; + +const AUTH_SECRET = 'test-auth-secret'; + +function signToken(payload: object, secret = AUTH_SECRET, options?: jsonwebtoken.SignOptions) { + return jsonwebtoken.sign(payload, secret, { expiresIn: '1h', ...options }); +} + +function createMockRunner() { + return { + state: 'running', + start: jest.fn().mockResolvedValue(undefined), + stop: jest.fn().mockResolvedValue(undefined), + triggerPoll: jest.fn().mockResolvedValue(undefined), + getRunStepExecutions: jest.fn().mockResolvedValue([]), + }; +} + +function createMockWorkflowPort() { + return { + getAvailableRuns: jest.fn().mockResolvedValue({ pending: [], malformed: [] }), + getAvailableRun: jest.fn(), + updateStepExecution: jest.fn().mockResolvedValue(undefined), + getCollectionSchema: jest.fn(), + getMcpServerConfigs: jest.fn().mockResolvedValue({}), + hasRunAccess: jest.fn().mockResolvedValue(true), + }; +} + +function createMockStore() { + return { + init: jest.fn().mockResolvedValue(undefined), + upsert: jest.fn().mockResolvedValue(undefined), + updateIfPresent: jest.fn().mockResolvedValue(undefined), + get: jest.fn().mockResolvedValue(null), + delete: jest.fn().mockResolvedValue(undefined), + close: jest.fn().mockResolvedValue(undefined), + }; +} + +function createMockEncryption() { + return { + // Deterministic stub: the route under test only needs an opaque blob back. + encrypt: jest.fn((plaintext: string) => ({ + ciphertext: Buffer.from(`enc(${plaintext})`), + })), + decrypt: jest.fn(), + }; +} + +function createMockFetcher() { + return { fetch: jest.fn().mockResolvedValue({ tools: [], mcpServerName: undefined }) }; +} + +function createMockTokenService() { + return { evict: jest.fn() }; +} + +function createServer(overrides: Record = {}) { + return new ExecutorHttpServer({ + port: 0, + runner: createMockRunner(), + authSecret: AUTH_SECRET, + workflowPort: createMockWorkflowPort(), + mcpOAuthCredentialsStore: createMockStore(), + credentialEncryption: createMockEncryption(), + remoteToolFetcher: createMockFetcher(), + ...overrides, + } as never); +} + +const validBody = { + mcpServerId: 'mcp-server-1', + refreshToken: 'refresh-token-xyz', + clientId: 'client-abc', + clientSecret: 'client-secret-123', + tokenEndpoint: 'https://auth.example.com/token', + tokenEndpointAuthMethod: 'client_secret_post', + scopes: 'read write', +}; + +describe('POST /mcp-oauth-credentials', () => { + describe('authentication', () => { + it('returns 401 when no token is provided', async () => { + const server = createServer(); + + const response = await request(server.callback) + .post('/mcp-oauth-credentials') + .send(validBody); + + expect(response.status).toBe(401); + expect(response.body).toEqual({ error: 'Unauthorized' }); + }); + + it('returns 401 when the token is signed with the wrong secret', async () => { + const server = createServer(); + const token = signToken({ id: 1 }, 'wrong-secret'); + + const response = await request(server.callback) + .post('/mcp-oauth-credentials') + .set('Authorization', `Bearer ${token}`) + .send(validBody); + + expect(response.status).toBe(401); + }); + + it('does not write to the store when unauthenticated', async () => { + const store = createMockStore(); + const server = createServer({ mcpOAuthCredentialsStore: store }); + + await request(server.callback).post('/mcp-oauth-credentials').send(validBody); + + expect(store.upsert).not.toHaveBeenCalled(); + }); + }); + + describe('user identity from token', () => { + it('upserts using the user id from the token', async () => { + const store = createMockStore(); + const server = createServer({ mcpOAuthCredentialsStore: store }); + const token = signToken({ id: 7 }); + + await request(server.callback) + .post('/mcp-oauth-credentials') + .set('Authorization', `Bearer ${token}`) + .send(validBody); + + expect(store.upsert).toHaveBeenCalledWith( + expect.objectContaining({ userId: 7, mcpServerId: 'mcp-server-1' }), + ); + }); + + it('evicts any cached access token on deposit so a reconnect takes effect immediately', async () => { + const tokenService = createMockTokenService(); + const server = createServer({ oauthTokenService: tokenService }); + const token = signToken({ id: 7 }); + + const response = await request(server.callback) + .post('/mcp-oauth-credentials') + .set('Authorization', `Bearer ${token}`) + .send(validBody); + + expect(response.status).toBe(200); + expect(tokenService.evict).toHaveBeenCalledWith(7, 'mcp-server-1'); + }); + + it('rejects a body that tries to supply a user id (the token is the only source)', async () => { + const store = createMockStore(); + const server = createServer({ mcpOAuthCredentialsStore: store }); + const token = signToken({ id: 7 }); + + const response = await request(server.callback) + .post('/mcp-oauth-credentials') + .set('Authorization', `Bearer ${token}`) + .send({ ...validBody, userId: 999, user_id: 999 }); + + expect(response.status).toBe(400); + expect(store.upsert).not.toHaveBeenCalled(); + }); + + it('returns 401 when the token carries no numeric id (rejected by the claims middleware)', async () => { + const store = createMockStore(); + const server = createServer({ mcpOAuthCredentialsStore: store }); + const token = signToken({ email: 'no-id@example.com' }); + + const response = await request(server.callback) + .post('/mcp-oauth-credentials') + .set('Authorization', `Bearer ${token}`) + .send(validBody); + + expect(response.status).toBe(401); + expect(store.upsert).not.toHaveBeenCalled(); + }); + }); + + describe('encryption before persistence', () => { + it('encrypts the refresh token and stores only the ciphertext (never plaintext)', async () => { + const store = createMockStore(); + const encryption = createMockEncryption(); + const server = createServer({ + mcpOAuthCredentialsStore: store, + credentialEncryption: encryption, + }); + const token = signToken({ id: 1 }); + + await request(server.callback) + .post('/mcp-oauth-credentials') + .set('Authorization', `Bearer ${token}`) + .send(validBody); + + expect(encryption.encrypt).toHaveBeenCalledWith('refresh-token-xyz'); + const persisted = store.upsert.mock.calls[0][0]; + expect(Buffer.isBuffer(persisted.refreshTokenEnc)).toBe(true); + expect(persisted.refreshTokenEnc.toString()).toBe('enc(refresh-token-xyz)'); + // The plaintext must not have been handed to the store under any field. + expect(JSON.stringify(persisted)).not.toContain('refresh-token-xyz'); + }); + + it('encrypts the client secret when one is provided', async () => { + const store = createMockStore(); + const encryption = createMockEncryption(); + const server = createServer({ + mcpOAuthCredentialsStore: store, + credentialEncryption: encryption, + }); + const token = signToken({ id: 1 }); + + await request(server.callback) + .post('/mcp-oauth-credentials') + .set('Authorization', `Bearer ${token}`) + .send(validBody); + + expect(encryption.encrypt).toHaveBeenCalledWith('client-secret-123'); + expect(store.upsert.mock.calls[0][0].clientSecretEnc.toString()).toBe( + 'enc(client-secret-123)', + ); + }); + + it('stores a null client secret for a public / PKCE client (no clientSecret in body)', async () => { + const store = createMockStore(); + const server = createServer({ mcpOAuthCredentialsStore: store }); + const token = signToken({ id: 1 }); + const { clientSecret, ...publicBody } = validBody; + + await request(server.callback) + .post('/mcp-oauth-credentials') + .set('Authorization', `Bearer ${token}`) + .send({ ...publicBody, tokenEndpointAuthMethod: 'none' }); + + expect(store.upsert).toHaveBeenCalledWith(expect.objectContaining({ clientSecretEnc: null })); + }); + }); + + describe('fail closed when the encryption key is missing', () => { + it('returns 503 with a typed executor_encryption_key_missing code', async () => { + const encryption = createMockEncryption(); + encryption.encrypt.mockImplementation(() => { + throw new ExecutorEncryptionKeyMissingError(); + }); + const server = createServer({ credentialEncryption: encryption }); + const token = signToken({ id: 1 }); + + const response = await request(server.callback) + .post('/mcp-oauth-credentials') + .set('Authorization', `Bearer ${token}`) + .send(validBody); + + expect(response.status).toBe(503); + expect(response.body).toEqual( + expect.objectContaining({ code: 'executor_encryption_key_missing' }), + ); + }); + + it('does not persist anything when the key is missing', async () => { + const store = createMockStore(); + const encryption = createMockEncryption(); + encryption.encrypt.mockImplementation(() => { + throw new ExecutorEncryptionKeyMissingError(); + }); + const server = createServer({ + mcpOAuthCredentialsStore: store, + credentialEncryption: encryption, + }); + const token = signToken({ id: 1 }); + + await request(server.callback) + .post('/mcp-oauth-credentials') + .set('Authorization', `Bearer ${token}`) + .send(validBody); + + expect(store.upsert).not.toHaveBeenCalled(); + }); + }); + + describe('body validation', () => { + it('returns 400 when the refresh token is missing', async () => { + const store = createMockStore(); + const server = createServer({ mcpOAuthCredentialsStore: store }); + const token = signToken({ id: 1 }); + const { refreshToken, ...noRefresh } = validBody; + + const response = await request(server.callback) + .post('/mcp-oauth-credentials') + .set('Authorization', `Bearer ${token}`) + .send(noRefresh); + + expect(response.status).toBe(400); + expect(store.upsert).not.toHaveBeenCalled(); + }); + + it('returns 400 when mcpServerId is missing', async () => { + const store = createMockStore(); + const server = createServer({ mcpOAuthCredentialsStore: store }); + const token = signToken({ id: 1 }); + const { mcpServerId, ...noServer } = validBody; + + const response = await request(server.callback) + .post('/mcp-oauth-credentials') + .set('Authorization', `Bearer ${token}`) + .send(noServer); + + expect(response.status).toBe(400); + expect(store.upsert).not.toHaveBeenCalled(); + }); + + it('returns 400 when tokenEndpoint is missing', async () => { + const store = createMockStore(); + const server = createServer({ mcpOAuthCredentialsStore: store }); + const token = signToken({ id: 1 }); + const { tokenEndpoint, ...noEndpoint } = validBody; + + const response = await request(server.callback) + .post('/mcp-oauth-credentials') + .set('Authorization', `Bearer ${token}`) + .send(noEndpoint); + + expect(response.status).toBe(400); + expect(store.upsert).not.toHaveBeenCalled(); + }); + + it('returns 400 when tokenEndpoint points at a link-local / metadata address (SSRF guard)', async () => { + const store = createMockStore(); + const server = createServer({ mcpOAuthCredentialsStore: store }); + const token = signToken({ id: 1 }); + + const response = await request(server.callback) + .post('/mcp-oauth-credentials') + .set('Authorization', `Bearer ${token}`) + .send({ ...validBody, tokenEndpoint: 'http://169.254.169.254/latest/meta-data' }); + + expect(response.status).toBe(400); + expect(store.upsert).not.toHaveBeenCalled(); + }); + + it('returns 400 when a field has the wrong type', async () => { + const store = createMockStore(); + const server = createServer({ mcpOAuthCredentialsStore: store }); + const token = signToken({ id: 1 }); + + const response = await request(server.callback) + .post('/mcp-oauth-credentials') + .set('Authorization', `Bearer ${token}`) + .send({ ...validBody, refreshToken: 12345 }); + + expect(response.status).toBe(400); + expect(store.upsert).not.toHaveBeenCalled(); + }); + + it('returns 400 when clientSecretExpiresAt is not a parseable date', async () => { + const store = createMockStore(); + const server = createServer({ mcpOAuthCredentialsStore: store }); + const token = signToken({ id: 1 }); + + const response = await request(server.callback) + .post('/mcp-oauth-credentials') + .set('Authorization', `Bearer ${token}`) + .send({ ...validBody, clientSecretExpiresAt: 'not-a-date' }); + + expect(response.status).toBe(400); + expect(store.upsert).not.toHaveBeenCalled(); + }); + + it('returns 400 when clientSecret is an empty string (not silently treated as a public client)', async () => { + const store = createMockStore(); + const server = createServer({ mcpOAuthCredentialsStore: store }); + const token = signToken({ id: 1 }); + + const response = await request(server.callback) + .post('/mcp-oauth-credentials') + .set('Authorization', `Bearer ${token}`) + .send({ ...validBody, clientSecret: '' }); + + expect(response.status).toBe(400); + expect(store.upsert).not.toHaveBeenCalled(); + }); + }); + + describe('store failure', () => { + it('returns 500 when the store rejects', async () => { + const store = createMockStore(); + store.upsert.mockRejectedValue(new Error('db down')); + const server = createServer({ mcpOAuthCredentialsStore: store }); + const token = signToken({ id: 1 }); + + const response = await request(server.callback) + .post('/mcp-oauth-credentials') + .set('Authorization', `Bearer ${token}`) + .send(validBody); + + expect(response.status).toBe(500); + }); + }); +}); + +describe('DELETE /mcp-oauth-credentials/:mcpServerId', () => { + it('returns 401 when no token is provided', async () => { + const server = createServer(); + + const response = await request(server.callback).delete('/mcp-oauth-credentials/mcp-server-1'); + + expect(response.status).toBe(401); + }); + + it('deletes the credential for (token user, mcpServerId) and returns 204 with no body', async () => { + const store = createMockStore(); + const server = createServer({ mcpOAuthCredentialsStore: store }); + const token = signToken({ id: 7 }); + + const response = await request(server.callback) + .delete('/mcp-oauth-credentials/mcp-server-1') + .set('Authorization', `Bearer ${token}`); + + expect(response.status).toBe(204); + expect(response.body).toEqual({}); + expect(store.delete).toHaveBeenCalledWith(7, 'mcp-server-1'); + }); + + it('evicts the in-process cached access token so the disconnect takes effect immediately', async () => { + const oauthTokenService = createMockTokenService(); + const server = createServer({ oauthTokenService }); + const token = signToken({ id: 7 }); + + await request(server.callback) + .delete('/mcp-oauth-credentials/mcp-server-1') + .set('Authorization', `Bearer ${token}`); + + expect(oauthTokenService.evict).toHaveBeenCalledWith(7, 'mcp-server-1'); + }); + + it("does not delete another user's credential", async () => { + const store = createMockStore(); + const server = createServer({ mcpOAuthCredentialsStore: store }); + const token = signToken({ id: 7 }); + + await request(server.callback) + .delete('/mcp-oauth-credentials/mcp-server-1') + .set('Authorization', `Bearer ${token}`); + + expect(store.delete).not.toHaveBeenCalledWith(999, expect.anything()); + }); + + it('returns 401 when the token carries no numeric id (rejected by the claims middleware)', async () => { + const store = createMockStore(); + const server = createServer({ mcpOAuthCredentialsStore: store }); + const token = signToken({ email: 'no-id@example.com' }); + + const response = await request(server.callback) + .delete('/mcp-oauth-credentials/mcp-server-1') + .set('Authorization', `Bearer ${token}`); + + expect(response.status).toBe(401); + expect(store.delete).not.toHaveBeenCalled(); + }); +}); + +describe('GET /list-mcp-tools', () => { + const toolDef = (name: string, description: string) => ({ base: { name, description } }); + + it('returns 401 when no token is provided', async () => { + const server = createServer(); + + const response = await request(server.callback).get('/list-mcp-tools?mcpServerId=mcp-server-1'); + + expect(response.status).toBe(401); + }); + + it('lists the tools for (token user, mcpServerId), resolving the vault credential', async () => { + const fetcher = { + fetch: jest.fn().mockResolvedValue({ + tools: [toolDef('search', 'Search things'), toolDef('create', 'Create a thing')], + mcpServerName: 'srv', + }), + }; + const server = createServer({ remoteToolFetcher: fetcher }); + const token = signToken({ id: 7 }); + + const response = await request(server.callback) + .get('/list-mcp-tools?mcpServerId=mcp-server-1') + .set('Authorization', `Bearer ${token}`); + + expect(response.status).toBe(200); + // user_id comes from the validated JWT (7), never the request. + expect(fetcher.fetch).toHaveBeenCalledWith('mcp-server-1', 7); + expect(response.body).toEqual({ + tools: [ + { name: 'search', description: 'Search things' }, + { name: 'create', description: 'Create a thing' }, + ], + }); + }); + + it('returns 400 when mcpServerId is missing', async () => { + const fetcher = { fetch: jest.fn() }; + const server = createServer({ remoteToolFetcher: fetcher }); + const token = signToken({ id: 7 }); + + const response = await request(server.callback) + .get('/list-mcp-tools') + .set('Authorization', `Bearer ${token}`); + + expect(response.status).toBe(400); + expect(fetcher.fetch).not.toHaveBeenCalled(); + }); + + it('returns a typed needs-oauth-reauth (409) when the credential is missing or unrefreshable', async () => { + const fetcher = { + fetch: jest.fn().mockRejectedValue(new OAuthReauthRequiredError('mcp-server-1')), + }; + const server = createServer({ remoteToolFetcher: fetcher }); + const token = signToken({ id: 7 }); + + const response = await request(server.callback) + .get('/list-mcp-tools?mcpServerId=mcp-server-1') + .set('Authorization', `Bearer ${token}`); + + expect(response.status).toBe(409); + expect(response.body).toEqual({ + awaitingInputReason: 'needs-oauth-reauth', + mcpServerId: 'mcp-server-1', + }); + }); + + it('returns the typed 503 executor_encryption_key_missing when the key is absent', async () => { + const fetcher = { + fetch: jest.fn().mockRejectedValue(new ExecutorEncryptionKeyMissingError()), + }; + const server = createServer({ remoteToolFetcher: fetcher }); + const token = signToken({ id: 7 }); + + const response = await request(server.callback) + .get('/list-mcp-tools?mcpServerId=mcp-server-1') + .set('Authorization', `Bearer ${token}`); + + expect(response.status).toBe(503); + expect(response.body).toEqual({ code: 'executor_encryption_key_missing' }); + }); +}); diff --git a/packages/workflow-executor/test/http/mcp-oauth-credentials.test.ts b/packages/workflow-executor/test/http/mcp-oauth-credentials.test.ts new file mode 100644 index 0000000000..093446111b --- /dev/null +++ b/packages/workflow-executor/test/http/mcp-oauth-credentials.test.ts @@ -0,0 +1,134 @@ +import type CredentialEncryption from '../../src/crypto/credential-encryption'; +import type { DepositCredentialsBody } from '../../src/http/mcp-oauth-credentials'; + +import { ExecutorEncryptionKeyMissingError } from '../../src/errors'; +import { + buildMcpOAuthCredentialInput, + depositCredentialsBodySchema, +} from '../../src/http/mcp-oauth-credentials'; + +function createEncryption(): CredentialEncryption { + return { + encrypt: jest.fn((plaintext: string) => ({ + ciphertext: Buffer.from(`enc(${plaintext})`), + })), + decrypt: jest.fn(), + } as unknown as CredentialEncryption; +} + +const fullBody: DepositCredentialsBody = { + mcpServerId: 'mcp-server-1', + refreshToken: 'refresh-token-xyz', + clientId: 'client-abc', + clientSecret: 'client-secret-123', + clientSecretExpiresAt: '2030-01-02T03:04:05.000Z', + tokenEndpoint: 'https://auth.example.com/token', + tokenEndpointAuthMethod: 'client_secret_post', + scopes: 'read write', +}; + +describe('buildMcpOAuthCredentialInput', () => { + it('encrypts the refresh token and maps the record for the given user', () => { + const encryption = createEncryption(); + + const input = buildMcpOAuthCredentialInput({ body: fullBody, userId: 7, encryption }); + + expect(encryption.encrypt).toHaveBeenCalledWith('refresh-token-xyz'); + expect(input.userId).toBe(7); + expect(input.mcpServerId).toBe('mcp-server-1'); + expect(input.refreshTokenEnc.toString()).toBe('enc(refresh-token-xyz)'); + expect(input.tokenEndpoint).toBe('https://auth.example.com/token'); + expect(input.tokenEndpointAuthMethod).toBe('client_secret_post'); + expect(input.scopes).toBe('read write'); + }); + + it('encrypts the client secret and parses the expiry when both are provided', () => { + const encryption = createEncryption(); + + const input = buildMcpOAuthCredentialInput({ body: fullBody, userId: 7, encryption }); + + expect(encryption.encrypt).toHaveBeenCalledWith('client-secret-123'); + expect(input.clientSecretEnc?.toString()).toBe('enc(client-secret-123)'); + expect(input.clientSecretExpiresAt).toEqual(new Date('2030-01-02T03:04:05.000Z')); + }); + + it('leaves optional client fields null for a public / PKCE client', () => { + const encryption = createEncryption(); + const publicBody: DepositCredentialsBody = { + mcpServerId: 'mcp-server-1', + refreshToken: 'refresh-token-xyz', + tokenEndpoint: 'https://auth.example.com/token', + }; + + const input = buildMcpOAuthCredentialInput({ body: publicBody, userId: 7, encryption }); + + expect(encryption.encrypt).toHaveBeenCalledTimes(1); + expect(input.clientId).toBeNull(); + expect(input.clientSecretEnc).toBeNull(); + expect(input.clientSecretExpiresAt).toBeNull(); + expect(input.tokenEndpointAuthMethod).toBeNull(); + expect(input.scopes).toBeNull(); + }); + + it('propagates ExecutorEncryptionKeyMissingError so the caller can fail closed', () => { + const encryption = createEncryption(); + (encryption.encrypt as jest.Mock).mockImplementation(() => { + throw new ExecutorEncryptionKeyMissingError(); + }); + + expect(() => buildMcpOAuthCredentialInput({ body: fullBody, userId: 7, encryption })).toThrow( + ExecutorEncryptionKeyMissingError, + ); + }); +}); + +describe('depositCredentialsBodySchema', () => { + const validBody = { + mcpServerId: 'mcp-server-1', + refreshToken: 'refresh-token-xyz', + tokenEndpoint: 'https://auth.example.com/token', + }; + + it('rejects a clientSecret supplied without a clientId', () => { + const result = depositCredentialsBodySchema.safeParse({ ...validBody, clientSecret: 'secret' }); + + expect(result.success).toBe(false); + }); + + it('rejects an empty-string clientId (would drop client_id at refresh time)', () => { + const result = depositCredentialsBodySchema.safeParse({ ...validBody, clientId: '' }); + + expect(result.success).toBe(false); + }); + + it('accepts a clientSecret paired with a clientId', () => { + const result = depositCredentialsBodySchema.safeParse({ + ...validBody, + clientId: 'client-abc', + clientSecret: 'secret', + }); + + expect(result.success).toBe(true); + }); + + it('rejects an unsupported tokenEndpointAuthMethod', () => { + const result = depositCredentialsBodySchema.safeParse({ + ...validBody, + tokenEndpointAuthMethod: 'client_secret_posst', + }); + + expect(result.success).toBe(false); + }); + + it.each(['client_secret_basic', 'client_secret_post', 'none'] as const)( + 'accepts the supported client-authentication method %s', + method => { + const result = depositCredentialsBodySchema.safeParse({ + ...validBody, + tokenEndpointAuthMethod: method, + }); + + expect(result.success).toBe(true); + }, + ); +}); diff --git a/packages/workflow-executor/test/integration/workflow-execution.test.ts b/packages/workflow-executor/test/integration/workflow-execution.test.ts index 7bfe84ff6c..5c8f7b7486 100644 --- a/packages/workflow-executor/test/integration/workflow-execution.test.ts +++ b/packages/workflow-executor/test/integration/workflow-execution.test.ts @@ -9,9 +9,14 @@ import jsonwebtoken from 'jsonwebtoken'; import request from 'supertest'; import { z } from 'zod'; +import createConsoleLogger from '../../src/adapters/console-logger'; +import CredentialEncryption from '../../src/crypto/credential-encryption'; import ExecutorHttpServer from '../../src/http/executor-http-server'; +import OAuthTokenService from '../../src/oauth/token-service'; +import RemoteToolFetcher from '../../src/remote-tool-fetcher'; import Runner from '../../src/runner'; import SchemaCache from '../../src/schema-cache'; +import InMemoryMcpOAuthCredentialsStore from '../../src/stores/in-memory-mcp-oauth-credentials-store'; import InMemoryStore from '../../src/stores/in-memory-store'; import { StepExecutionMode, StepType } from '../../src/types/validated/step-definition'; @@ -192,6 +197,13 @@ function createIntegrationSetup(overrides?: { const agentPort = overrides?.agentPort ?? createMockAgentPort(); const runStore = new InMemoryStore(); const schemaCache = new SchemaCache(); + const mcpOAuthCredentialsStore = new InMemoryMcpOAuthCredentialsStore(); + const credentialEncryption = new CredentialEncryption(); + const mcpOAuthTokenService = new OAuthTokenService({ + store: mcpOAuthCredentialsStore, + encryption: credentialEncryption, + logger: createConsoleLogger(), + }); const runner = new Runner({ agentPort, @@ -210,13 +222,24 @@ function createIntegrationSetup(overrides?: { pollingIntervalS: overrides?.pollingIntervalS ?? 60, envSecret: ENV_SECRET, authSecret: AUTH_SECRET, + mcpOAuthTokenService, }); + const remoteToolFetcher = new RemoteToolFetcher( + workflowPort, + aiClient, + createConsoleLogger(), + mcpOAuthTokenService, + ); + const server = new ExecutorHttpServer({ port: 0, runner, authSecret: AUTH_SECRET, workflowPort, + mcpOAuthCredentialsStore, + credentialEncryption, + remoteToolFetcher, }); return { runner, server, workflowPort, agentPort, runStore, aiClient, model }; diff --git a/packages/workflow-executor/test/oauth/keyed-mutex.test.ts b/packages/workflow-executor/test/oauth/keyed-mutex.test.ts new file mode 100644 index 0000000000..02d0e756e9 --- /dev/null +++ b/packages/workflow-executor/test/oauth/keyed-mutex.test.ts @@ -0,0 +1,79 @@ +import KeyedMutex from '../../src/oauth/keyed-mutex'; + +function deferred() { + let resolve!: (value: T) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((res, rej) => { + resolve = res; + reject = rej; + }); + + return { promise, resolve, reject }; +} + +describe('KeyedMutex', () => { + describe('runExclusive', () => { + it('returns the value the task resolves with', async () => { + const mutex = new KeyedMutex(); + + await expect(mutex.runExclusive('k', async () => 42)).resolves.toBe(42); + }); + + it('runs tasks sharing a key one at a time, in order', async () => { + const mutex = new KeyedMutex(); + const events: string[] = []; + const first = deferred(); + + const firstRun = mutex.runExclusive('user:1', async () => { + events.push('first:start'); + await first.promise; + events.push('first:end'); + }); + + const secondRun = mutex.runExclusive('user:1', async () => { + events.push('second:start'); + }); + + // The second task must not start while the first holds the key. + await Promise.resolve(); + expect(events).toEqual(['first:start']); + + first.resolve(); + await Promise.all([firstRun, secondRun]); + + expect(events).toEqual(['first:start', 'first:end', 'second:start']); + }); + + it('runs tasks for different keys concurrently', async () => { + const mutex = new KeyedMutex(); + const events: string[] = []; + const blocker = deferred(); + + const blocked = mutex.runExclusive('user:1', async () => { + events.push('a:start'); + await blocker.promise; + }); + const other = mutex.runExclusive('user:2', async () => { + events.push('b:start'); + }); + + await other; + // 'user:2' completed without waiting for the still-blocked 'user:1'. + expect(events).toEqual(['a:start', 'b:start']); + + blocker.resolve(); + await blocked; + }); + + it('lets a later task on the same key run after an earlier task rejects', async () => { + const mutex = new KeyedMutex(); + + const failed = mutex.runExclusive('user:1', async () => { + throw new Error('boom'); + }); + + await expect(failed).rejects.toThrow('boom'); + await expect(mutex.runExclusive('user:1', async () => 'ok')).resolves.toBe('ok'); + }); + }); +}); diff --git a/packages/workflow-executor/test/oauth/refresh-grant.test.ts b/packages/workflow-executor/test/oauth/refresh-grant.test.ts new file mode 100644 index 0000000000..76cfc50807 --- /dev/null +++ b/packages/workflow-executor/test/oauth/refresh-grant.test.ts @@ -0,0 +1,295 @@ +import { + InvalidTokenEndpointError, + OAuthInvalidGrantError, + OAuthRefreshError, +} from '../../src/errors'; +import refreshAccessToken from '../../src/oauth/refresh-grant'; + +function mockResponse(options: { + ok: boolean; + status: number; + payload?: unknown; + nonJson?: boolean; +}) { + return { + ok: options.ok, + status: options.status, + json: options.nonJson + ? () => Promise.reject(new Error('not json')) + : () => Promise.resolve(options.payload ?? {}), + } as unknown as Response; +} + +describe('refreshAccessToken', () => { + let fetchSpy: jest.SpyInstance; + + beforeEach(() => { + fetchSpy = jest.spyOn(global, 'fetch'); + }); + + afterEach(() => { + fetchSpy.mockRestore(); + }); + + function lastRequest() { + const [url, init] = fetchSpy.mock.calls[0]; + const body = init.body as URLSearchParams; + + return { url, headers: init.headers as Record, body }; + } + + it('rejects a link-local token endpoint without making a request (SSRF guard)', async () => { + await expect( + refreshAccessToken({ tokenEndpoint: 'http://169.254.169.254/token', refreshToken: 'rt-1' }), + ).rejects.toBeInstanceOf(InvalidTokenEndpointError); + + expect(fetchSpy).not.toHaveBeenCalled(); + }); + + it('sends the request with redirect:manual so redirects are never followed', async () => { + fetchSpy.mockResolvedValue( + mockResponse({ ok: true, status: 200, payload: { access_token: 'at' } }), + ); + + await refreshAccessToken({ tokenEndpoint: 'https://idp/token', refreshToken: 'rt-1' }); + + expect(fetchSpy.mock.calls[0][1].redirect).toBe('manual'); + }); + + it('rejects a redirect response without following it or reading the body (SSRF via redirect)', async () => { + const jsonSpy = jest.fn(); + fetchSpy.mockResolvedValue({ + ok: false, + status: 302, + type: 'opaqueredirect', + json: jsonSpy, + } as unknown as Response); + + await expect( + refreshAccessToken({ tokenEndpoint: 'https://idp/token', refreshToken: 'rt-1' }), + ).rejects.toBeInstanceOf(OAuthRefreshError); + + // The grant body (refresh token / client secret) must not be re-sent nor the reply parsed. + expect(jsonSpy).not.toHaveBeenCalled(); + }); + + it('posts a refresh_token grant with the refresh token to the token endpoint', async () => { + fetchSpy.mockResolvedValue( + mockResponse({ ok: true, status: 200, payload: { access_token: 'at' } }), + ); + + await refreshAccessToken({ tokenEndpoint: 'https://idp/token', refreshToken: 'rt-1' }); + + const { url, body, headers } = lastRequest(); + expect(url).toBe('https://idp/token'); + expect(fetchSpy.mock.calls[0][1].method).toBe('POST'); + expect(headers['content-type']).toBe('application/x-www-form-urlencoded'); + expect(body.get('grant_type')).toBe('refresh_token'); + expect(body.get('refresh_token')).toBe('rt-1'); + }); + + it('throws OAuthRefreshError (not a TypeError) when the token endpoint body is literal null', async () => { + fetchSpy.mockResolvedValue({ + ok: false, + status: 500, + json: () => Promise.resolve(null), + } as unknown as Response); + + await expect( + refreshAccessToken({ tokenEndpoint: 'https://idp/token', refreshToken: 'rt-1' }), + ).rejects.toBeInstanceOf(OAuthRefreshError); + }); + + it('sends client credentials via Basic auth for client_secret_basic (the default with a secret)', async () => { + fetchSpy.mockResolvedValue( + mockResponse({ ok: true, status: 200, payload: { access_token: 'at' } }), + ); + + await refreshAccessToken({ + tokenEndpoint: 'https://idp/token', + refreshToken: 'rt-1', + clientId: 'cid', + clientSecret: 'secret', + }); + + const { headers, body } = lastRequest(); + expect(headers.authorization).toBe(`Basic ${Buffer.from('cid:secret').toString('base64')}`); + expect(body.get('client_secret')).toBeNull(); + }); + + it('form-encodes the Basic credentials per RFC 6749 (a space becomes +, not %20)', async () => { + fetchSpy.mockResolvedValue( + mockResponse({ ok: true, status: 200, payload: { access_token: 'at' } }), + ); + + await refreshAccessToken({ + tokenEndpoint: 'https://idp/token', + refreshToken: 'rt-1', + clientId: 'client id', + clientSecret: 'sec ret', + }); + + const { headers } = lastRequest(); + expect(headers.authorization).toBe( + `Basic ${Buffer.from('client+id:sec+ret').toString('base64')}`, + ); + }); + + it('sends client credentials in the body for client_secret_post', async () => { + fetchSpy.mockResolvedValue( + mockResponse({ ok: true, status: 200, payload: { access_token: 'at' } }), + ); + + await refreshAccessToken({ + tokenEndpoint: 'https://idp/token', + refreshToken: 'rt-1', + clientId: 'cid', + clientSecret: 'secret', + tokenEndpointAuthMethod: 'client_secret_post', + }); + + const { headers, body } = lastRequest(); + expect(headers.authorization).toBeUndefined(); + expect(body.get('client_id')).toBe('cid'); + expect(body.get('client_secret')).toBe('secret'); + }); + + it('sends only client_id for a public client (no secret)', async () => { + fetchSpy.mockResolvedValue( + mockResponse({ ok: true, status: 200, payload: { access_token: 'at' } }), + ); + + await refreshAccessToken({ + tokenEndpoint: 'https://idp/token', + refreshToken: 'rt-1', + clientId: 'cid', + }); + + const { headers, body } = lastRequest(); + expect(headers.authorization).toBeUndefined(); + expect(body.get('client_id')).toBe('cid'); + expect(body.get('client_secret')).toBeNull(); + }); + + it('includes the scope parameter only when scopes are provided', async () => { + fetchSpy.mockResolvedValue( + mockResponse({ ok: true, status: 200, payload: { access_token: 'at' } }), + ); + + await refreshAccessToken({ + tokenEndpoint: 'https://idp/token', + refreshToken: 'rt-1', + scopes: 'a b', + }); + expect(lastRequest().body.get('scope')).toBe('a b'); + + fetchSpy.mockClear(); + await refreshAccessToken({ tokenEndpoint: 'https://idp/token', refreshToken: 'rt-1' }); + expect(lastRequest().body.get('scope')).toBeNull(); + }); + + it('returns the access token, expiry, and a rotated refresh token on success', async () => { + fetchSpy.mockResolvedValue( + mockResponse({ + ok: true, + status: 200, + payload: { access_token: 'at-1', expires_in: 3600, refresh_token: 'rt-2' }, + }), + ); + + const result = await refreshAccessToken({ + tokenEndpoint: 'https://idp/token', + refreshToken: 'rt-1', + }); + + expect(result).toEqual({ accessToken: 'at-1', expiresInS: 3600, refreshToken: 'rt-2' }); + }); + + it('leaves expiresInS undefined when the response omits expires_in', async () => { + fetchSpy.mockResolvedValue( + mockResponse({ ok: true, status: 200, payload: { access_token: 'at-1' } }), + ); + + const result = await refreshAccessToken({ + tokenEndpoint: 'https://idp/token', + refreshToken: 'rt-1', + }); + + expect(result.expiresInS).toBeUndefined(); + expect(result.refreshToken).toBeUndefined(); + }); + + it('coerces a numeric-string expires_in so the token is still cacheable', async () => { + fetchSpy.mockResolvedValue( + mockResponse({ + ok: true, + status: 200, + payload: { access_token: 'at-1', expires_in: '3600' }, + }), + ); + + const result = await refreshAccessToken({ + tokenEndpoint: 'https://idp/token', + refreshToken: 'rt-1', + }); + + expect(result.expiresInS).toBe(3600); + }); + + it('throws OAuthInvalidGrantError when the endpoint returns error invalid_grant', async () => { + fetchSpy.mockResolvedValue( + mockResponse({ ok: false, status: 400, payload: { error: 'invalid_grant' } }), + ); + + await expect( + refreshAccessToken({ tokenEndpoint: 'https://idp/token', refreshToken: 'rt-1' }), + ).rejects.toBeInstanceOf(OAuthInvalidGrantError); + }); + + it('throws OAuthRefreshError for a generic error response', async () => { + fetchSpy.mockResolvedValue( + mockResponse({ ok: false, status: 400, payload: { error: 'invalid_scope' } }), + ); + + await expect( + refreshAccessToken({ tokenEndpoint: 'https://idp/token', refreshToken: 'rt-1' }), + ).rejects.toBeInstanceOf(OAuthRefreshError); + }); + + it.each(['invalid_client', 'unauthorized_client'])( + 'maps a dead client credential (%s) to re-auth', + async error => { + fetchSpy.mockResolvedValue(mockResponse({ ok: false, status: 401, payload: { error } })); + + await expect( + refreshAccessToken({ tokenEndpoint: 'https://idp/token', refreshToken: 'rt-1' }), + ).rejects.toBeInstanceOf(OAuthInvalidGrantError); + }, + ); + + it('throws OAuthRefreshError on a 5xx from the token endpoint', async () => { + fetchSpy.mockResolvedValue(mockResponse({ ok: false, status: 503, nonJson: true })); + + await expect( + refreshAccessToken({ tokenEndpoint: 'https://idp/token', refreshToken: 'rt-1' }), + ).rejects.toBeInstanceOf(OAuthRefreshError); + }); + + it('throws OAuthRefreshError when the endpoint is unreachable', async () => { + fetchSpy.mockRejectedValue(new Error('ECONNREFUSED')); + + await expect( + refreshAccessToken({ tokenEndpoint: 'https://idp/token', refreshToken: 'rt-1' }), + ).rejects.toBeInstanceOf(OAuthRefreshError); + }); + + it('throws OAuthRefreshError when a 200 response carries no access_token', async () => { + fetchSpy.mockResolvedValue( + mockResponse({ ok: true, status: 200, payload: { token_type: 'bearer' } }), + ); + + await expect( + refreshAccessToken({ tokenEndpoint: 'https://idp/token', refreshToken: 'rt-1' }), + ).rejects.toBeInstanceOf(OAuthRefreshError); + }); +}); diff --git a/packages/workflow-executor/test/oauth/token-endpoint-url.test.ts b/packages/workflow-executor/test/oauth/token-endpoint-url.test.ts new file mode 100644 index 0000000000..7eb80920a9 --- /dev/null +++ b/packages/workflow-executor/test/oauth/token-endpoint-url.test.ts @@ -0,0 +1,170 @@ +import { InvalidTokenEndpointError } from '../../src/errors'; +import assertSafeTokenEndpoint from '../../src/oauth/token-endpoint-url'; + +describe('assertSafeTokenEndpoint', () => { + const ORIGINAL_NODE_ENV = process.env.NODE_ENV; + + afterEach(() => { + process.env.NODE_ENV = ORIGINAL_NODE_ENV; + }); + + describe('any environment', () => { + beforeEach(() => { + process.env.NODE_ENV = 'development'; + }); + + it('accepts an https endpoint with a public host', () => { + expect(() => assertSafeTokenEndpoint('https://idp.example.com/oauth/token')).not.toThrow(); + }); + + it('accepts http on loopback off-production (the local OAuth sim)', () => { + expect(() => assertSafeTokenEndpoint('http://127.0.0.1:9100/token')).not.toThrow(); + expect(() => assertSafeTokenEndpoint('http://localhost:9100/token')).not.toThrow(); + }); + + it('rejects a value that is not a valid absolute URL', () => { + expect(() => assertSafeTokenEndpoint('not a url')).toThrow(InvalidTokenEndpointError); + }); + + it('rejects a non-http(s) scheme', () => { + expect(() => assertSafeTokenEndpoint('file:///etc/passwd')).toThrow( + InvalidTokenEndpointError, + ); + expect(() => assertSafeTokenEndpoint('ftp://idp/token')).toThrow(InvalidTokenEndpointError); + }); + + it('rejects a link-local / cloud-metadata address even off-production', () => { + expect(() => assertSafeTokenEndpoint('http://169.254.169.254/latest/meta-data')).toThrow( + InvalidTokenEndpointError, + ); + expect(() => assertSafeTokenEndpoint('https://169.254.169.254/token')).toThrow( + InvalidTokenEndpointError, + ); + }); + + it('rejects an integer-encoded cloud-metadata address (2852039166 -> 169.254.169.254)', () => { + expect(() => assertSafeTokenEndpoint('https://2852039166/token')).toThrow( + InvalidTokenEndpointError, + ); + }); + + it('rejects an IPv6 link-local address', () => { + expect(() => assertSafeTokenEndpoint('http://[fe80::1]/token')).toThrow( + InvalidTokenEndpointError, + ); + }); + + it('rejects an IPv4-mapped IPv6 metadata address (e.g. ::ffff:169.254.169.254)', () => { + expect(() => assertSafeTokenEndpoint('http://[::ffff:169.254.169.254]/token')).toThrow( + InvalidTokenEndpointError, + ); + }); + + it('rejects the unspecified address, which routes to loopback (0.0.0.0 / ::)', () => { + expect(() => assertSafeTokenEndpoint('https://0.0.0.0/token')).toThrow( + InvalidTokenEndpointError, + ); + expect(() => assertSafeTokenEndpoint('https://[::]/token')).toThrow( + InvalidTokenEndpointError, + ); + }); + }); + + describe('in production', () => { + beforeEach(() => { + process.env.NODE_ENV = 'production'; + }); + + it('requires https (rejects http)', () => { + expect(() => assertSafeTokenEndpoint('http://idp.example.com/token')).toThrow( + InvalidTokenEndpointError, + ); + }); + + it('rejects loopback (the executor host itself)', () => { + expect(() => assertSafeTokenEndpoint('https://127.0.0.1/token')).toThrow( + InvalidTokenEndpointError, + ); + expect(() => assertSafeTokenEndpoint('https://localhost/token')).toThrow( + InvalidTokenEndpointError, + ); + expect(() => assertSafeTokenEndpoint('https://[::1]/token')).toThrow( + InvalidTokenEndpointError, + ); + }); + + it('rejects alternate-encoding loopback IPs (integer/hex/octal/short-form -> 127.0.0.1)', () => { + for (const host of ['2130706433', '0x7f000001', '0177.0.0.1', '127.1']) { + expect(() => assertSafeTokenEndpoint(`https://${host}/token`)).toThrow( + InvalidTokenEndpointError, + ); + } + }); + + it('rejects loopback hostname aliases (localhost. and the .localhost TLD)', () => { + expect(() => assertSafeTokenEndpoint('https://localhost./token')).toThrow( + InvalidTokenEndpointError, + ); + expect(() => assertSafeTokenEndpoint('https://foo.localhost/token')).toThrow( + InvalidTokenEndpointError, + ); + }); + + it('does not over-block a real domain that merely contains "localhost"', () => { + expect(() => + assertSafeTokenEndpoint('https://auth.localhost.example.com/token'), + ).not.toThrow(); + }); + + it('rejects trailing-dot IP literals (URL parsing normalizes the dot before classification)', () => { + expect(() => assertSafeTokenEndpoint('https://127.0.0.1./token')).toThrow( + InvalidTokenEndpointError, + ); + expect(() => assertSafeTokenEndpoint('https://169.254.169.254./token')).toThrow( + InvalidTokenEndpointError, + ); + }); + + it('rejects an IPv4-mapped IPv6 loopback (e.g. ::ffff:127.0.0.1)', () => { + expect(() => assertSafeTokenEndpoint('https://[::ffff:127.0.0.1]/token')).toThrow( + InvalidTokenEndpointError, + ); + }); + + it('still allows a private RFC1918 host over https (private providers are supported)', () => { + expect(() => assertSafeTokenEndpoint('https://10.0.0.5/token')).not.toThrow(); + expect(() => assertSafeTokenEndpoint('https://192.168.1.10:8443/token')).not.toThrow(); + }); + + it('accepts a public https endpoint', () => { + expect(() => assertSafeTokenEndpoint('https://idp.example.com/oauth/token')).not.toThrow(); + }); + }); + + describe('fails closed on a non-explicit environment', () => { + it('applies the strict rules when NODE_ENV is unset (not silently relaxed)', () => { + delete process.env.NODE_ENV; + expect(() => assertSafeTokenEndpoint('http://idp.example.com/token')).toThrow( + InvalidTokenEndpointError, + ); + expect(() => assertSafeTokenEndpoint('https://127.0.0.1/token')).toThrow( + InvalidTokenEndpointError, + ); + }); + + it('applies the strict rules when NODE_ENV is an unexpected value (e.g. misspelled)', () => { + process.env.NODE_ENV = 'produciton'; + expect(() => assertSafeTokenEndpoint('http://idp.example.com/token')).toThrow( + InvalidTokenEndpointError, + ); + expect(() => assertSafeTokenEndpoint('https://localhost/token')).toThrow( + InvalidTokenEndpointError, + ); + }); + + it('relaxes only for an explicit development/test opt-in', () => { + process.env.NODE_ENV = 'test'; + expect(() => assertSafeTokenEndpoint('http://127.0.0.1:9101/token')).not.toThrow(); + }); + }); +}); diff --git a/packages/workflow-executor/test/oauth/token-service.test.ts b/packages/workflow-executor/test/oauth/token-service.test.ts new file mode 100644 index 0000000000..948f9232b8 --- /dev/null +++ b/packages/workflow-executor/test/oauth/token-service.test.ts @@ -0,0 +1,464 @@ +import type CredentialEncryption from '../../src/crypto/credential-encryption'; +import type { RefreshGrantParams, RefreshGrantResult } from '../../src/oauth/refresh-grant'; +import type { + McpOAuthCredentialInput, + McpOAuthCredentialsStore, + StoredMcpOAuthCredential, +} from '../../src/ports/mcp-oauth-credentials-store'; + +import { + ExecutorEncryptionKeyMissingError, + OAuthInvalidGrantError, + OAuthReauthRequiredError, + OAuthRefreshError, +} from '../../src/errors'; +import OAuthTokenService from '../../src/oauth/token-service'; + +const USER_ID = 7; +const SERVER_ID = 'srv-1'; + +function makeCredential( + overrides: Partial = {}, +): StoredMcpOAuthCredential { + return { + id: 1, + userId: USER_ID, + mcpServerId: SERVER_ID, + refreshTokenEnc: Buffer.from('enc-rt-1'), + clientId: 'cid', + clientSecretEnc: Buffer.from('enc-secret'), + clientSecretExpiresAt: null, + tokenEndpoint: 'https://idp/token', + tokenEndpointAuthMethod: 'client_secret_basic', + scopes: 'a b', + ...overrides, + }; +} + +function setup(options?: { + credential?: StoredMcpOAuthCredential | null; + refresh?: jest.Mock, [RefreshGrantParams]>; + expirySkewS?: number; +}) { + const credential = options?.credential === undefined ? makeCredential() : options.credential; + const get = jest.fn().mockResolvedValue(credential); + const upsert = jest.fn().mockResolvedValue(undefined); + const updateIfPresent = jest.fn().mockResolvedValue(undefined); + const store = { get, upsert, updateIfPresent } as unknown as McpOAuthCredentialsStore; + + const decrypt = jest.fn((buf: Buffer) => `decrypted:${buf.toString()}`); + const encrypt = jest.fn((plain: string) => ({ + ciphertext: Buffer.from(`enc:${plain}`), + })); + const encryption = { decrypt, encrypt } as unknown as CredentialEncryption; + + const refresh = + options?.refresh ?? jest.fn().mockResolvedValue({ accessToken: 'at-1', expiresInS: 3600 }); + + let clock = 1_000_000; + const now = () => clock; + + const service = new OAuthTokenService({ + store, + encryption, + refreshAccessToken: refresh, + expirySkewS: options?.expirySkewS ?? 60, + now, + }); + + return { + service, + get, + upsert, + updateIfPresent, + decrypt, + encrypt, + refresh, + advance: (ms: number) => { + clock += ms; + }, + }; +} + +describe('OAuthTokenService.getAccessToken', () => { + describe('acquisition', () => { + it('looks up the credential by user and server, refreshes, and returns the access token', async () => { + const { service, get, refresh } = setup(); + + const token = await service.getAccessToken(USER_ID, SERVER_ID); + + expect(token).toBe('at-1'); + expect(get).toHaveBeenCalledWith(USER_ID, SERVER_ID); + expect(refresh).toHaveBeenCalledTimes(1); + }); + + it('decrypts the refresh token and client secret and passes the stored endpoint to the grant', async () => { + const { service, refresh } = setup(); + + await service.getAccessToken(USER_ID, SERVER_ID); + + expect(refresh).toHaveBeenCalledWith({ + tokenEndpoint: 'https://idp/token', + refreshToken: 'decrypted:enc-rt-1', + clientId: 'cid', + clientSecret: 'decrypted:enc-secret', + tokenEndpointAuthMethod: 'client_secret_basic', + scopes: 'a b', + }); + }); + + it('passes a null client secret to the grant when none is stored', async () => { + const { service, refresh } = setup({ credential: makeCredential({ clientSecretEnc: null }) }); + + await service.getAccessToken(USER_ID, SERVER_ID); + + expect(refresh).toHaveBeenCalledWith(expect.objectContaining({ clientSecret: null })); + }); + + it('raises OAuthReauthRequiredError and never refreshes when no credential is stored', async () => { + const { service, refresh } = setup({ credential: null }); + + await expect(service.getAccessToken(USER_ID, SERVER_ID)).rejects.toBeInstanceOf( + OAuthReauthRequiredError, + ); + expect(refresh).not.toHaveBeenCalled(); + }); + + it('propagates a transient refresh failure without converting it to a re-auth pause', async () => { + const refresh = jest.fn().mockRejectedValue(new OAuthRefreshError('5xx')); + const { service } = setup({ refresh }); + + await expect(service.getAccessToken(USER_ID, SERVER_ID)).rejects.toBeInstanceOf( + OAuthRefreshError, + ); + }); + }); + + describe('expiry-skew cache', () => { + it('serves the cached token on a second call within the skew window', async () => { + const { service, refresh } = setup(); + + const first = await service.getAccessToken(USER_ID, SERVER_ID); + const second = await service.getAccessToken(USER_ID, SERVER_ID); + + expect(first).toBe(second); + expect(refresh).toHaveBeenCalledTimes(1); + }); + + it('evict drops the cached token so the next acquire refreshes again', async () => { + const { service, refresh } = setup(); + + await service.getAccessToken(USER_ID, SERVER_ID); + await service.getAccessToken(USER_ID, SERVER_ID); + expect(refresh).toHaveBeenCalledTimes(1); // second call served from cache + + service.evict(USER_ID, SERVER_ID); + await service.getAccessToken(USER_ID, SERVER_ID); + + expect(refresh).toHaveBeenCalledTimes(2); // cache dropped → refreshed again + }); + + it('refreshes again once the token is within the skew window of expiry', async () => { + const { service, refresh, advance } = setup({ expirySkewS: 60 }); + + await service.getAccessToken(USER_ID, SERVER_ID); // expires_in 3600s, skew 60s + advance((3600 - 60) * 1000); // now exactly at the skew threshold + + await service.getAccessToken(USER_ID, SERVER_ID); + + expect(refresh).toHaveBeenCalledTimes(2); + }); + + it('does not cache a token when the grant omits expires_in', async () => { + const refresh = jest.fn().mockResolvedValue({ accessToken: 'at-1' }); + const { service } = setup({ refresh }); + + await service.getAccessToken(USER_ID, SERVER_ID); + await service.getAccessToken(USER_ID, SERVER_ID); + + expect(refresh).toHaveBeenCalledTimes(2); + }); + + it('forceRefresh bypasses a still-valid cached token', async () => { + const { service, refresh } = setup(); + + await service.getAccessToken(USER_ID, SERVER_ID); + await service.getAccessToken(USER_ID, SERVER_ID, { forceRefresh: true }); + + expect(refresh).toHaveBeenCalledTimes(2); + }); + + it('caches independently per (user, server)', async () => { + const { service, refresh } = setup(); + + await service.getAccessToken(USER_ID, SERVER_ID); + await service.getAccessToken(USER_ID, 'other-server'); + + expect(refresh).toHaveBeenCalledTimes(2); + }); + }); + + describe('refresh-token rotation', () => { + it('persists the rotated refresh token, encrypted, when the grant returns a new one', async () => { + const refresh = jest + .fn() + .mockResolvedValue({ accessToken: 'at-1', expiresInS: 3600, refreshToken: 'rt-2' }); + const { service, updateIfPresent, encrypt } = setup({ refresh }); + + await service.getAccessToken(USER_ID, SERVER_ID); + + expect(encrypt).toHaveBeenCalledWith('rt-2'); + expect(updateIfPresent).toHaveBeenCalledWith( + 1, + expect.objectContaining({ + userId: USER_ID, + mcpServerId: SERVER_ID, + refreshTokenEnc: Buffer.from('enc:rt-2'), + }), + ); + }); + + it('does not write back when the grant returns no new refresh token', async () => { + const { service, updateIfPresent } = setup(); + + await service.getAccessToken(USER_ID, SERVER_ID); + + expect(updateIfPresent).not.toHaveBeenCalled(); + }); + + it('still returns the token when the rotation write-back fails', async () => { + const refresh = jest + .fn() + .mockResolvedValue({ accessToken: 'at-1', expiresInS: 3600, refreshToken: 'rt-2' }); + const { service, updateIfPresent } = setup({ refresh }); + (updateIfPresent as jest.Mock).mockRejectedValue(new Error('db down')); + + await expect(service.getAccessToken(USER_ID, SERVER_ID)).resolves.toBe('at-1'); + }); + + it('still returns the token when encrypting the rotated refresh token throws', async () => { + const refresh = jest + .fn() + .mockResolvedValue({ accessToken: 'at-1', expiresInS: 3600, refreshToken: 'rt-2' }); + const { service, encrypt, updateIfPresent } = setup({ refresh }); + (encrypt as jest.Mock).mockImplementation(() => { + throw new Error('key unavailable'); + }); + + await expect(service.getAccessToken(USER_ID, SERVER_ID)).resolves.toBe('at-1'); + expect(updateIfPresent).not.toHaveBeenCalled(); + }); + }); + + describe('invalid_grant — concurrent rotation recovery', () => { + it('re-reads the credential and retries once with the rotated token', async () => { + const rotated = makeCredential({ refreshTokenEnc: Buffer.from('enc-rt-rotated') }); + const get = jest.fn().mockResolvedValueOnce(makeCredential()).mockResolvedValueOnce(rotated); + const refresh = jest + .fn() + .mockRejectedValueOnce(new OAuthInvalidGrantError()) + .mockResolvedValueOnce({ accessToken: 'at-after-retry', expiresInS: 3600 }); + + const service = new OAuthTokenService({ + store: { get, upsert: jest.fn() } as unknown as McpOAuthCredentialsStore, + encryption: { + decrypt: (buf: Buffer) => `decrypted:${buf.toString()}`, + encrypt: jest.fn(), + } as unknown as CredentialEncryption, + refreshAccessToken: refresh, + }); + + const token = await service.getAccessToken(USER_ID, SERVER_ID); + + expect(token).toBe('at-after-retry'); + expect(refresh).toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ refreshToken: 'decrypted:enc-rt-rotated' }), + ); + }); + + it('persists the rotated token onto the re-read credential fields, not the stale ones', async () => { + const original = makeCredential({ scopes: 'a b' }); + const latest = makeCredential({ + refreshTokenEnc: Buffer.from('enc-rt-rotated'), + scopes: 'a b c', + }); + const get = jest.fn().mockResolvedValueOnce(original).mockResolvedValueOnce(latest); + const updateIfPresent = jest.fn().mockResolvedValue(undefined); + const refresh = jest + .fn() + .mockRejectedValueOnce(new OAuthInvalidGrantError()) + .mockResolvedValueOnce({ accessToken: 'at', expiresInS: 3600, refreshToken: 'rt-3' }); + + const service = new OAuthTokenService({ + store: { get, upsert: jest.fn(), updateIfPresent } as unknown as McpOAuthCredentialsStore, + encryption: { + decrypt: (buf: Buffer) => buf.toString(), + encrypt: () => ({ ciphertext: Buffer.from('enc:rt-3') }), + } as unknown as CredentialEncryption, + refreshAccessToken: refresh, + }); + + await service.getAccessToken(USER_ID, SERVER_ID); + + expect(updateIfPresent).toHaveBeenCalledWith(1, expect.objectContaining({ scopes: 'a b c' })); + }); + + it('raises OAuthReauthRequiredError when the re-read shows the same (unrotated) token', async () => { + const refresh = jest.fn().mockRejectedValue(new OAuthInvalidGrantError()); + const { service } = setup({ refresh }); + + await expect(service.getAccessToken(USER_ID, SERVER_ID)).rejects.toBeInstanceOf( + OAuthReauthRequiredError, + ); + expect(refresh).toHaveBeenCalledTimes(1); + }); + + it('raises OAuthReauthRequiredError when the retry also returns invalid_grant', async () => { + const rotated = makeCredential({ refreshTokenEnc: Buffer.from('enc-rt-rotated') }); + const get = jest.fn().mockResolvedValueOnce(makeCredential()).mockResolvedValueOnce(rotated); + const refresh = jest.fn().mockRejectedValue(new OAuthInvalidGrantError()); + + const service = new OAuthTokenService({ + store: { get, upsert: jest.fn() } as unknown as McpOAuthCredentialsStore, + encryption: { + decrypt: (buf: Buffer) => buf.toString(), + encrypt: jest.fn(), + } as unknown as CredentialEncryption, + refreshAccessToken: refresh, + }); + + await expect(service.getAccessToken(USER_ID, SERVER_ID)).rejects.toBeInstanceOf( + OAuthReauthRequiredError, + ); + expect(refresh).toHaveBeenCalledTimes(2); + }); + }); + + describe('serialization', () => { + it('collapses concurrent requests for the same (user, server) into a single refresh', async () => { + let resolveRefresh!: (value: RefreshGrantResult) => void; + const refresh = jest.fn().mockReturnValue( + new Promise(resolve => { + resolveRefresh = resolve; + }), + ); + const { service } = setup({ refresh }); + + const a = service.getAccessToken(USER_ID, SERVER_ID); + const b = service.getAccessToken(USER_ID, SERVER_ID); + resolveRefresh({ accessToken: 'at-shared', expiresInS: 3600 }); + + await expect(a).resolves.toBe('at-shared'); + await expect(b).resolves.toBe('at-shared'); + expect(refresh).toHaveBeenCalledTimes(1); + }); + }); + + describe('decrypt failure classification', () => { + it('surfaces a decrypt failure with the key present as needs-oauth-reauth (recoverable)', async () => { + const service = new OAuthTokenService({ + store: { + get: jest.fn().mockResolvedValue(makeCredential()), + upsert: jest.fn(), + } as unknown as McpOAuthCredentialsStore, + encryption: { + // A since-rotated/hard-swapped key fails GCM auth-tag verification with a generic error. + decrypt: () => { + throw new Error('Unsupported state or unable to authenticate data'); + }, + encrypt: jest.fn(), + } as unknown as CredentialEncryption, + refreshAccessToken: jest.fn(), + }); + + await expect(service.getAccessToken(USER_ID, SERVER_ID)).rejects.toBeInstanceOf( + OAuthReauthRequiredError, + ); + }); + + it('rethrows a missing-key error as terminal, never reaching the grant', async () => { + const refreshAccessToken = jest.fn(); + const service = new OAuthTokenService({ + store: { + get: jest.fn().mockResolvedValue(makeCredential()), + upsert: jest.fn(), + } as unknown as McpOAuthCredentialsStore, + encryption: { + decrypt: () => { + throw new ExecutorEncryptionKeyMissingError(); + }, + encrypt: jest.fn(), + } as unknown as CredentialEncryption, + refreshAccessToken, + }); + + await expect(service.getAccessToken(USER_ID, SERVER_ID)).rejects.toBeInstanceOf( + ExecutorEncryptionKeyMissingError, + ); + expect(refreshAccessToken).not.toHaveBeenCalled(); + }); + }); +}); + +// A disconnect (DELETE) landing between the grant read and the write-back must not re-create the +// row: the atomic update-only path leaves a deleted credential deleted, not silently resurrected. +describe('OAuthTokenService — concurrent disconnect during refresh write-back', () => { + it('does not resurrect a credential deleted between the snapshot read and the rotated-token write-back', async () => { + // GIVEN a stored credential and a refresh that rotates the token while a concurrent disconnect + // (the row is DELETED) lands before the write-back; updateIfPresent updates only an existing row. + let current: StoredMcpOAuthCredential | null = makeCredential(); + const store = { + get: jest.fn(async () => current), + updateIfPresent: jest.fn(async (id: number, credential: McpOAuthCredentialInput) => { + if (current && current.id === id) current = { ...credential, id }; + }), + delete: jest.fn(async () => { + current = null; + }), + } as unknown as McpOAuthCredentialsStore; + + const refresh = jest.fn(async () => { + // The disconnect lands after refreshAndCache's snapshot read, before persistRotatedRefreshToken. + await store.delete(USER_ID, SERVER_ID); + + return { accessToken: 'at-1', expiresInS: 3600, refreshToken: 'rotated-rt' }; + }); + + const encryption = { + decrypt: (buf: Buffer) => buf.toString(), + encrypt: (plain: string) => ({ ciphertext: Buffer.from(`enc:${plain}`) }), + } as unknown as CredentialEncryption; + + const service = new OAuthTokenService({ store, encryption, refreshAccessToken: refresh }); + + // WHEN a token is acquired, triggering the refresh and the rotated-token write-back. + await service.getAccessToken(USER_ID, SERVER_ID); + + // THEN the deleted credential must stay deleted — the write-back must not recreate it. + expect(await store.get(USER_ID, SERVER_ID)).toBeNull(); + }); +}); + +// A refresh already in flight when evict() fires (disconnect) must not repopulate the cache +// afterwards, or later calls keep serving a token for a disconnected credential. +describe('OAuthTokenService — evict during in-flight refresh', () => { + it('does not cache the freshly-minted token when a disconnect evicts mid-refresh', async () => { + const refresh = jest.fn().mockResolvedValue({ accessToken: 'at-2', expiresInS: 3600 }); + const { service } = setup({ refresh }); + // The disconnect (evict) lands during the first refresh — after the snapshot read, before caching. + refresh.mockImplementationOnce(async () => { + service.evict(USER_ID, SERVER_ID); + + return { accessToken: 'at-1', expiresInS: 3600 }; + }); + + const first = await service.getAccessToken(USER_ID, SERVER_ID); + const second = await service.getAccessToken(USER_ID, SERVER_ID); + + // The in-flight caller still gets its token, but it was not cached — the next call refreshes anew. + expect(first).toBe('at-1'); + expect(second).toBe('at-2'); + expect(refresh).toHaveBeenCalledTimes(2); + }); +}); diff --git a/packages/workflow-executor/test/remote-tool-fetcher.test.ts b/packages/workflow-executor/test/remote-tool-fetcher.test.ts index 7527816adc..2da8bfe4a9 100644 --- a/packages/workflow-executor/test/remote-tool-fetcher.test.ts +++ b/packages/workflow-executor/test/remote-tool-fetcher.test.ts @@ -1,16 +1,25 @@ +import type OAuthTokenService from '../src/oauth/token-service'; import type { AiModelPort } from '../src/ports/ai-model-port'; import type { Logger } from '../src/ports/logger-port'; import type { WorkflowPort } from '../src/ports/workflow-port'; import type { RemoteTool, ToolConfig } from '@forestadmin/ai-proxy'; +import { OAuthReauthRequiredError } from '../src/errors'; import RemoteToolFetcher, { scopeConfigsToServer } from '../src/remote-tool-fetcher'; +const USER_ID = 1; + +type AiModelMethods = Pick; + function createMockWorkflowPort(): jest.Mocked> { return { getMcpServerConfigs: jest.fn().mockResolvedValue({}) }; } -function createMockAiModelPort(): jest.Mocked> { - return { loadRemoteTools: jest.fn().mockResolvedValue([]) }; +function createMockAiModelPort(): jest.Mocked { + return { + loadRemoteTools: jest.fn().mockResolvedValue([]), + loadRemoteToolsWithFailures: jest.fn().mockResolvedValue({ tools: [], failures: [] }), + }; } function createMockLogger(): jest.MockedFunction { @@ -23,16 +32,23 @@ function makeRemoteTool(sourceId: string, mcpServerId?: string): RemoteTool { function makeFetcher(overrides?: { workflowPort?: Partial>>; - aiModelPort?: Partial>>; + aiModelPort?: Partial>; logger?: jest.MockedFunction; + tokenService?: OAuthTokenService; }) { const workflowPort = { ...createMockWorkflowPort(), ...overrides?.workflowPort }; const aiModelPort = { ...createMockAiModelPort(), ...overrides?.aiModelPort }; const logger = overrides?.logger ?? createMockLogger(); + const tokenService = + overrides?.tokenService ?? + ({ + getAccessToken: jest.fn().mockResolvedValue('access-token'), + } as unknown as OAuthTokenService); const fetcher = new RemoteToolFetcher( workflowPort as unknown as WorkflowPort, aiModelPort as unknown as AiModelPort, logger, + tokenService, ); return { fetcher, workflowPort, aiModelPort, logger }; @@ -41,6 +57,15 @@ function makeFetcher(overrides?: { const cfg = (id: string | undefined): ToolConfig => ({ id, url: 'https://x.example', type: 'http' as const, headers: {} } as unknown as ToolConfig); +const oauthCfg = (id: string): ToolConfig => + ({ + id, + authType: 'oauth2', + url: 'https://x.example', + type: 'http' as const, + headers: {}, + } as unknown as ToolConfig); + // --------------------------------------------------------------------------- // scopeConfigsToServer (pure) // --------------------------------------------------------------------------- @@ -68,6 +93,24 @@ describe('scopeConfigsToServer', () => { }); }); +// --------------------------------------------------------------------------- +// RemoteToolFetcher construction +// --------------------------------------------------------------------------- + +describe('RemoteToolFetcher construction', () => { + it('throws when constructed without an OAuth token service', () => { + expect( + () => + new RemoteToolFetcher( + createMockWorkflowPort() as unknown as WorkflowPort, + createMockAiModelPort() as unknown as AiModelPort, + createMockLogger(), + undefined as unknown as OAuthTokenService, + ), + ).toThrow('requires an OAuth token service'); + }); +}); + // --------------------------------------------------------------------------- // RemoteToolFetcher.fetch // --------------------------------------------------------------------------- @@ -82,7 +125,7 @@ describe('RemoteToolFetcher.fetch', () => { }, }); - await fetcher.fetch('id-A'); + await fetcher.fetch('id-A', USER_ID); expect(aiModelPort.loadRemoteTools).toHaveBeenCalledWith({ 'srv-a': cfg('id-A') }); }); @@ -92,7 +135,7 @@ describe('RemoteToolFetcher.fetch', () => { workflowPort: { getMcpServerConfigs: jest.fn().mockResolvedValue({}) }, }); - const result = await fetcher.fetch('id-A'); + const result = await fetcher.fetch('id-A', USER_ID); expect(result).toEqual({ tools: [], mcpServerName: undefined }); expect(aiModelPort.loadRemoteTools).not.toHaveBeenCalled(); @@ -107,9 +150,9 @@ describe('RemoteToolFetcher.fetch', () => { aiModelPort: { loadRemoteTools: jest.fn().mockResolvedValue(remoteTools) }, }); - const result = await fetcher.fetch('id-A'); + const result = await fetcher.fetch('id-A', USER_ID); - expect(result).toEqual({ tools: remoteTools, mcpServerName: 'srv-a' }); + expect(result).toEqual({ tools: remoteTools, mcpServerName: 'srv-a', loadFailed: false }); }); it('warns about the missing target with the list of advertised ids when no config matches', async () => { @@ -121,7 +164,7 @@ describe('RemoteToolFetcher.fetch', () => { }, }); - await fetcher.fetch('id-missing'); + await fetcher.fetch('id-missing', USER_ID); expect(logger).toHaveBeenCalledWith( 'Warn', @@ -139,7 +182,7 @@ describe('RemoteToolFetcher.fetch', () => { workflowPort: { getMcpServerConfigs: jest.fn().mockResolvedValue({}) }, }); - await fetcher.fetch('id-A'); + await fetcher.fetch('id-A', USER_ID); expect(logger).toHaveBeenCalledWith( 'Warn', @@ -160,7 +203,7 @@ describe('RemoteToolFetcher.fetch', () => { }, }); - await fetcher.fetch('id-A'); + await fetcher.fetch('id-A', USER_ID); expect(logger.mock.calls.find(c => c[0] === 'Warn')).toBeUndefined(); }); @@ -173,7 +216,7 @@ describe('RemoteToolFetcher.fetch', () => { aiModelPort: { loadRemoteTools: jest.fn().mockResolvedValue([]) }, }); - await fetcher.fetch('id-A'); + await fetcher.fetch('id-A', USER_ID); expect(logger).toHaveBeenCalledWith('Error', 'MCP servers failed to load tools', { requestedMcpServerId: 'id-A', @@ -182,6 +225,30 @@ describe('RemoteToolFetcher.fetch', () => { }); }); + it('sets loadFailed when the scoped server produced no tools', async () => { + const { fetcher } = makeFetcher({ + workflowPort: { + getMcpServerConfigs: jest.fn().mockResolvedValue({ 'srv-a': cfg('id-A') }), + }, + aiModelPort: { loadRemoteTools: jest.fn().mockResolvedValue([]) }, + }); + + expect((await fetcher.fetch('id-A', USER_ID)).loadFailed).toBe(true); + }); + + it('does not set loadFailed when tools load successfully', async () => { + const { fetcher } = makeFetcher({ + workflowPort: { + getMcpServerConfigs: jest.fn().mockResolvedValue({ 'srv-a': cfg('id-A') }), + }, + aiModelPort: { + loadRemoteTools: jest.fn().mockResolvedValue([makeRemoteTool('srv-a', 'id-A')]), + }, + }); + + expect((await fetcher.fetch('id-A', USER_ID)).loadFailed).toBeFalsy(); + }); + it('does not log a partial-failure error when a tool carries the scoped config id', async () => { const { fetcher, logger } = makeFetcher({ workflowPort: { @@ -192,7 +259,7 @@ describe('RemoteToolFetcher.fetch', () => { }, }); - await fetcher.fetch('id-A'); + await fetcher.fetch('id-A', USER_ID); expect(logger.mock.calls.find(c => c[0] === 'Error')).toBeUndefined(); }); @@ -214,7 +281,7 @@ describe('RemoteToolFetcher.fetch', () => { }, }); - await fetcher.fetch('id-zendesk'); + await fetcher.fetch('id-zendesk', USER_ID); expect(logger.mock.calls.find(c => c[0] === 'Error')).toBeUndefined(); }); @@ -232,7 +299,7 @@ describe('RemoteToolFetcher.fetch', () => { aiModelPort: { loadRemoteTools: jest.fn().mockResolvedValue([]) }, }); - await fetcher.fetch('id-zendesk'); + await fetcher.fetch('id-zendesk', USER_ID); expect(logger).toHaveBeenCalledWith('Error', 'MCP servers failed to load tools', { requestedMcpServerId: 'id-zendesk', @@ -250,7 +317,7 @@ describe('RemoteToolFetcher.fetch', () => { aiModelPort: { loadRemoteTools: jest.fn().mockResolvedValue(remoteTools) }, }); - const result = await fetcher.fetch('id-A'); + const result = await fetcher.fetch('id-A', USER_ID); expect(result.tools).toBe(remoteTools); }); @@ -265,7 +332,7 @@ describe('RemoteToolFetcher.fetch', () => { }, }); - await expect(fetcher.fetch('id-A')).rejects.toThrow('MCP unreachable'); + await expect(fetcher.fetch('id-A', USER_ID)).rejects.toThrow('MCP unreachable'); expect(logger.mock.calls.find(c => c[0] === 'Error')).toBeUndefined(); }); @@ -276,7 +343,153 @@ describe('RemoteToolFetcher.fetch', () => { }, }); - await expect(fetcher.fetch('id-A')).rejects.toThrow('orchestrator down'); + await expect(fetcher.fetch('id-A', USER_ID)).rejects.toThrow('orchestrator down'); expect(aiModelPort.loadRemoteTools).not.toHaveBeenCalled(); }); }); + +describe('RemoteToolFetcher.fetch — OAuth2 servers', () => { + const makeTokenService = (getAccessToken: jest.Mock): OAuthTokenService => + ({ getAccessToken } as unknown as OAuthTokenService); + + const authFailure = { + server: 'srv-a', + mcpServerId: 'id-A', + kind: 'auth', + error: new Error('401'), + }; + + it('acquires a token, injects it as a Bearer header, and returns the loaded tools + a reload hook', async () => { + const tool = makeRemoteTool('srv-a', 'id-A'); + const getAccessToken = jest.fn().mockResolvedValue('tok-1'); + const loadRemoteToolsWithFailures = jest + .fn() + .mockResolvedValue({ tools: [tool], failures: [] }); + const { fetcher } = makeFetcher({ + workflowPort: { + getMcpServerConfigs: jest.fn().mockResolvedValue({ 'srv-a': oauthCfg('id-A') }), + }, + aiModelPort: { loadRemoteToolsWithFailures }, + tokenService: makeTokenService(getAccessToken), + }); + + const result = await fetcher.fetch('id-A', USER_ID); + + expect(getAccessToken).toHaveBeenCalledWith(USER_ID, 'id-A', { forceRefresh: false }); + expect(loadRemoteToolsWithFailures).toHaveBeenCalledWith({ + 'srv-a': expect.objectContaining({ headers: { Authorization: 'Bearer tok-1' } }), + }); + expect(result.tools).toEqual([tool]); + expect(typeof result.reloadWithFreshAuth).toBe('function'); + }); + + it('injects the Bearer token for every scoped config that shares the mcpServerId', async () => { + const getAccessToken = jest.fn().mockResolvedValue('tok-1'); + const loadRemoteToolsWithFailures = jest.fn().mockResolvedValue({ tools: [], failures: [] }); + const { fetcher } = makeFetcher({ + workflowPort: { + getMcpServerConfigs: jest + .fn() + .mockResolvedValue({ 'srv-a': oauthCfg('id-A'), 'srv-a-dup': oauthCfg('id-A') }), + }, + aiModelPort: { loadRemoteToolsWithFailures }, + tokenService: makeTokenService(getAccessToken), + }); + + await fetcher.fetch('id-A', USER_ID); + + expect(loadRemoteToolsWithFailures).toHaveBeenCalledWith({ + 'srv-a': expect.objectContaining({ headers: { Authorization: 'Bearer tok-1' } }), + 'srv-a-dup': expect.objectContaining({ headers: { Authorization: 'Bearer tok-1' } }), + }); + }); + + it('force-refreshes and retries list-tools once on an auth failure, then succeeds', async () => { + const tool = makeRemoteTool('srv-a', 'id-A'); + const getAccessToken = jest.fn().mockResolvedValue('tok'); + const loadRemoteToolsWithFailures = jest + .fn() + .mockResolvedValueOnce({ tools: [], failures: [authFailure] }) + .mockResolvedValueOnce({ tools: [tool], failures: [] }); + const { fetcher } = makeFetcher({ + workflowPort: { + getMcpServerConfigs: jest.fn().mockResolvedValue({ 'srv-a': oauthCfg('id-A') }), + }, + aiModelPort: { loadRemoteToolsWithFailures }, + tokenService: makeTokenService(getAccessToken), + }); + + const result = await fetcher.fetch('id-A', USER_ID); + + expect(getAccessToken).toHaveBeenNthCalledWith(1, USER_ID, 'id-A', { forceRefresh: false }); + expect(getAccessToken).toHaveBeenNthCalledWith(2, USER_ID, 'id-A', { forceRefresh: true }); + expect(result.tools).toEqual([tool]); + }); + + it('raises OAuthReauthRequiredError when the auth failure persists after a forced refresh', async () => { + const getAccessToken = jest.fn().mockResolvedValue('tok'); + const loadRemoteToolsWithFailures = jest + .fn() + .mockResolvedValue({ tools: [], failures: [authFailure] }); + const { fetcher } = makeFetcher({ + workflowPort: { + getMcpServerConfigs: jest.fn().mockResolvedValue({ 'srv-a': oauthCfg('id-A') }), + }, + aiModelPort: { loadRemoteToolsWithFailures }, + tokenService: makeTokenService(getAccessToken), + }); + + await expect(fetcher.fetch('id-A', USER_ID)).rejects.toBeInstanceOf(OAuthReauthRequiredError); + }); + + it('propagates OAuthReauthRequiredError from token acquisition (no stored credential)', async () => { + const getAccessToken = jest.fn().mockRejectedValue(new OAuthReauthRequiredError('id-A')); + const { fetcher } = makeFetcher({ + workflowPort: { + getMcpServerConfigs: jest.fn().mockResolvedValue({ 'srv-a': oauthCfg('id-A') }), + }, + tokenService: makeTokenService(getAccessToken), + }); + + await expect(fetcher.fetch('id-A', USER_ID)).rejects.toBeInstanceOf(OAuthReauthRequiredError); + }); + + it('reloadWithFreshAuth forces a refresh and returns freshly-authed tools', async () => { + const tool = makeRemoteTool('srv-a', 'id-A'); + const getAccessToken = jest.fn().mockResolvedValue('tok'); + const loadRemoteToolsWithFailures = jest + .fn() + .mockResolvedValue({ tools: [tool], failures: [] }); + const { fetcher } = makeFetcher({ + workflowPort: { + getMcpServerConfigs: jest.fn().mockResolvedValue({ 'srv-a': oauthCfg('id-A') }), + }, + aiModelPort: { loadRemoteToolsWithFailures }, + tokenService: makeTokenService(getAccessToken), + }); + + const { reloadWithFreshAuth } = await fetcher.fetch('id-A', USER_ID); + if (!reloadWithFreshAuth) throw new Error('expected reloadWithFreshAuth to be defined'); + getAccessToken.mockClear(); + const reloaded = await reloadWithFreshAuth(); + + expect(getAccessToken).toHaveBeenCalledWith(USER_ID, 'id-A', { forceRefresh: true }); + expect(reloaded).toEqual([tool]); + }); + + it('leaves the token service untouched for a bearer/none server', async () => { + const getAccessToken = jest.fn(); + const tool = makeRemoteTool('srv-a', 'id-A'); + const { fetcher, aiModelPort } = makeFetcher({ + workflowPort: { getMcpServerConfigs: jest.fn().mockResolvedValue({ 'srv-a': cfg('id-A') }) }, + aiModelPort: { loadRemoteTools: jest.fn().mockResolvedValue([tool]) }, + tokenService: makeTokenService(getAccessToken), + }); + + const result = await fetcher.fetch('id-A', USER_ID); + + expect(getAccessToken).not.toHaveBeenCalled(); + expect(aiModelPort.loadRemoteTools).toHaveBeenCalled(); + expect(result.reloadWithFreshAuth).toBeUndefined(); + }); +}); diff --git a/packages/workflow-executor/test/runner.test.ts b/packages/workflow-executor/test/runner.test.ts index d18cccef6d..383ff2e0d0 100644 --- a/packages/workflow-executor/test/runner.test.ts +++ b/packages/workflow-executor/test/runner.test.ts @@ -1,4 +1,5 @@ import type { StepContextConfig } from '../src/executors/step-executor-factory'; +import type OAuthTokenService from '../src/oauth/token-service'; import type { AgentPort } from '../src/ports/agent-port'; import type { AiModelPort } from '../src/ports/ai-model-port'; import type { Logger } from '../src/ports/logger-port'; @@ -79,6 +80,7 @@ function createMockRunStore(overrides: Partial = {}): jest.Mocked; } @@ -105,6 +107,7 @@ function createRunnerConfig( close: jest.fn().mockResolvedValue(undefined), getStepExecutions: jest.fn().mockResolvedValue([]), saveStepExecution: jest.fn().mockResolvedValue(undefined), + deleteStepExecution: jest.fn().mockResolvedValue(undefined), } as unknown as RunStore, pollingIntervalS: POLLING_INTERVAL_S, aiModelPort: createMockAiClient() as unknown as AiModelPort, @@ -120,6 +123,9 @@ function createRunnerConfig( schemaCache: new SchemaCache(), envSecret: VALID_ENV_SECRET, authSecret: VALID_AUTH_SECRET, + mcpOAuthTokenService: { + getAccessToken: jest.fn().mockResolvedValue('access-token'), + } as unknown as OAuthTokenService, ...overrides, }; } @@ -1713,7 +1719,7 @@ describe('StepExecutorFactory.create — factory', () => { fetchRemoteTools, ); expect(executor).toBeInstanceOf(McpStepExecutor); - expect(fetchRemoteTools).toHaveBeenCalledWith('srv-42'); + expect(fetchRemoteTools).toHaveBeenCalledWith('srv-42', 1); expect( ( executor as unknown as { getExtraLogContext(): Record } diff --git a/packages/workflow-executor/test/stores/database-mcp-oauth-credentials-store.pg.test.ts b/packages/workflow-executor/test/stores/database-mcp-oauth-credentials-store.pg.test.ts new file mode 100644 index 0000000000..d7c3d72068 --- /dev/null +++ b/packages/workflow-executor/test/stores/database-mcp-oauth-credentials-store.pg.test.ts @@ -0,0 +1,114 @@ +import { Sequelize } from 'sequelize'; + +import DatabaseMcpOAuthCredentialsStore from '../../src/stores/database-mcp-oauth-credentials-store'; + +// Real-Postgres integration test. Skipped unless WORKFLOW_EXECUTOR_TEST_DATABASE_URL points at a +// reachable Postgres (CI has none), e.g. +// WORKFLOW_EXECUTOR_TEST_DATABASE_URL=postgres://forest:secret@localhost:5435/forest +// Guards the shared-database safety the SQLite suite can't exercise: the table + its umzug registry +// live under the `forest` schema (never `public`), and the advisory lock serializes concurrent boots. +const PG_URL = process.env.WORKFLOW_EXECUTOR_TEST_DATABASE_URL; +const describePg = PG_URL ? describe : describe.skip; + +const SCHEMA = 'forest'; +const TABLE = 'ai_mcp_oauth_credentials'; + +const makeSequelize = () => new Sequelize(PG_URL as string, { logging: false }); + +const count = async (sequelize: Sequelize, sql: string): Promise => { + const [rows] = await sequelize.query(sql); + + return (rows as Array<{ n: number }>)[0].n; +}; + +describePg('DatabaseMcpOAuthCredentialsStore — Postgres shared-schema integration', () => { + let admin: Sequelize; + + beforeEach(async () => { + admin = makeSequelize(); + await admin.query(`DROP SCHEMA IF EXISTS "${SCHEMA}" CASCADE`); + }); + + afterEach(async () => { + await admin.query(`DROP SCHEMA IF EXISTS "${SCHEMA}" CASCADE`); + await admin.close(); + }); + + it('creates the table and its migration registry under the forest schema, never public', async () => { + const store = new DatabaseMcpOAuthCredentialsStore({ sequelize: makeSequelize() }); + + try { + await store.init(); + + expect( + await count( + admin, + `SELECT count(*)::int AS n FROM information_schema.tables + WHERE table_schema = '${SCHEMA}' AND table_name = '${TABLE}'`, + ), + ).toBe(1); + // Never leaks into public — that is the shared-database safety guarantee. + expect( + await count( + admin, + `SELECT count(*)::int AS n FROM information_schema.tables + WHERE table_schema = 'public' AND table_name = '${TABLE}'`, + ), + ).toBe(0); + // The umzug registry lives in forest too, with 002 recorded. + expect(await count(admin, `SELECT count(*)::int AS n FROM "${SCHEMA}"."SequelizeMeta"`)).toBe( + 1, + ); + } finally { + await store.close(); + } + }); + + it('round-trips an encrypted credential with the BYTEA blob preserved byte-for-byte', async () => { + const store = new DatabaseMcpOAuthCredentialsStore({ sequelize: makeSequelize() }); + + try { + await store.init(); + const refreshTokenEnc = Buffer.from([0x00, 0x01, 0xfe, 0xff, 0x10]); + + await store.upsert({ + userId: 7, + mcpServerId: 'mcp-server-1', + refreshTokenEnc, + clientId: null, + clientSecretEnc: null, + clientSecretExpiresAt: null, + tokenEndpoint: 'https://auth.example.com/token', + tokenEndpointAuthMethod: null, + scopes: null, + }); + + const row = await store.get(7, 'mcp-server-1'); + expect(row?.refreshTokenEnc.toString('hex')).toBe(refreshTokenEnc.toString('hex')); + } finally { + await store.close(); + } + }); + + it('lets N instances init() concurrently on an empty schema without crashing, migrating once', async () => { + const stores = Array.from( + { length: 5 }, + () => new DatabaseMcpOAuthCredentialsStore({ sequelize: makeSequelize() }), + ); + + try { + const results = await Promise.allSettled(stores.map(store => store.init())); + + expect(results.filter(result => result.status === 'rejected')).toEqual([]); + expect( + await count( + admin, + `SELECT count(*)::int AS n FROM information_schema.tables + WHERE table_schema = '${SCHEMA}' AND table_name = '${TABLE}'`, + ), + ).toBe(1); + } finally { + await Promise.all(stores.map(store => store.close())); + } + }); +}); diff --git a/packages/workflow-executor/test/stores/database-mcp-oauth-credentials-store.test.ts b/packages/workflow-executor/test/stores/database-mcp-oauth-credentials-store.test.ts new file mode 100644 index 0000000000..d428a01da8 --- /dev/null +++ b/packages/workflow-executor/test/stores/database-mcp-oauth-credentials-store.test.ts @@ -0,0 +1,338 @@ +/** + * Spec for the database (Sequelize) MCP OAuth credentials store + its Umzug migration. + * + * Behaviour: + * - One row per (user_id, mcp_server_id) — UNIQUE (user_id, mcp_server_id); upsert in place. + * - Refresh token + client secret are stored as encrypted BLOBs; the store persists opaque bytes + * (encryption itself is exercised in credential-encryption.test.ts — the store does not encrypt). + * - client_id, client_secret_enc, client_secret_expires_at, scopes are nullable + * (null for public / PKCE clients). + * - Deleted on disconnect / permanent refresh failure. + * - Migration `002_create_mcp_oauth_credentials` is added alongside `001_create_workflow_step_executions`. + * + * Store contract: + * import DatabaseMcpOAuthCredentialsStore from '../../src/stores/database-mcp-oauth-credentials-store'; + * const store = new DatabaseMcpOAuthCredentialsStore({ sequelize }); + * await store.init(); // runs the 002 migration (table exists after) + * await store.upsert(credential); // keyed by (userId, mcpServerId) + * const row = await store.get(userId, mcpServerId); // StoredCredential | null + * await store.delete(userId, mcpServerId); + * await store.close(); + * + * Field names are camelCase, mapping to the snake_case columns. + */ +import type { Sequelize as SequelizeType } from 'sequelize'; + +import { Sequelize } from 'sequelize'; + +import DatabaseMcpOAuthCredentialsStore from '../../src/stores/database-mcp-oauth-credentials-store'; + +interface CredentialInput { + userId: number; + mcpServerId: string; + refreshTokenEnc: Buffer; + clientId?: string | null; + clientSecretEnc?: Buffer | null; + clientSecretExpiresAt?: Date | null; + tokenEndpoint: string; + tokenEndpointAuthMethod?: string | null; + scopes?: string | null; +} + +function makeCredential(overrides: Partial = {}): CredentialInput { + return { + userId: 42, + mcpServerId: 'mcp-server-1', + refreshTokenEnc: Buffer.from('enc-refresh-token'), + clientId: 'client-abc', + clientSecretEnc: Buffer.from('enc-client-secret'), + clientSecretExpiresAt: null, + tokenEndpoint: 'https://auth.example.com/token', + tokenEndpointAuthMethod: 'client_secret_post', + scopes: 'read write', + ...overrides, + }; +} + +// Asserts presence and narrows the type — avoids non-null assertions (`!`), which the codebase avoids. +function unwrap(value: T | null | undefined): T { + if (value === null || value === undefined) { + throw new Error('expected a stored credential, got null/undefined'); + } + + return value; +} + +describe('DatabaseMcpOAuthCredentialsStore (SQLite)', () => { + let sequelize: SequelizeType; + let store: DatabaseMcpOAuthCredentialsStore; + + beforeEach(async () => { + sequelize = new Sequelize({ dialect: 'sqlite', storage: ':memory:', logging: false }); + store = new DatabaseMcpOAuthCredentialsStore({ sequelize }); + await store.init(); + }); + + afterEach(async () => { + await store.close(); + }); + + describe('get', () => { + it('returns null for an unknown (userId, mcpServerId)', async () => { + expect(await store.get(999, 'no-such-server')).toBeNull(); + }); + + it('returns the stored credential for a known (userId, mcpServerId)', async () => { + const credential = makeCredential(); + + await store.upsert(credential); + const row = await store.get(credential.userId, credential.mcpServerId); + + expect(row).toEqual( + expect.objectContaining({ + userId: 42, + mcpServerId: 'mcp-server-1', + clientId: 'client-abc', + tokenEndpoint: 'https://auth.example.com/token', + tokenEndpointAuthMethod: 'client_secret_post', + scopes: 'read write', + }), + ); + }); + + it('preserves the encrypted blobs byte-for-byte', async () => { + const refreshTokenEnc = Buffer.from([0x00, 0x01, 0xfe, 0xff, 0x10]); + const clientSecretEnc = Buffer.from([0xde, 0xad, 0xbe, 0xef]); + + await store.upsert(makeCredential({ refreshTokenEnc, clientSecretEnc })); + const row = unwrap(await store.get(42, 'mcp-server-1')); + + expect(row.refreshTokenEnc.toString('hex')).toBe(refreshTokenEnc.toString('hex')); + expect(unwrap(row.clientSecretEnc).toString('hex')).toBe(clientSecretEnc.toString('hex')); + }); + }); + + describe('upsert', () => { + it('updates the existing row in place for the same (userId, mcpServerId)', async () => { + await store.upsert(makeCredential({ refreshTokenEnc: Buffer.from('old') })); + await store.upsert(makeCredential({ refreshTokenEnc: Buffer.from('new') })); + + const row = unwrap(await store.get(42, 'mcp-server-1')); + + expect(row.refreshTokenEnc.toString()).toBe('new'); + }); + + it('keeps exactly one row after re-upserting the same key (UNIQUE constraint)', async () => { + await store.upsert(makeCredential({ refreshTokenEnc: Buffer.from('v1') })); + await store.upsert(makeCredential({ refreshTokenEnc: Buffer.from('v2') })); + + // Pollution-proof: count only rows for this known key, never the whole table. + const [rows] = await sequelize.query( + 'SELECT COUNT(*) AS c FROM ai_mcp_oauth_credentials WHERE user_id = 42 AND mcp_server_id = :id', + { replacements: { id: 'mcp-server-1' } }, + ); + expect(Number((rows[0] as { c: number }).c)).toBe(1); + }); + + it('stores nullable client fields as null for a public / PKCE client', async () => { + await store.upsert( + makeCredential({ + clientId: null, + clientSecretEnc: null, + clientSecretExpiresAt: null, + tokenEndpointAuthMethod: 'none', + scopes: null, + }), + ); + + const row = await store.get(42, 'mcp-server-1'); + + expect(row).toEqual( + expect.objectContaining({ + clientId: null, + clientSecretEnc: null, + clientSecretExpiresAt: null, + scopes: null, + }), + ); + }); + + it('persists client_secret_expires_at when provided', async () => { + const expiresAt = new Date('2030-01-02T03:04:05.000Z'); + + await store.upsert(makeCredential({ clientSecretExpiresAt: expiresAt })); + const row = unwrap(await store.get(42, 'mcp-server-1')); + + expect(new Date(unwrap(row.clientSecretExpiresAt)).toISOString()).toBe( + expiresAt.toISOString(), + ); + }); + }); + + describe('updateIfPresent', () => { + it('updates the row matching the given id in place', async () => { + await store.upsert(makeCredential({ refreshTokenEnc: Buffer.from('old') })); + const { id } = unwrap(await store.get(42, 'mcp-server-1')); + + await store.updateIfPresent(id, makeCredential({ refreshTokenEnc: Buffer.from('rotated') })); + + const row = unwrap(await store.get(42, 'mcp-server-1')); + expect(row.refreshTokenEnc.toString()).toBe('rotated'); + }); + + it('does not insert a row when none exists for the key', async () => { + await store.updateIfPresent(1, makeCredential()); + + expect(await store.get(42, 'mcp-server-1')).toBeNull(); + }); + + it('does not touch a row that was re-created with a different id', async () => { + await store.upsert(makeCredential({ refreshTokenEnc: Buffer.from('old') })); + const staleId = unwrap(await store.get(42, 'mcp-server-1')).id; + await store.delete(42, 'mcp-server-1'); + await store.upsert(makeCredential({ refreshTokenEnc: Buffer.from('reauthorized') })); + + await store.updateIfPresent(staleId, makeCredential({ refreshTokenEnc: Buffer.from('rot') })); + + const row = unwrap(await store.get(42, 'mcp-server-1')); + expect(row.refreshTokenEnc.toString()).toBe('reauthorized'); + }); + }); + + describe('isolation', () => { + it('keeps credentials for the same server but different users separate', async () => { + await store.upsert(makeCredential({ userId: 1, refreshTokenEnc: Buffer.from('user-1') })); + await store.upsert(makeCredential({ userId: 2, refreshTokenEnc: Buffer.from('user-2') })); + + const rowOne = unwrap(await store.get(1, 'mcp-server-1')); + const rowTwo = unwrap(await store.get(2, 'mcp-server-1')); + + expect(rowOne.refreshTokenEnc.toString()).toBe('user-1'); + expect(rowTwo.refreshTokenEnc.toString()).toBe('user-2'); + }); + + it('keeps credentials for the same user but different servers separate', async () => { + await store.upsert( + makeCredential({ mcpServerId: 'server-a', refreshTokenEnc: Buffer.from('a') }), + ); + await store.upsert( + makeCredential({ mcpServerId: 'server-b', refreshTokenEnc: Buffer.from('b') }), + ); + + const rowA = unwrap(await store.get(42, 'server-a')); + const rowB = unwrap(await store.get(42, 'server-b')); + + expect(rowA.refreshTokenEnc.toString()).toBe('a'); + expect(rowB.refreshTokenEnc.toString()).toBe('b'); + }); + }); + + describe('delete', () => { + it('removes the credential for a (userId, mcpServerId)', async () => { + await store.upsert(makeCredential()); + + await store.delete(42, 'mcp-server-1'); + + expect(await store.get(42, 'mcp-server-1')).toBeNull(); + }); + + it('does not affect other users when deleting one user', async () => { + await store.upsert(makeCredential({ userId: 1 })); + await store.upsert(makeCredential({ userId: 2 })); + + await store.delete(1, 'mcp-server-1'); + + expect(await store.get(1, 'mcp-server-1')).toBeNull(); + expect(await store.get(2, 'mcp-server-1')).not.toBeNull(); + }); + + it('is a no-op (does not throw) when deleting a non-existent credential', async () => { + await expect(store.delete(999, 'no-such-server')).resolves.toBeUndefined(); + }); + }); + + describe('migration / init', () => { + it('creates the ai_mcp_oauth_credentials table on init', async () => { + const [rows] = await sequelize.query( + "SELECT name FROM sqlite_master WHERE type='table' AND name='ai_mcp_oauth_credentials'", + ); + + expect(rows).toHaveLength(1); + }); + + it('runs init idempotently', async () => { + await expect(store.init()).resolves.toBeUndefined(); + }); + + it('rejects an insert with a null token_endpoint at the DB level', async () => { + // token_endpoint is NOT NULL: the refresh grant has nowhere to go without it. + await expect( + sequelize.query( + 'INSERT INTO ai_mcp_oauth_credentials ' + + '(user_id, mcp_server_id, refresh_token_enc, created_at, updated_at) ' + + "VALUES (7, 'mcp-server-1', :blob, :now, :now)", + { replacements: { blob: Buffer.from('no-endpoint'), now: new Date() } }, + ), + ).rejects.toThrow(); + }); + + it('enforces the UNIQUE (user_id, mcp_server_id) constraint at the DB level', async () => { + // Direct insert bypassing upsert proves the constraint exists in the schema, not just app logic. + await store.upsert(makeCredential()); + + await expect( + sequelize.query( + 'INSERT INTO ai_mcp_oauth_credentials ' + + '(user_id, mcp_server_id, refresh_token_enc, token_endpoint, ' + + 'created_at, updated_at) ' + + "VALUES (42, 'mcp-server-1', :blob, :tokenEndpoint, :now, :now)", + { + replacements: { + blob: Buffer.from('dup'), + tokenEndpoint: 'https://auth.example.com/token', + now: new Date(), + }, + }, + ), + ).rejects.toThrow(); + }); + }); + + describe('failure handling', () => { + it('logs and rethrows when the migration fails', async () => { + const badSequelize = new Sequelize({ + dialect: 'sqlite', + storage: ':memory:', + logging: false, + }); + const badStore = new DatabaseMcpOAuthCredentialsStore({ sequelize: badSequelize }); + + // Break the query interface so createTable fails mid-migration. + jest + .spyOn(badSequelize.getQueryInterface(), 'createTable') + .mockRejectedValueOnce(new Error('disk full')); + + const logger = jest.fn(); + await expect(badStore.init(logger)).rejects.toThrow('disk full'); + expect(logger).toHaveBeenCalledWith( + 'Error', + 'MCP OAuth credentials migration failed', + expect.objectContaining({ error: 'disk full' }), + ); + + await badSequelize.close(); + }); + + it('close() catches and logs the error instead of throwing', async () => { + const logger = jest.fn(); + jest.spyOn(sequelize, 'close').mockRejectedValueOnce(new Error('close failed')); + + await expect(store.close(logger)).resolves.toBeUndefined(); + expect(logger).toHaveBeenCalledWith( + 'Error', + 'Failed to close database connection', + expect.objectContaining({ error: 'close failed' }), + ); + }); + }); +}); diff --git a/packages/workflow-executor/test/stores/database-store-schema.test.ts b/packages/workflow-executor/test/stores/database-store-schema.test.ts index 87e3d7cb4a..b80d33136d 100644 --- a/packages/workflow-executor/test/stores/database-store-schema.test.ts +++ b/packages/workflow-executor/test/stores/database-store-schema.test.ts @@ -238,6 +238,27 @@ describe('DatabaseStore — schema namespacing', () => { await expect(new DatabaseStore({ sequelize }).init()).rejects.toThrow('pool.max >= 2'); }); + it('sizes the pool check against the write sub-pool when replication makes maxSize undefined', async () => { + const { sequelize, calls } = setup('postgres'); + (sequelize as unknown as { connectionManager: { pool: unknown } }).connectionManager.pool = { + write: { maxSize: 5 }, + read: { maxSize: 5 }, + }; + + await new DatabaseStore({ sequelize }).init(); + + expect(calls).toContain('migrate'); // passed the pool>=2 guard and proceeded to migrate + }); + + it('throws when the replication write-pool max is < 2', async () => { + const { sequelize } = setup('postgres'); + (sequelize as unknown as { connectionManager: { pool: unknown } }).connectionManager.pool = { + write: { maxSize: 1 }, + }; + + await expect(new DatabaseStore({ sequelize }).init()).rejects.toThrow('pool.max >= 2'); + }); + it('does not open a transaction or take a lock on SQLite', async () => { const { sequelize, query } = setup('sqlite'); diff --git a/packages/workflow-executor/test/stores/database-store.test.ts b/packages/workflow-executor/test/stores/database-store.test.ts index cd385b49a0..6e3ce0d4b3 100644 --- a/packages/workflow-executor/test/stores/database-store.test.ts +++ b/packages/workflow-executor/test/stores/database-store.test.ts @@ -27,6 +27,32 @@ describe('DatabaseStore (SQLite)', () => { await store.close(); }); + it('clears the step execution for a given (runId, stepIndex)', async () => { + await store.saveStepExecution('run-1', makeStepExecution({ stepIndex: 0 })); + + await store.deleteStepExecution('run-1', 0); + + expect(await store.getStepExecutions('run-1')).toEqual([]); + }); + + it('leaves other steps in the same run untouched when clearing one step', async () => { + const kept = makeStepExecution({ stepIndex: 1, type: 'read-record' } as never); + await store.saveStepExecution('run-1', makeStepExecution({ stepIndex: 0 })); + await store.saveStepExecution('run-1', kept); + + await store.deleteStepExecution('run-1', 0); + + expect(await store.getStepExecutions('run-1')).toEqual([kept]); + }); + + it('is a no-op when no execution exists for that (runId, stepIndex)', async () => { + await store.saveStepExecution('run-1', makeStepExecution({ stepIndex: 0 })); + + await store.deleteStepExecution('run-1', 99); + + expect((await store.getStepExecutions('run-1')).map(s => s.stepIndex)).toEqual([0]); + }); + it('returns empty array for unknown runId', async () => { const result = await store.getStepExecutions('unknown'); expect(result).toEqual([]); diff --git a/packages/workflow-executor/test/stores/in-memory-mcp-oauth-credentials-store.test.ts b/packages/workflow-executor/test/stores/in-memory-mcp-oauth-credentials-store.test.ts new file mode 100644 index 0000000000..6f1a82f1d0 --- /dev/null +++ b/packages/workflow-executor/test/stores/in-memory-mcp-oauth-credentials-store.test.ts @@ -0,0 +1,165 @@ +/** + * Spec for the in-memory MCP OAuth credentials store (dev / --in-memory). + * + * Behaviour mirrors the database store's contract, minus persistence: + * - One entry per (userId, mcpServerId) — upsert overwrites in place. + * - Stores opaque encrypted bytes (no crypto in the store itself). + * - get returns null for unknown keys; delete is a no-op for unknown keys. + * - State is process-local and lost on restart (same throwaway semantics as InMemoryStore). + */ +import type { McpOAuthCredentialInput } from '../../src/ports/mcp-oauth-credentials-store'; + +import InMemoryMcpOAuthCredentialsStore from '../../src/stores/in-memory-mcp-oauth-credentials-store'; + +function makeCredential(overrides: Partial = {}): McpOAuthCredentialInput { + return { + userId: 42, + mcpServerId: 'mcp-server-1', + refreshTokenEnc: Buffer.from('enc-refresh-token'), + clientId: 'client-abc', + clientSecretEnc: Buffer.from('enc-client-secret'), + clientSecretExpiresAt: null, + tokenEndpoint: 'https://auth.example.com/token', + tokenEndpointAuthMethod: 'client_secret_post', + scopes: 'read write', + ...overrides, + }; +} + +// Asserts presence and narrows the type — avoids non-null assertions (`!`), which the codebase avoids. +function unwrap(value: T | null | undefined): T { + if (value === null || value === undefined) { + throw new Error('expected a stored credential, got null/undefined'); + } + + return value; +} + +describe('InMemoryMcpOAuthCredentialsStore', () => { + let store: InMemoryMcpOAuthCredentialsStore; + + beforeEach(() => { + store = new InMemoryMcpOAuthCredentialsStore(); + }); + + describe('get', () => { + it('returns null for an unknown (userId, mcpServerId)', async () => { + expect(await store.get(999, 'no-such-server')).toBeNull(); + }); + + it('returns the stored credential for a known (userId, mcpServerId)', async () => { + await store.upsert(makeCredential()); + + expect(await store.get(42, 'mcp-server-1')).toEqual( + expect.objectContaining({ + userId: 42, + mcpServerId: 'mcp-server-1', + tokenEndpoint: 'https://auth.example.com/token', + }), + ); + }); + + it('preserves the encrypted blobs byte-for-byte', async () => { + const refreshTokenEnc = Buffer.from([0x00, 0x01, 0xfe, 0xff]); + await store.upsert(makeCredential({ refreshTokenEnc })); + + const row = unwrap(await store.get(42, 'mcp-server-1')); + + expect(row.refreshTokenEnc.toString('hex')).toBe(refreshTokenEnc.toString('hex')); + }); + }); + + describe('upsert', () => { + it('overwrites the existing entry in place for the same key', async () => { + await store.upsert(makeCredential({ refreshTokenEnc: Buffer.from('old') })); + await store.upsert(makeCredential({ refreshTokenEnc: Buffer.from('new') })); + + const row = unwrap(await store.get(42, 'mcp-server-1')); + + expect(row.refreshTokenEnc.toString()).toBe('new'); + }); + + it('assigns a positive integer id', async () => { + await store.upsert(makeCredential()); + + expect(unwrap(await store.get(42, 'mcp-server-1')).id).toBeGreaterThanOrEqual(1); + }); + }); + + describe('updateIfPresent', () => { + it('updates the row matching the given id in place', async () => { + await store.upsert(makeCredential({ refreshTokenEnc: Buffer.from('old') })); + const { id } = unwrap(await store.get(42, 'mcp-server-1')); + + await store.updateIfPresent(id, makeCredential({ refreshTokenEnc: Buffer.from('rotated') })); + + expect(unwrap(await store.get(42, 'mcp-server-1')).refreshTokenEnc.toString()).toBe( + 'rotated', + ); + }); + + it('does not insert when the row is absent', async () => { + await store.updateIfPresent(1, makeCredential()); + + expect(await store.get(42, 'mcp-server-1')).toBeNull(); + }); + + it('does not touch a row that was re-created with a different id', async () => { + await store.upsert(makeCredential({ refreshTokenEnc: Buffer.from('old') })); + const staleId = unwrap(await store.get(42, 'mcp-server-1')).id; + await store.delete(42, 'mcp-server-1'); + await store.upsert(makeCredential({ refreshTokenEnc: Buffer.from('reauthorized') })); + + await store.updateIfPresent(staleId, makeCredential({ refreshTokenEnc: Buffer.from('rot') })); + + expect(unwrap(await store.get(42, 'mcp-server-1')).refreshTokenEnc.toString()).toBe( + 'reauthorized', + ); + }); + }); + + describe('isolation', () => { + it('keeps entries for different users and servers separate', async () => { + await store.upsert(makeCredential({ userId: 1, refreshTokenEnc: Buffer.from('user-1') })); + await store.upsert(makeCredential({ userId: 2, refreshTokenEnc: Buffer.from('user-2') })); + await store.upsert( + makeCredential({ mcpServerId: 'server-b', refreshTokenEnc: Buffer.from('b') }), + ); + + expect(unwrap(await store.get(1, 'mcp-server-1')).refreshTokenEnc.toString()).toBe('user-1'); + expect(unwrap(await store.get(2, 'mcp-server-1')).refreshTokenEnc.toString()).toBe('user-2'); + expect(unwrap(await store.get(42, 'server-b')).refreshTokenEnc.toString()).toBe('b'); + }); + }); + + describe('delete', () => { + it('removes the credential for a (userId, mcpServerId)', async () => { + await store.upsert(makeCredential()); + + await store.delete(42, 'mcp-server-1'); + + expect(await store.get(42, 'mcp-server-1')).toBeNull(); + }); + + it('is a no-op (does not throw) for an unknown credential', async () => { + await expect(store.delete(999, 'no-such-server')).resolves.toBeUndefined(); + }); + + it('does not affect other users when deleting one user', async () => { + await store.upsert(makeCredential({ userId: 1 })); + await store.upsert(makeCredential({ userId: 2 })); + + await store.delete(1, 'mcp-server-1'); + + expect(await store.get(1, 'mcp-server-1')).toBeNull(); + expect(await store.get(2, 'mcp-server-1')).not.toBeNull(); + }); + }); + + describe('lifecycle', () => { + it('init and close are no-ops that resolve', async () => { + await expect(store.init()).resolves.toBeUndefined(); + await expect(store.close()).resolves.toBeUndefined(); + }); + }); +}); diff --git a/packages/workflow-executor/test/stores/in-memory-store.test.ts b/packages/workflow-executor/test/stores/in-memory-store.test.ts index adfe353da8..d9aeb2497d 100644 --- a/packages/workflow-executor/test/stores/in-memory-store.test.ts +++ b/packages/workflow-executor/test/stores/in-memory-store.test.ts @@ -19,6 +19,32 @@ describe('InMemoryStore', () => { store = new InMemoryStore(); }); + it('clears the step execution for a given (runId, stepIndex)', async () => { + await store.saveStepExecution('run-1', makeStepExecution({ stepIndex: 0 })); + + await store.deleteStepExecution('run-1', 0); + + expect(await store.getStepExecutions('run-1')).toEqual([]); + }); + + it('leaves other steps in the same run untouched when clearing one step', async () => { + const kept = makeStepExecution({ stepIndex: 1, type: 'read-record' } as never); + await store.saveStepExecution('run-1', makeStepExecution({ stepIndex: 0 })); + await store.saveStepExecution('run-1', kept); + + await store.deleteStepExecution('run-1', 0); + + expect(await store.getStepExecutions('run-1')).toEqual([kept]); + }); + + it('is a no-op when no execution exists for that (runId, stepIndex)', async () => { + await store.saveStepExecution('run-1', makeStepExecution({ stepIndex: 0 })); + + await store.deleteStepExecution('run-1', 99); + + expect((await store.getStepExecutions('run-1')).map(s => s.stepIndex)).toEqual([0]); + }); + it('returns empty array for unknown runId', async () => { const result = await store.getStepExecutions('unknown'); expect(result).toEqual([]); diff --git a/packages/workflow-executor/test/types/step-outcome.test.ts b/packages/workflow-executor/test/types/step-outcome.test.ts index b137ea0bc6..419077e64b 100644 --- a/packages/workflow-executor/test/types/step-outcome.test.ts +++ b/packages/workflow-executor/test/types/step-outcome.test.ts @@ -1,5 +1,8 @@ import { StepType } from '../../src/types/validated/step-definition'; -import { stepTypeToOutcomeType } from '../../src/types/validated/step-outcome'; +import { + McpStepOutcomeSchema, + stepTypeToOutcomeType, +} from '../../src/types/validated/step-outcome'; describe('stepTypeToOutcomeType', () => { it('maps Condition to condition', () => { @@ -34,3 +37,43 @@ describe('stepTypeToOutcomeType', () => { expect(stepTypeToOutcomeType('future-step-type' as StepType)).toBe('record'); }); }); + +describe('McpStepOutcomeSchema — awaitingInputReason', () => { + const base = { type: 'mcp' as const, stepId: 'step-1', stepIndex: 0 }; + + it("accepts an awaiting-input outcome carrying awaitingInputReason 'needs-oauth-reauth'", () => { + const parsed = McpStepOutcomeSchema.parse({ + ...base, + status: 'awaiting-input', + awaitingInputReason: 'needs-oauth-reauth', + }); + + expect(parsed.awaitingInputReason).toBe('needs-oauth-reauth'); + }); + + it('allows awaitingInputReason to be omitted', () => { + const parsed = McpStepOutcomeSchema.parse({ ...base, status: 'awaiting-input' }); + + expect(parsed.awaitingInputReason).toBeUndefined(); + }); + + it('rejects an unknown awaitingInputReason value', () => { + expect(() => + McpStepOutcomeSchema.parse({ + ...base, + status: 'awaiting-input', + awaitingInputReason: 'nope', + }), + ).toThrow(); + }); + + it("rejects the legacy 'reason' key under the strict schema", () => { + expect(() => + McpStepOutcomeSchema.parse({ + ...base, + status: 'awaiting-input', + reason: 'needs-oauth-reauth', + }), + ).toThrow(); + }); +});