From d462bf4113f2e8ca1798cc4cdd91d0800b3d07f1 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Fri, 22 May 2026 20:43:34 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=AA=AA=20fix:=20Consolidate=20MCP=20OAuth?= =?UTF-8?q?=20Policy=20(#13254)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🦉 feat: Implement proactive OAuth flow for connections without stored tokens * 🤝 fix: Enhance proactive OAuth flow handling in MCPConnectionFactory * fix: Add timeout handling for proactive OAuth flow in MCPConnectionFactory Co-authored-by: Copilot * fix: Refine proactive MCP OAuth flow * test: Cover proactive OAuth missing handler * fix: Require explicit MCP OAuth signal * fix: Consolidate MCP OAuth policy --------- Co-authored-by: Gil Assunção Co-authored-by: Copilot --- packages/api/src/mcp/MCPConnectionFactory.ts | 89 +++- packages/api/src/mcp/MCPManager.ts | 4 +- packages/api/src/mcp/UserConnectionManager.ts | 60 ++- .../__tests__/MCPConnectionFactory.test.ts | 450 ++++++++++++++++++ .../api/src/mcp/__tests__/MCPManager.test.ts | 114 +++++ packages/api/src/mcp/types/index.ts | 13 + packages/api/src/mcp/utils.ts | 10 + 7 files changed, 710 insertions(+), 30 deletions(-) diff --git a/packages/api/src/mcp/MCPConnectionFactory.ts b/packages/api/src/mcp/MCPConnectionFactory.ts index aa2dd5d6d3..68e3abd776 100644 --- a/packages/api/src/mcp/MCPConnectionFactory.ts +++ b/packages/api/src/mcp/MCPConnectionFactory.ts @@ -7,7 +7,7 @@ import type { FlowStateManager } from '~/flow/manager'; import type * as t from './types'; import { MCPTokenStorage, MCPOAuthHandler, ReauthenticationRequiredError } from '~/mcp/oauth'; import { PENDING_STALE_MS, normalizeExpiresAt } from '~/flow/manager'; -import { sanitizeUrlForLogging, isClientRejectionMessage } from './utils'; +import { sanitizeUrlForLogging, isClientRejectionMessage, isOAuthServer } from './utils'; import { withTimeout } from '~/utils/promise'; import { MCPConnection } from './connection'; import { processMCPEnv } from '~/utils'; @@ -46,7 +46,7 @@ export class MCPConnectionFactory { /** Creates a new MCP connection with optional OAuth support */ static async create( basic: t.BasicConnectionOptions, - oauth?: t.OAuthConnectionOptions, + oauth?: t.OAuthConnectionOptions | t.UserConnectionContext, ): Promise { const factory = new this(basic, oauth); return factory.createConnection(); @@ -237,9 +237,23 @@ export class MCPConnectionFactory { let cleanupOAuthHandlers: (() => void) | null = null; if (this.useOAuth) { cleanupOAuthHandlers = this.handleOAuthEvents(connection); + } else { + const nonOAuthHandler = () => { + logger.info( + `${this.logPrefix} Server does not use OAuth; treating 401/403 as auth failure`, + ); + connection.emit('oauthFailed', new Error('Server does not use OAuth')); + }; + connection.on('oauthRequired', nonOAuthHandler); + cleanupOAuthHandlers = () => { + connection.removeListener('oauthRequired', nonOAuthHandler); + }; } try { + if (this.shouldInitiateOAuthBeforeConnect(oauthTokens)) { + await this.initiateOAuthBeforeConnect(connection); + } await this.attemptToConnect(connection); if (cleanupOAuthHandlers) { cleanupOAuthHandlers(); @@ -253,6 +267,77 @@ export class MCPConnectionFactory { } } + private shouldInitiateOAuthBeforeConnect(oauthTokens: MCPOAuthTokens | null): boolean { + if (!this.useOAuth || oauthTokens) { + return false; + } + return isOAuthServer(this.serverConfig); + } + + private getServerUrl(): string | undefined { + return 'url' in this.serverConfig ? this.serverConfig.url : undefined; + } + + private async initiateOAuthBeforeConnect(connection: MCPConnection): Promise { + const serverUrl = this.getServerUrl(); + if (!serverUrl) { + throw new Error(`${this.logPrefix} OAuth required but server URL is missing from config`); + } + + const oauthTimeout = this.connectionTimeout ?? 60000 * 2; + logger.info( + `${this.logPrefix} No stored tokens, proactively triggering OAuth flow before connecting (timeout: ${oauthTimeout}ms)`, + ); + + await new Promise((resolve, reject) => { + let timeoutId: ReturnType | null = null; + let oauthHandledListener: (() => void) | null = null; + let oauthFailedListener: ((error: Error) => void) | null = null; + + const cleanup = () => { + if (timeoutId) { + clearTimeout(timeoutId); + } + if (oauthHandledListener) { + connection.off('oauthHandled', oauthHandledListener); + } + if (oauthFailedListener) { + connection.off('oauthFailed', oauthFailedListener); + } + }; + + oauthHandledListener = () => { + cleanup(); + resolve(); + }; + + oauthFailedListener = (error: Error) => { + cleanup(); + reject(error); + }; + + timeoutId = setTimeout(() => { + cleanup(); + reject(new Error(`Proactive OAuth flow timeout after ${oauthTimeout}ms`)); + }, oauthTimeout); + + connection.once('oauthHandled', oauthHandledListener); + connection.once('oauthFailed', oauthFailedListener); + + const emitted = connection.emit('oauthRequired', { + serverName: this.serverName, + error: new Error('OAuth tokens missing before connection'), + serverUrl, + userId: this.userId, + }); + + if (!emitted) { + cleanup(); + reject(new Error('OAuth required but no handler is registered')); + } + }); + } + /** Retrieves existing OAuth tokens from storage or returns null */ protected async getOAuthTokens(): Promise { if (!this.tokenMethods?.findToken) return null; diff --git a/packages/api/src/mcp/MCPManager.ts b/packages/api/src/mcp/MCPManager.ts index c0a7acfd06..7b0127b475 100644 --- a/packages/api/src/mcp/MCPManager.ts +++ b/packages/api/src/mcp/MCPManager.ts @@ -18,7 +18,7 @@ import { preProcessGraphTokens } from '~/utils/graph'; import { formatToolContent } from './parsers'; import { MCPConnection } from './connection'; import { processMCPEnv } from '~/utils/env'; -import { isUserSourced } from './utils'; +import { isUserSourced, isOAuthServer } from './utils'; /** * Centralized manager for MCP server connections and tool execution. @@ -102,7 +102,7 @@ export class MCPManager extends UserConnectionManager { return { tools: null, oauthRequired: false, oauthUrl: null }; } - const useOAuth = Boolean(serverConfig.requiresOAuth || serverConfig.oauthMetadata); + const useOAuth = isOAuthServer(serverConfig); const registry = MCPServersRegistry.getInstance(); const useSSRFProtection = registry.shouldEnableSSRFProtection(); diff --git a/packages/api/src/mcp/UserConnectionManager.ts b/packages/api/src/mcp/UserConnectionManager.ts index f352497f78..2a420861c3 100644 --- a/packages/api/src/mcp/UserConnectionManager.ts +++ b/packages/api/src/mcp/UserConnectionManager.ts @@ -4,7 +4,7 @@ import type * as t from './types'; import { MCPServersRegistry } from '~/mcp/registry/MCPServersRegistry'; import { ConnectionsRepository } from '~/mcp/ConnectionsRepository'; import { MCPConnectionFactory } from '~/mcp/MCPConnectionFactory'; -import { isUserSourced } from './utils'; +import { isUserSourced, isOAuthServer } from './utils'; import { MCPConnection } from './connection'; import { mcpConfig } from './mcpConfig'; @@ -35,14 +35,7 @@ export abstract class UserConnectionManager { } /** Gets or creates a connection for a specific user, coalescing concurrent attempts */ - public async getUserConnection( - opts: { - serverName: string; - forceNew?: boolean; - /** Pre-resolved config for config-source servers not in YAML/DB */ - serverConfig?: t.ParsedServerConfig; - } & Omit, - ): Promise { + public async getUserConnection(opts: t.UserMCPConnectionOptions): Promise { const { serverName, forceNew, user } = opts; const userId = user?.id; if (!userId) { @@ -89,11 +82,7 @@ export abstract class UserConnectionManager { returnOnOAuth = false, connectionTimeout, serverConfig: providedConfig, - }: { - serverName: string; - forceNew?: boolean; - serverConfig?: t.ParsedServerConfig; - } & Omit, + }: t.UserMCPConnectionOptions, userId: string, ): Promise { if (await this.appConnections!.has(serverName)) { @@ -161,16 +150,26 @@ export abstract class UserConnectionManager { try { const registry = MCPServersRegistry.getInstance(); - connection = await MCPConnectionFactory.create( - { - serverConfig: config, - serverName: serverName, - dbSourced: isUserSourced(config), - useSSRFProtection: registry.shouldEnableSSRFProtection(), - allowedDomains: registry.getAllowedDomains(), - allowedAddresses: registry.getAllowedAddresses(), - }, - { + const basic: t.BasicConnectionOptions = { + serverConfig: config, + serverName: serverName, + dbSourced: isUserSourced(config), + useSSRFProtection: registry.shouldEnableSSRFProtection(), + allowedDomains: registry.getAllowedDomains(), + allowedAddresses: registry.getAllowedAddresses(), + }; + + const useOAuth = isOAuthServer(config); + let connectionOptions: t.OAuthConnectionOptions | t.UserConnectionContext; + if (useOAuth) { + if (!flowManager) { + throw new McpError( + ErrorCode.InvalidRequest, + `[MCP][User: ${userId}] OAuth server "${serverName}" requires a flowManager`, + ); + } + + connectionOptions = { useOAuth: true, user: user, customUserVars: customUserVars, @@ -182,8 +181,17 @@ export abstract class UserConnectionManager { returnOnOAuth: returnOnOAuth, requestBody: requestBody, connectionTimeout: connectionTimeout, - }, - ); + }; + } else { + connectionOptions = { + user, + customUserVars, + requestBody, + connectionTimeout, + }; + } + + connection = await MCPConnectionFactory.create(basic, connectionOptions); if (!(await connection?.isConnected())) { throw new Error('Failed to establish connection after initialization attempt.'); diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts index 0467a38c27..5a0787868d 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts @@ -95,6 +95,34 @@ describe('MCPConnectionFactory', () => { expect(mockConnectionInstance.connect).toHaveBeenCalled(); }); + it('should register fallback oauthRequired handler for non-OAuth connections', async () => { + const basicOptions = { + serverName: 'test-server', + serverConfig: mockServerConfig, + }; + + mockConnectionInstance.isConnected.mockResolvedValue(true); + + await MCPConnectionFactory.create(basicOptions); + + expect(mockConnectionInstance.on).toHaveBeenCalledWith('oauthRequired', expect.any(Function)); + + const onCall = (mockConnectionInstance.on as jest.Mock).mock.calls.find( + ([event]: [string]) => event === 'oauthRequired', + ); + const handler = onCall![1] as () => void; + handler(); + + expect(mockConnectionInstance.emit).toHaveBeenCalledWith( + 'oauthFailed', + expect.objectContaining({ message: 'Server does not use OAuth' }), + ); + expect(mockConnectionInstance.removeListener).toHaveBeenCalledWith( + 'oauthRequired', + expect.any(Function), + ); + }); + it('should create a connection with OAuth', async () => { const basicOptions = { serverName: 'test-server', @@ -1002,4 +1030,426 @@ describe('MCPConnectionFactory', () => { expect(mockLogger.debug).toHaveBeenCalled(); }); }); + + describe('proactive OAuth flow', () => { + const makeOAuthServerConfig = (): t.MCPOptions => + ({ + type: 'streamable-http' as const, + url: 'https://bigquery.googleapis.com/mcp', + initTimeout: 5000, + requiresOAuth: true, + }) as unknown as t.MCPOptions; + + const makeOAuthOptions = () => ({ + useOAuth: true as const, + user: mockUser, + flowManager: mockFlowManager, + tokenMethods: { + findToken: jest.fn(), + createToken: jest.fn(), + updateToken: jest.fn(), + deleteTokens: jest.fn(), + }, + }); + + function wireEventHandlers(instance: jest.Mocked) { + type Listener = (...args: unknown[]) => void; + + const handlers: Record = {}; + const onceWrappers = new Map(); + const key = (event: string | symbol): string => + typeof event === 'symbol' ? event.toString() : event; + const addHandler = (event: string | symbol, handler: Listener) => { + (handlers[key(event)] ??= []).push(handler); + }; + const removeHandler = (event: string | symbol, handler: Listener) => { + const list = handlers[key(event)]; + if (!list) { + return; + } + const wrapped = onceWrappers.get(handler); + const handlerToRemove = wrapped ?? handler; + const index = list.indexOf(handlerToRemove); + if (index !== -1) { + list.splice(index, 1); + } + if (wrapped) { + onceWrappers.delete(handler); + } + }; + + instance.on.mockImplementation((event: string | symbol, handler: Listener) => { + addHandler(event, handler); + return instance; + }); + + instance.once.mockImplementation((event: string | symbol, handler: Listener) => { + const wrapped: Listener = (...args) => { + removeHandler(event, handler); + handler(...args); + }; + onceWrappers.set(handler, wrapped); + addHandler(event, wrapped); + return instance; + }); + + instance.off.mockImplementation((event: string | symbol, handler: Listener) => { + removeHandler(event, handler); + return instance; + }); + + instance.removeListener.mockImplementation((event: string | symbol, handler: Listener) => { + removeHandler(event, handler); + return instance; + }); + + instance.emit.mockImplementation((event: string | symbol, ...args: unknown[]) => { + const list = handlers[key(event)]; + if (!list || list.length === 0) { + return false; + } + for (const fn of [...list]) { + fn(...args); + } + return true; + }); + + return handlers; + } + + it('should trigger proactive OAuth when requiresOAuth and no tokens', async () => { + const serverConfig = makeOAuthServerConfig(); + const oauthOptions = makeOAuthOptions(); + + mockProcessMCPEnv.mockReturnValue(serverConfig); + mockFlowManager.createFlowWithHandler.mockRejectedValue(new Error('no tokens')); + + const mockTokens: MCPOAuthTokens = { + access_token: 'bq-token', + token_type: 'Bearer', + obtained_at: Date.now(), + }; + + const mockFlowData = { + authorizationUrl: 'https://accounts.google.com/o/oauth2/auth?state=xyz', + flowId: 'flow-bq', + flowMetadata: { + serverName: 'bigquery', + userId: 'user123', + serverUrl: 'https://bigquery.googleapis.com/mcp', + state: 'state-xyz', + clientInfo: { client_id: 'bq-client' }, + }, + }; + + mockMCPOAuthHandler.generateFlowId.mockReturnValue('flow-bq'); + mockMCPOAuthHandler.initiateOAuthFlow.mockResolvedValue(mockFlowData); + mockFlowManager.getFlowState.mockResolvedValue(null); + mockFlowManager.createFlow.mockResolvedValue(mockTokens); + + wireEventHandlers(mockConnectionInstance); + mockConnectionInstance.isConnected.mockResolvedValue(true); + + const connection = await MCPConnectionFactory.create( + { serverName: 'bigquery', serverConfig }, + oauthOptions, + ); + + expect(connection).toBe(mockConnectionInstance); + expect(mockLogger.info).toHaveBeenCalledWith( + expect.stringContaining('proactively triggering OAuth flow'), + ); + expect(mockConnectionInstance.setOAuthTokens).toHaveBeenCalledWith(mockTokens); + expect(mockConnectionInstance.connect).toHaveBeenCalled(); + }); + + it('should trigger proactive OAuth when oauth is configured without requiresOAuth', async () => { + const serverConfig = { + type: 'streamable-http' as const, + url: 'https://drivemcp.googleapis.com/mcp/v1', + initTimeout: 5000, + oauth: { + authorization_url: 'https://accounts.google.com/o/oauth2/v2/auth', + token_url: 'https://oauth2.googleapis.com/token', + }, + } as t.ParsedServerConfig; + const oauthOptions = makeOAuthOptions(); + + mockProcessMCPEnv.mockReturnValue(serverConfig); + mockFlowManager.createFlowWithHandler.mockResolvedValue(null); + + const mockTokens: MCPOAuthTokens = { + access_token: 'drive-token', + token_type: 'Bearer', + obtained_at: Date.now(), + }; + + mockMCPOAuthHandler.generateFlowId.mockReturnValue('flow-drive'); + mockMCPOAuthHandler.initiateOAuthFlow.mockResolvedValue({ + authorizationUrl: 'https://accounts.google.com/o/oauth2/auth?state=drive', + flowId: 'flow-drive', + flowMetadata: { + serverName: 'drive', + userId: 'user123', + serverUrl: 'https://drivemcp.googleapis.com/mcp/v1', + state: 'state-drive', + }, + }); + mockFlowManager.getFlowState.mockResolvedValue(null); + mockFlowManager.createFlow.mockResolvedValue(mockTokens); + + wireEventHandlers(mockConnectionInstance); + mockConnectionInstance.isConnected.mockResolvedValue(true); + + const connection = await MCPConnectionFactory.create( + { serverName: 'drive', serverConfig }, + oauthOptions, + ); + + expect(connection).toBe(mockConnectionInstance); + expect(mockMCPOAuthHandler.initiateOAuthFlow).toHaveBeenCalled(); + expect(mockConnectionInstance.setOAuthTokens).toHaveBeenCalledWith(mockTokens); + expect(mockConnectionInstance.connect).toHaveBeenCalled(); + }); + + it('should not trigger proactive OAuth when only OAuth metadata is present', async () => { + const serverConfig = { + type: 'streamable-http' as const, + url: 'https://metadata-only.example.com/mcp', + initTimeout: 5000, + oauthMetadata: { + authorization_servers: ['https://auth.example.com/'], + }, + } as t.ParsedServerConfig; + const oauthOptions = makeOAuthOptions(); + + mockProcessMCPEnv.mockReturnValue(serverConfig); + mockFlowManager.createFlowWithHandler.mockResolvedValue(null); + mockConnectionInstance.isConnected.mockResolvedValue(true); + + const connection = await MCPConnectionFactory.create( + { serverName: 'metadata-only', serverConfig }, + oauthOptions, + ); + + expect(connection).toBe(mockConnectionInstance); + expect(mockMCPOAuthHandler.initiateOAuthFlow).not.toHaveBeenCalled(); + expect(mockConnectionInstance.connect).toHaveBeenCalled(); + }); + + it('should NOT trigger proactive OAuth when useOAuth is true but requiresOAuth is absent', async () => { + const serverConfig = { + command: 'node', + args: ['server.js'], + initTimeout: 5000, + } as t.MCPOptions; + + const oauthOptions = makeOAuthOptions(); + + mockProcessMCPEnv.mockReturnValue(serverConfig); + mockFlowManager.createFlowWithHandler.mockRejectedValue(new Error('no tokens')); + mockConnectionInstance.isConnected.mockResolvedValue(true); + + const connection = await MCPConnectionFactory.create( + { serverName: 'test-server', serverConfig }, + oauthOptions, + ); + + expect(connection).toBe(mockConnectionInstance); + expect(mockLogger.info).not.toHaveBeenCalledWith( + expect.stringContaining('proactively triggering OAuth flow'), + ); + }); + + it('should not trigger proactive OAuth when requiresOAuth is explicitly false', async () => { + const serverConfig = { + type: 'streamable-http' as const, + url: 'https://api.example.com/mcp', + initTimeout: 5000, + requiresOAuth: false, + oauth: { + authorization_url: 'https://auth.example.com/oauth/authorize', + token_url: 'https://auth.example.com/oauth/token', + }, + } as t.MCPOptions; + const oauthOptions = makeOAuthOptions(); + + mockProcessMCPEnv.mockReturnValue(serverConfig); + mockFlowManager.createFlowWithHandler.mockResolvedValue(null); + mockConnectionInstance.isConnected.mockResolvedValue(true); + + const connection = await MCPConnectionFactory.create( + { serverName: 'test-server', serverConfig }, + oauthOptions, + ); + + expect(connection).toBe(mockConnectionInstance); + expect(mockMCPOAuthHandler.initiateOAuthFlow).not.toHaveBeenCalled(); + expect(mockLogger.info).not.toHaveBeenCalledWith( + expect.stringContaining('proactively triggering OAuth flow'), + ); + }); + + it('should skip proactive OAuth when tokens already exist', async () => { + const serverConfig = makeOAuthServerConfig(); + const oauthOptions = makeOAuthOptions(); + + mockProcessMCPEnv.mockReturnValue(serverConfig); + const existingTokens: MCPOAuthTokens = { + access_token: 'existing-token', + token_type: 'Bearer', + obtained_at: Date.now(), + }; + mockFlowManager.createFlowWithHandler.mockResolvedValue(existingTokens); + mockConnectionInstance.isConnected.mockResolvedValue(true); + + const connection = await MCPConnectionFactory.create( + { serverName: 'bigquery', serverConfig }, + oauthOptions, + ); + + expect(connection).toBe(mockConnectionInstance); + expect(mockLogger.info).not.toHaveBeenCalledWith( + expect.stringContaining('proactively triggering OAuth flow'), + ); + }); + + it('should reject when proactive OAuth flow fails', async () => { + const serverConfig = makeOAuthServerConfig(); + const oauthOptions = { + ...makeOAuthOptions(), + returnOnOAuth: true, + oauthStart: jest.fn(), + }; + + mockProcessMCPEnv.mockReturnValue(serverConfig); + mockFlowManager.createFlowWithHandler.mockRejectedValue(new Error('no tokens')); + + const mockFlowData = { + authorizationUrl: 'https://accounts.google.com/o/oauth2/auth', + flowId: 'flow-bq', + flowMetadata: { + serverName: 'bigquery', + userId: 'user123', + serverUrl: 'https://bigquery.googleapis.com/mcp', + state: 'state-xyz', + }, + }; + + mockMCPOAuthHandler.generateFlowId.mockReturnValue('flow-bq'); + mockMCPOAuthHandler.initiateOAuthFlow.mockResolvedValue(mockFlowData); + mockFlowManager.getFlowState.mockResolvedValue(null); + mockFlowManager.createFlow.mockReturnValue(new Promise(() => {})); + + wireEventHandlers(mockConnectionInstance); + mockConnectionInstance.isConnected.mockResolvedValue(false); + + await expect( + MCPConnectionFactory.create({ serverName: 'bigquery', serverConfig }, oauthOptions), + ).rejects.toThrow('OAuth flow initiated - return early'); + }); + + it('should throw when requiresOAuth is true but url is missing', async () => { + const serverConfig = { + type: 'streamable-http' as const, + initTimeout: 5000, + requiresOAuth: true, + } as unknown as t.MCPOptions; + + const oauthOptions = makeOAuthOptions(); + + mockProcessMCPEnv.mockReturnValue(serverConfig); + mockFlowManager.createFlowWithHandler.mockRejectedValue(new Error('no tokens')); + + wireEventHandlers(mockConnectionInstance); + mockConnectionInstance.isConnected.mockResolvedValue(false); + + await expect( + MCPConnectionFactory.create({ serverName: 'no-url', serverConfig }, oauthOptions), + ).rejects.toThrow('server URL is missing'); + }); + + it('should reject when proactive OAuth has no registered handler', async () => { + const serverConfig = makeOAuthServerConfig(); + const oauthOptions = makeOAuthOptions(); + + mockProcessMCPEnv.mockReturnValue(serverConfig); + mockFlowManager.createFlowWithHandler.mockResolvedValue(null); + mockConnectionInstance.on.mockReturnValue(mockConnectionInstance); + mockConnectionInstance.once.mockReturnValue(mockConnectionInstance); + mockConnectionInstance.off.mockReturnValue(mockConnectionInstance); + mockConnectionInstance.emit.mockReturnValue(false); + + await expect( + MCPConnectionFactory.create({ serverName: 'bigquery', serverConfig }, oauthOptions), + ).rejects.toThrow('OAuth required but no handler is registered'); + + expect(mockConnectionInstance.connect).not.toHaveBeenCalled(); + }); + + it('should clean up cross-listeners when oauthHandled fires', async () => { + const serverConfig = makeOAuthServerConfig(); + const oauthOptions = makeOAuthOptions(); + + mockProcessMCPEnv.mockReturnValue(serverConfig); + mockFlowManager.createFlowWithHandler.mockRejectedValue(new Error('no tokens')); + + const mockTokens: MCPOAuthTokens = { + access_token: 'cleanup-token', + token_type: 'Bearer', + obtained_at: Date.now(), + }; + + mockMCPOAuthHandler.generateFlowId.mockReturnValue('flow-cleanup'); + mockMCPOAuthHandler.initiateOAuthFlow.mockResolvedValue({ + authorizationUrl: 'https://auth.example.com', + flowId: 'flow-cleanup', + flowMetadata: { + serverName: 'bigquery', + userId: 'user123', + serverUrl: 'https://bigquery.googleapis.com/mcp', + state: 'state-cleanup', + clientInfo: { client_id: 'client-cleanup' }, + }, + }); + mockFlowManager.getFlowState.mockResolvedValue(null); + mockFlowManager.createFlow.mockResolvedValue(mockTokens); + + const handlers = wireEventHandlers(mockConnectionInstance); + mockConnectionInstance.isConnected.mockResolvedValue(true); + + await MCPConnectionFactory.create({ serverName: 'bigquery', serverConfig }, oauthOptions); + + // After oauthHandled resolved, the oauthFailed listener should have been removed + const failedListeners = handlers['oauthFailed'] ?? []; + expect(failedListeners.length).toBe(0); + }); + + it('should not trigger proactive OAuth during tool discovery', async () => { + const serverConfig = makeOAuthServerConfig(); + const oauthOptions = { + ...makeOAuthOptions(), + oauthStart: jest.fn(), + }; + const mockTools = [ + { name: 'tool1', description: 'First tool', inputSchema: { type: 'object' } }, + ]; + + mockProcessMCPEnv.mockReturnValue(serverConfig); + mockFlowManager.createFlowWithHandler.mockResolvedValue(null); + mockConnectionInstance.connect.mockResolvedValue(undefined); + mockConnectionInstance.isConnected.mockResolvedValue(true); + mockConnectionInstance.fetchTools = jest.fn().mockResolvedValue(mockTools); + + const result = await MCPConnectionFactory.discoverTools( + { serverName: 'bigquery', serverConfig }, + oauthOptions, + ); + + expect(result.tools).toEqual(mockTools); + expect(result.oauthRequired).toBe(false); + expect(oauthOptions.oauthStart).not.toHaveBeenCalled(); + expect(mockMCPOAuthHandler.initiateOAuthFlow).not.toHaveBeenCalled(); + }); + }); }); diff --git a/packages/api/src/mcp/__tests__/MCPManager.test.ts b/packages/api/src/mcp/__tests__/MCPManager.test.ts index 45a51e01ac..cd30a75318 100644 --- a/packages/api/src/mcp/__tests__/MCPManager.test.ts +++ b/packages/api/src/mcp/__tests__/MCPManager.test.ts @@ -905,6 +905,28 @@ describe('MCPManager', () => { ); }); + it('should treat configured oauth as OAuth when requiresOAuth is unset', async () => { + mockAppConnections({ + get: jest.fn().mockResolvedValue(null), + }); + + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue({ + type: 'sse', + url: 'https://api.example.com', + oauth: { + authorization_url: 'https://auth.example.com/oauth/authorize', + token_url: 'https://auth.example.com/oauth/token', + }, + }); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + const result = await manager.discoverServerTools({ serverName }); + + expect(result.tools).toBeNull(); + expect(result.oauthRequired).toBe(true); + expect(MCPConnectionFactory.discoverTools).not.toHaveBeenCalled(); + }); + it('should return OAuth info when server requires OAuth but no user provided', async () => { mockAppConnections({ get: jest.fn().mockResolvedValue(null), @@ -967,4 +989,96 @@ describe('MCPManager', () => { ); }); }); + + describe('getUserConnection - useOAuth derivation', () => { + const mockUser = { id: userId, email: 'test@example.com' } as unknown as IUser; + const mockFlowManager = { + createFlow: jest.fn(), + getFlowState: jest.fn(), + deleteFlow: jest.fn(), + }; + const mockConnection = { + isConnected: jest.fn().mockResolvedValue(true), + isStale: jest.fn().mockReturnValue(false), + disconnect: jest.fn(), + } as unknown as MCPConnection; + + it('should pass useOAuth for servers with configured oauth and no requiresOAuth value', async () => { + mockAppConnections({ + has: jest.fn().mockResolvedValue(false), + }); + + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue({ + type: 'sse', + url: 'https://oauth-mcp.example.com', + oauth: { + authorization_url: 'https://auth.example.com/oauth/authorize', + token_url: 'https://auth.example.com/oauth/token', + }, + }); + + (MCPConnectionFactory.create as jest.Mock).mockResolvedValue(mockConnection); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + await manager.getUserConnection({ + serverName, + user: mockUser, + flowManager: mockFlowManager as unknown as t.UserMCPConnectionOptions['flowManager'], + }); + + expect(MCPConnectionFactory.create).toHaveBeenCalledWith( + expect.objectContaining({ serverName }), + expect.objectContaining({ useOAuth: true }), + ); + }); + + it('should not pass useOAuth for servers with requiresOAuth: false', async () => { + mockAppConnections({ + has: jest.fn().mockResolvedValue(false), + }); + + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue({ + type: 'streamable-http', + url: 'http://private-mcp.svc:5446/mcp', + requiresOAuth: false, + oauth: { + authorization_url: 'https://auth.example.com/oauth/authorize', + token_url: 'https://auth.example.com/oauth/token', + }, + }); + + (MCPConnectionFactory.create as jest.Mock).mockResolvedValue(mockConnection); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + await manager.getUserConnection({ + serverName, + user: mockUser, + }); + + expect(MCPConnectionFactory.create).toHaveBeenCalledWith( + expect.objectContaining({ serverName }), + expect.not.objectContaining({ useOAuth: true }), + ); + }); + + it('should throw when OAuth server lacks flowManager', async () => { + mockAppConnections({ + has: jest.fn().mockResolvedValue(false), + }); + + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue({ + type: 'sse', + url: 'https://oauth-mcp.example.com', + requiresOAuth: true, + }); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + await expect( + manager.getUserConnection({ + serverName, + user: mockUser, + }), + ).rejects.toThrow('requires a flowManager'); + }); + }); }); diff --git a/packages/api/src/mcp/types/index.ts b/packages/api/src/mcp/types/index.ts index 6bc593c1eb..ed85c13fa9 100644 --- a/packages/api/src/mcp/types/index.ts +++ b/packages/api/src/mcp/types/index.ts @@ -204,6 +204,19 @@ export interface OAuthConnectionOptions extends UserConnectionContext { returnOnOAuth?: boolean; } +/** Options accepted by UserConnectionManager.getUserConnection. OAuth fields are optional. */ +export interface UserMCPConnectionOptions extends UserConnectionContext { + serverName: string; + forceNew?: boolean; + serverConfig?: ParsedServerConfig; + flowManager?: FlowStateManager; + tokenMethods?: TokenMethods; + signal?: AbortSignal; + oauthStart?: (authURL: string) => Promise; + oauthEnd?: () => Promise; + returnOnOAuth?: boolean; +} + export interface ToolDiscoveryOptions { serverName: string; user?: IUser; diff --git a/packages/api/src/mcp/utils.ts b/packages/api/src/mcp/utils.ts index 9befc3c549..872f43002c 100644 --- a/packages/api/src/mcp/utils.ts +++ b/packages/api/src/mcp/utils.ts @@ -3,6 +3,16 @@ import type { ParsedServerConfig } from '~/mcp/types'; export const mcpToolPattern = new RegExp(`^.+${Constants.mcp_delimiter}.+$`); +/** Whether a server should use MCP OAuth handling. */ +export function isOAuthServer( + config: Pick, +): boolean { + if (config.requiresOAuth === false) { + return false; + } + return config.requiresOAuth === true || config.oauth != null; +} + /** Checks that `customUserVars` is present AND non-empty (guards against truthy `{}`) */ export function hasCustomUserVars(config: Pick): boolean { return !!config.customUserVars && Object.keys(config.customUserVars).length > 0;