From 83c7494c161e7c831c857e69a67cb53cd092d241 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Vincent=20Molini=C3=A9?= Date: Tue, 27 Jan 2026 14:51:11 +0100 Subject: [PATCH 1/5] feat(mcp server): add support for oauth authentication --- packages/agent/src/routes/ai/ai-proxy.ts | 4 + .../agent/test/routes/ai/ai-proxy.test.ts | 27 ++- packages/ai-proxy/src/index.ts | 7 +- packages/ai-proxy/src/mcp-client.ts | 74 ++++++- packages/ai-proxy/src/router.ts | 16 +- .../ai-proxy/src/types/mcp-config-checker.ts | 4 +- packages/ai-proxy/test/mcp-client.test.ts | 185 ++++++++++++++++++ packages/ai-proxy/test/router.test.ts | 21 ++ 8 files changed, 327 insertions(+), 11 deletions(-) diff --git a/packages/agent/src/routes/ai/ai-proxy.ts b/packages/agent/src/routes/ai/ai-proxy.ts index 85eb69521..0321fde3b 100644 --- a/packages/agent/src/routes/ai/ai-proxy.ts +++ b/packages/agent/src/routes/ai/ai-proxy.ts @@ -40,11 +40,15 @@ export default class AiProxyRoute extends BaseRoute { private async handleAiProxy(context: Context): Promise { try { + const mcpOauthTokensHeader = context.request.headers['x-mcp-oauth-tokens'] as string; + const mcpOAuthTokens = mcpOauthTokensHeader ? JSON.parse(mcpOauthTokensHeader) : undefined; + context.response.body = await this.aiProxyRouter.route({ route: context.params.route, body: context.request.body, query: context.query, mcpConfigs: await this.options.forestAdminClient.mcpServerConfigService.getConfiguration(), + mcpOAuthTokens, }); context.response.status = HttpCode.Ok; } catch (error) { diff --git a/packages/agent/test/routes/ai/ai-proxy.test.ts b/packages/agent/test/routes/ai/ai-proxy.test.ts index 0034914d1..5201d0b9c 100644 --- a/packages/agent/test/routes/ai/ai-proxy.test.ts +++ b/packages/agent/test/routes/ai/ai-proxy.test.ts @@ -85,7 +85,7 @@ describe('AiProxyRoute', () => { expect(context.response.body).toEqual(expectedResponse); }); - test('should pass route, body, query and mcpConfigs to router', async () => { + test('should pass route, body, query, mcpConfigs and headers to router', async () => { const route = new AiProxyRoute(services, options, aiConfigurations); mockRoute.mockResolvedValueOnce({}); @@ -105,9 +105,33 @@ describe('AiProxyRoute', () => { body: { messages: [{ role: 'user', content: 'Hello' }] }, query: { 'ai-name': 'gpt4' }, mcpConfigs: undefined, // mcpServerConfigService.getConfiguration returns undefined in test + headers: { mcpOauthTokens: undefined }, }); }); + test('should parse x-mcp-oauth-tokens header as JSON and pass to router', async () => { + const route = new AiProxyRoute(services, options, aiConfigurations); + mockRoute.mockResolvedValueOnce({}); + + const tokens = { server1: 'token1', server2: 'token2' }; + const context = createMockContext({ + customProperties: { + params: { route: 'ai-query' }, + }, + requestBody: { messages: [] }, + headers: { 'x-mcp-oauth-tokens': JSON.stringify(tokens) }, + }); + context.query = {}; + + await (route as any).handleAiProxy(context); + + expect(mockRoute).toHaveBeenCalledWith( + expect.objectContaining({ + headers: { mcpOauthTokens: tokens }, + }), + ); + }); + describe('error handling', () => { test('should convert AINotConfiguredError to UnprocessableError', async () => { const route = new AiProxyRoute(services, options, aiConfigurations); @@ -212,7 +236,6 @@ describe('AiProxyRoute', () => { }); context.query = {}; - // eslint-disable-next-line @typescript-eslint/no-explicit-any const promise = (route as any).handleAiProxy(context); await expect(promise).rejects.toBe(unknownError); diff --git a/packages/ai-proxy/src/index.ts b/packages/ai-proxy/src/index.ts index 2cc9c9892..d5915732d 100644 --- a/packages/ai-proxy/src/index.ts +++ b/packages/ai-proxy/src/index.ts @@ -9,6 +9,9 @@ export * from './mcp-client'; export * from './types/errors'; -export function validMcpConfigurationOrThrow(mcpConfig: McpConfiguration) { - return McpConfigChecker.check(mcpConfig); +export function validMcpConfigurationOrThrow( + mcpConfig: McpConfiguration, + mcpOauthTokens?: Record, +) { + return McpConfigChecker.check(mcpConfig, mcpOauthTokens); } diff --git a/packages/ai-proxy/src/mcp-client.ts b/packages/ai-proxy/src/mcp-client.ts index b30212a8f..3d78954a6 100644 --- a/packages/ai-proxy/src/mcp-client.ts +++ b/packages/ai-proxy/src/mcp-client.ts @@ -5,28 +5,62 @@ import { MultiServerMCPClient } from '@langchain/mcp-adapters'; import { McpConnectionError } from './types/errors'; import McpServerRemoteTool from './types/mcp-server-remote-tool'; +export type McpAuthenticationType = 'none' | 'bearer' | 'oauth2'; + export type McpConfiguration = { configs: MultiServerMCPClient['config']['mcpServers']; + authenticationTypes?: Record; } & Omit; export default class McpClient { private readonly mcpClients: Record = {}; + private readonly authenticationTypes: Record; private readonly logger?: Logger; readonly tools: McpServerRemoteTool[] = []; - constructor(config: McpConfiguration, logger?: Logger) { + constructor(config: McpConfiguration, logger?: Logger, mcpOauthTokens?: Record) { this.logger = logger; + this.authenticationTypes = config.authenticationTypes ?? {}; // 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 token = mcpOauthTokens?.[name]; + const configWithToken = this.injectOauthToken(serverConfig, token); this.mcpClients[name] = new MultiServerMCPClient({ - mcpServers: { [name]: serverConfig }, + mcpServers: { [name]: configWithToken }, ...config, }); }); } + /** + * Injects the OAuth token as Authorization header into HTTP-based transport configurations. + * The frontend sends 'x-mcp-oauth-tokens' with tokens keyed by server name. + * For stdio transports, returns the config unchanged. + */ + private injectOauthToken( + serverConfig: MultiServerMCPClient['config']['mcpServers'][string], + token?: string, + ): MultiServerMCPClient['config']['mcpServers'][string] { + if (!token) return serverConfig; + + // Only inject token for HTTP-based transports (sse, http) + if (serverConfig.type === 'http') { + const { oauthConfig, ...headers } = serverConfig.headers || {}; + + return { + ...serverConfig, + headers: { + ...headers, + Authorization: token, + }, + }; + } + + return serverConfig; + } + async loadTools(): Promise { const errors: Array<{ server: string; error: Error }> = []; @@ -61,7 +95,26 @@ export default class McpClient { async testConnections(): Promise { try { await Promise.all( - Object.values(this.mcpClients).map(client => client.initializeConnections()), + Object.entries(this.mcpClients).map(async ([name, client]) => { + const isOAuth = this.authenticationTypes[name] === 'oauth2'; + + try { + await client.initializeConnections(); + } catch (error) { + // For OAuth servers without a token, we expect auth errors (401/403) + // The server is reachable, just not authenticated yet + if (isOAuth && this.isAuthenticationError(error)) { + this.logger?.( + 'Info', + `OAuth server ${name} is reachable (authentication will be required)`, + ); + + return; + } + + throw error; + } + }), ); return true; @@ -77,6 +130,21 @@ export default class McpClient { } } + /** + * Checks if an error is an authentication error (401/403). + * These errors indicate the server is reachable but requires authentication. + */ + private isAuthenticationError(error: unknown): boolean { + const message = (error as Error)?.message?.toLowerCase() ?? ''; + + return ( + message.includes('401') || + message.includes('403') || + message.includes('unauthorized') || + message.includes('forbidden') + ); + } + async closeConnections(): Promise { const entries = Object.entries(this.mcpClients); const results = await Promise.allSettled(entries.map(([, client]) => client.close())); diff --git a/packages/ai-proxy/src/router.ts b/packages/ai-proxy/src/router.ts index 049d760a2..8d16ba9d0 100644 --- a/packages/ai-proxy/src/router.ts +++ b/packages/ai-proxy/src/router.ts @@ -16,6 +16,12 @@ export type Query = { }; export type ApiKeys = RemoteToolsApiKeys; +/** Headers that can be passed to the route method */ +export type RouteHeaders = { + /** OAuth tokens for MCP server authentication, keyed by server name */ + mcpOauthTokens?: Record; +}; + export class Router { private readonly localToolsApiKeys?: ApiKeys; private readonly aiConfigurations: AiConfiguration[]; @@ -61,12 +67,18 @@ export class Router { * - invoke-remote-tool: Execute a remote tool by name with the provided inputs * - remote-tools: Return the list of available remote tools definitions */ - async route(args: { body?: Body; route: Route; query?: Query; mcpConfigs?: McpConfiguration }) { + async route(args: { + body?: Body; + route: Route; + query?: Query; + mcpConfigs?: McpConfiguration; + mcpOAuthTokens?: Record; + }) { let mcpClient: McpClient | undefined; try { if (args.mcpConfigs) { - mcpClient = new McpClient(args.mcpConfigs, this.logger); + mcpClient = new McpClient(args.mcpConfigs, this.logger, args.mcpOAuthTokens); } const remoteTools = new RemoteTools( diff --git a/packages/ai-proxy/src/types/mcp-config-checker.ts b/packages/ai-proxy/src/types/mcp-config-checker.ts index 1474354ac..7ce26bf7b 100644 --- a/packages/ai-proxy/src/types/mcp-config-checker.ts +++ b/packages/ai-proxy/src/types/mcp-config-checker.ts @@ -3,7 +3,7 @@ import type { McpConfiguration } from '../mcp-client'; import McpClient from '../mcp-client'; export default class McpConfigChecker { - static check(mcpConfig: McpConfiguration) { - return new McpClient(mcpConfig).testConnections(); + static check(mcpConfig: McpConfiguration, mcpOauthTokens?: Record) { + return new McpClient(mcpConfig, undefined, mcpOauthTokens).testConnections(); } } diff --git a/packages/ai-proxy/test/mcp-client.test.ts b/packages/ai-proxy/test/mcp-client.test.ts index caa09eaf2..9668c78d7 100644 --- a/packages/ai-proxy/test/mcp-client.test.ts +++ b/packages/ai-proxy/test/mcp-client.test.ts @@ -19,6 +19,13 @@ jest.mock('@langchain/mcp-adapters', () => { }; }); +// eslint-disable-next-line import/first +import { MultiServerMCPClient } from '@langchain/mcp-adapters'; + +const MockedMultiServerMCPClient = MultiServerMCPClient as jest.MockedClass< + typeof MultiServerMCPClient +>; + describe('McpClient', () => { beforeEach(() => { jest.clearAllMocks(); @@ -236,4 +243,182 @@ describe('McpClient', () => { }); }); }); + + describe('OAuth token injection', () => { + it('should inject OAuth token as Authorization header into SSE transport', () => { + const sseConfig: McpConfiguration = { + configs: { + remote: { + transport: 'sse', + url: 'https://example.com/mcp', + }, + }, + }; + + // eslint-disable-next-line no-new + new McpClient(sseConfig, undefined, { remote: 'my-oauth-token' }); + + expect(MockedMultiServerMCPClient).toHaveBeenCalledWith( + expect.objectContaining({ + mcpServers: { + remote: { + transport: 'sse', + url: 'https://example.com/mcp', + headers: { + Authorization: 'Bearer my-oauth-token', + }, + }, + }, + }), + ); + }); + + it('should inject OAuth token as Authorization header into HTTP transport', () => { + const httpConfig: McpConfiguration = { + configs: { + remote: { + transport: 'http', + url: 'https://example.com/mcp', + }, + }, + }; + + // eslint-disable-next-line no-new + new McpClient(httpConfig, undefined, { remote: 'my-oauth-token' }); + + expect(MockedMultiServerMCPClient).toHaveBeenCalledWith( + expect.objectContaining({ + mcpServers: { + remote: { + transport: 'http', + url: 'https://example.com/mcp', + headers: { + Authorization: 'Bearer my-oauth-token', + }, + }, + }, + }), + ); + }); + + it('should merge Authorization header with existing headers', () => { + const sseConfig: McpConfiguration = { + configs: { + remote: { + transport: 'sse', + url: 'https://example.com/mcp', + headers: { + 'x-custom-header': 'custom-value', + }, + }, + }, + }; + + // eslint-disable-next-line no-new + new McpClient(sseConfig, undefined, { remote: 'my-oauth-token' }); + + expect(MockedMultiServerMCPClient).toHaveBeenCalledWith( + expect.objectContaining({ + mcpServers: { + remote: { + transport: 'sse', + url: 'https://example.com/mcp', + headers: { + 'x-custom-header': 'custom-value', + Authorization: 'Bearer my-oauth-token', + }, + }, + }, + }), + ); + }); + + it('should not inject OAuth token into stdio transport even if token provided', () => { + // eslint-disable-next-line no-new + new McpClient(aConfig, undefined, { slack: 'my-oauth-token' }); + + expect(MockedMultiServerMCPClient).toHaveBeenCalledWith( + expect.objectContaining({ + mcpServers: { + slack: { + transport: 'stdio', + command: 'npx', + args: ['-y', '@modelcontextprotocol/server-slack'], + env: {}, + }, + }, + }), + ); + }); + + it('should only inject token for servers that have a matching token', () => { + const multiServerConfig: McpConfiguration = { + configs: { + server1: { + transport: 'sse', + url: 'https://server1.com/mcp', + }, + server2: { + transport: 'sse', + url: 'https://server2.com/mcp', + }, + }, + }; + + // eslint-disable-next-line no-new + new McpClient(multiServerConfig, undefined, { server1: 'token-for-server1' }); + + // server1 should have the token + expect(MockedMultiServerMCPClient).toHaveBeenCalledWith( + expect.objectContaining({ + mcpServers: { + server1: { + transport: 'sse', + url: 'https://server1.com/mcp', + headers: { + Authorization: 'Bearer token-for-server1', + }, + }, + }, + }), + ); + + // server2 should not have any headers injected + expect(MockedMultiServerMCPClient).toHaveBeenCalledWith( + expect.objectContaining({ + mcpServers: { + server2: { + transport: 'sse', + url: 'https://server2.com/mcp', + }, + }, + }), + ); + }); + + it('should not modify config when no OAuth tokens are provided', () => { + const sseConfig: McpConfiguration = { + configs: { + remote: { + transport: 'sse', + url: 'https://example.com/mcp', + }, + }, + }; + + // eslint-disable-next-line no-new + new McpClient(sseConfig); + + expect(MockedMultiServerMCPClient).toHaveBeenCalledWith( + expect.objectContaining({ + mcpServers: { + remote: { + transport: 'sse', + url: 'https://example.com/mcp', + }, + }, + }), + ); + }); + }); }); diff --git a/packages/ai-proxy/test/router.test.ts b/packages/ai-proxy/test/router.test.ts index a5800e3a7..2e7bfa626 100644 --- a/packages/ai-proxy/test/router.test.ts +++ b/packages/ai-proxy/test/router.test.ts @@ -362,6 +362,27 @@ describe('route', () => { expect(MockedMcpClient).toHaveBeenCalledWith( { configs: { server1: { command: 'test', args: [] } } }, customLogger, + undefined, + ); + }); + + it('passes mcpOauthTokens to McpClient when provided in headers', async () => { + const customLogger: Logger = jest.fn(); + const router = new Router({ + logger: customLogger, + }); + + const tokens = { server1: 'token-for-server1' }; + await router.route({ + route: 'remote-tools', + mcpConfigs: { configs: { server1: { command: 'test', args: [] } } }, + headers: { mcpOauthTokens: tokens }, + }); + + expect(MockedMcpClient).toHaveBeenCalledWith( + { configs: { server1: { command: 'test', args: [] } } }, + customLogger, + tokens, ); }); }); From 49068ad678d022e660252d7016bfae8a8356a043 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Vincent=20Molini=C3=A9?= Date: Wed, 28 Jan 2026 15:58:20 +0100 Subject: [PATCH 2/5] test: fix tests --- .../agent/test/routes/ai/ai-proxy.test.ts | 6 +- packages/ai-proxy/test/mcp-client.test.ts | 67 ++++++------------- packages/ai-proxy/test/router.test.ts | 4 +- 3 files changed, 25 insertions(+), 52 deletions(-) diff --git a/packages/agent/test/routes/ai/ai-proxy.test.ts b/packages/agent/test/routes/ai/ai-proxy.test.ts index 5201d0b9c..9a9d53959 100644 --- a/packages/agent/test/routes/ai/ai-proxy.test.ts +++ b/packages/agent/test/routes/ai/ai-proxy.test.ts @@ -85,7 +85,7 @@ describe('AiProxyRoute', () => { expect(context.response.body).toEqual(expectedResponse); }); - test('should pass route, body, query, mcpConfigs and headers to router', async () => { + test('should pass route, body, query, mcpConfigs and mcpOAuthTokens to router', async () => { const route = new AiProxyRoute(services, options, aiConfigurations); mockRoute.mockResolvedValueOnce({}); @@ -105,7 +105,7 @@ describe('AiProxyRoute', () => { body: { messages: [{ role: 'user', content: 'Hello' }] }, query: { 'ai-name': 'gpt4' }, mcpConfigs: undefined, // mcpServerConfigService.getConfiguration returns undefined in test - headers: { mcpOauthTokens: undefined }, + mcpOAuthTokens: undefined, }); }); @@ -127,7 +127,7 @@ describe('AiProxyRoute', () => { expect(mockRoute).toHaveBeenCalledWith( expect.objectContaining({ - headers: { mcpOauthTokens: tokens }, + mcpOAuthTokens: tokens, }), ); }); diff --git a/packages/ai-proxy/test/mcp-client.test.ts b/packages/ai-proxy/test/mcp-client.test.ts index 9668c78d7..5652af533 100644 --- a/packages/ai-proxy/test/mcp-client.test.ts +++ b/packages/ai-proxy/test/mcp-client.test.ts @@ -245,39 +245,11 @@ describe('McpClient', () => { }); describe('OAuth token injection', () => { - it('should inject OAuth token as Authorization header into SSE transport', () => { - const sseConfig: McpConfiguration = { - configs: { - remote: { - transport: 'sse', - url: 'https://example.com/mcp', - }, - }, - }; - - // eslint-disable-next-line no-new - new McpClient(sseConfig, undefined, { remote: 'my-oauth-token' }); - - expect(MockedMultiServerMCPClient).toHaveBeenCalledWith( - expect.objectContaining({ - mcpServers: { - remote: { - transport: 'sse', - url: 'https://example.com/mcp', - headers: { - Authorization: 'Bearer my-oauth-token', - }, - }, - }, - }), - ); - }); - - it('should inject OAuth token as Authorization header into HTTP transport', () => { + it('should inject OAuth token as Authorization header into HTTP type transport', () => { const httpConfig: McpConfiguration = { configs: { remote: { - transport: 'http', + type: 'http', url: 'https://example.com/mcp', }, }, @@ -290,10 +262,10 @@ describe('McpClient', () => { expect.objectContaining({ mcpServers: { remote: { - transport: 'http', + type: 'http', url: 'https://example.com/mcp', headers: { - Authorization: 'Bearer my-oauth-token', + Authorization: 'my-oauth-token', }, }, }, @@ -301,31 +273,32 @@ describe('McpClient', () => { ); }); - it('should merge Authorization header with existing headers', () => { - const sseConfig: McpConfiguration = { + it('should merge Authorization header with existing headers and strip oauthConfig', () => { + const httpConfig: McpConfiguration = { configs: { remote: { - transport: 'sse', + type: 'http', url: 'https://example.com/mcp', headers: { 'x-custom-header': 'custom-value', + oauthConfig: { clientId: 'test' } as unknown as string, }, }, }, }; // eslint-disable-next-line no-new - new McpClient(sseConfig, undefined, { remote: 'my-oauth-token' }); + new McpClient(httpConfig, undefined, { remote: 'my-oauth-token' }); expect(MockedMultiServerMCPClient).toHaveBeenCalledWith( expect.objectContaining({ mcpServers: { remote: { - transport: 'sse', + type: 'http', url: 'https://example.com/mcp', headers: { 'x-custom-header': 'custom-value', - Authorization: 'Bearer my-oauth-token', + Authorization: 'my-oauth-token', }, }, }, @@ -355,11 +328,11 @@ describe('McpClient', () => { const multiServerConfig: McpConfiguration = { configs: { server1: { - transport: 'sse', + type: 'http', url: 'https://server1.com/mcp', }, server2: { - transport: 'sse', + type: 'http', url: 'https://server2.com/mcp', }, }, @@ -373,10 +346,10 @@ describe('McpClient', () => { expect.objectContaining({ mcpServers: { server1: { - transport: 'sse', + type: 'http', url: 'https://server1.com/mcp', headers: { - Authorization: 'Bearer token-for-server1', + Authorization: 'token-for-server1', }, }, }, @@ -388,7 +361,7 @@ describe('McpClient', () => { expect.objectContaining({ mcpServers: { server2: { - transport: 'sse', + type: 'http', url: 'https://server2.com/mcp', }, }, @@ -397,23 +370,23 @@ describe('McpClient', () => { }); it('should not modify config when no OAuth tokens are provided', () => { - const sseConfig: McpConfiguration = { + const httpConfig: McpConfiguration = { configs: { remote: { - transport: 'sse', + type: 'http', url: 'https://example.com/mcp', }, }, }; // eslint-disable-next-line no-new - new McpClient(sseConfig); + new McpClient(httpConfig); expect(MockedMultiServerMCPClient).toHaveBeenCalledWith( expect.objectContaining({ mcpServers: { remote: { - transport: 'sse', + type: 'http', url: 'https://example.com/mcp', }, }, diff --git a/packages/ai-proxy/test/router.test.ts b/packages/ai-proxy/test/router.test.ts index 2e7bfa626..2606d3aaa 100644 --- a/packages/ai-proxy/test/router.test.ts +++ b/packages/ai-proxy/test/router.test.ts @@ -366,7 +366,7 @@ describe('route', () => { ); }); - it('passes mcpOauthTokens to McpClient when provided in headers', async () => { + it('passes mcpOAuthTokens to McpClient when provided', async () => { const customLogger: Logger = jest.fn(); const router = new Router({ logger: customLogger, @@ -376,7 +376,7 @@ describe('route', () => { await router.route({ route: 'remote-tools', mcpConfigs: { configs: { server1: { command: 'test', args: [] } } }, - headers: { mcpOauthTokens: tokens }, + mcpOAuthTokens: tokens, }); expect(MockedMcpClient).toHaveBeenCalledWith( From e64624f4cedfca2a62fc9acccf6355d191860a53 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Vincent=20Molini=C3=A9?= Date: Wed, 28 Jan 2026 16:06:26 +0100 Subject: [PATCH 3/5] fix: remove useless code --- packages/ai-proxy/src/mcp-client.ts | 36 +---------------------------- 1 file changed, 1 insertion(+), 35 deletions(-) diff --git a/packages/ai-proxy/src/mcp-client.ts b/packages/ai-proxy/src/mcp-client.ts index 3d78954a6..eaf3a4981 100644 --- a/packages/ai-proxy/src/mcp-client.ts +++ b/packages/ai-proxy/src/mcp-client.ts @@ -95,26 +95,7 @@ export default class McpClient { async testConnections(): Promise { try { await Promise.all( - Object.entries(this.mcpClients).map(async ([name, client]) => { - const isOAuth = this.authenticationTypes[name] === 'oauth2'; - - try { - await client.initializeConnections(); - } catch (error) { - // For OAuth servers without a token, we expect auth errors (401/403) - // The server is reachable, just not authenticated yet - if (isOAuth && this.isAuthenticationError(error)) { - this.logger?.( - 'Info', - `OAuth server ${name} is reachable (authentication will be required)`, - ); - - return; - } - - throw error; - } - }), + Object.values(this.mcpClients).map(client => client.initializeConnections()), ); return true; @@ -130,21 +111,6 @@ export default class McpClient { } } - /** - * Checks if an error is an authentication error (401/403). - * These errors indicate the server is reachable but requires authentication. - */ - private isAuthenticationError(error: unknown): boolean { - const message = (error as Error)?.message?.toLowerCase() ?? ''; - - return ( - message.includes('401') || - message.includes('403') || - message.includes('unauthorized') || - message.includes('forbidden') - ); - } - async closeConnections(): Promise { const entries = Object.entries(this.mcpClients); const results = await Promise.allSettled(entries.map(([, client]) => client.close())); From 5dfd9338a3d028bbee03b928b2ae1beda432a6b3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Vincent=20Molini=C3=A9?= Date: Wed, 28 Jan 2026 16:27:32 +0100 Subject: [PATCH 4/5] test: fix tests --- packages/agent/src/routes/ai/ai-proxy.ts | 10 ++++++- .../agent/test/routes/ai/ai-proxy.test.ts | 18 ++++++++++++ packages/ai-proxy/src/mcp-client.ts | 2 +- packages/ai-proxy/test/mcp-client.test.ts | 28 +++++++++++++++++++ 4 files changed, 56 insertions(+), 2 deletions(-) diff --git a/packages/agent/src/routes/ai/ai-proxy.ts b/packages/agent/src/routes/ai/ai-proxy.ts index 0321fde3b..a933bc800 100644 --- a/packages/agent/src/routes/ai/ai-proxy.ts +++ b/packages/agent/src/routes/ai/ai-proxy.ts @@ -41,7 +41,15 @@ export default class AiProxyRoute extends BaseRoute { private async handleAiProxy(context: Context): Promise { try { const mcpOauthTokensHeader = context.request.headers['x-mcp-oauth-tokens'] as string; - const mcpOAuthTokens = mcpOauthTokensHeader ? JSON.parse(mcpOauthTokensHeader) : undefined; + let mcpOAuthTokens: Record | undefined; + + if (mcpOauthTokensHeader) { + try { + mcpOAuthTokens = JSON.parse(mcpOauthTokensHeader); + } catch { + throw new BadRequestError('Invalid JSON in x-mcp-oauth-tokens header'); + } + } context.response.body = await this.aiProxyRouter.route({ route: context.params.route, diff --git a/packages/agent/test/routes/ai/ai-proxy.test.ts b/packages/agent/test/routes/ai/ai-proxy.test.ts index 9a9d53959..d86975bd6 100644 --- a/packages/agent/test/routes/ai/ai-proxy.test.ts +++ b/packages/agent/test/routes/ai/ai-proxy.test.ts @@ -132,6 +132,24 @@ describe('AiProxyRoute', () => { ); }); + test('should throw BadRequestError when x-mcp-oauth-tokens header contains invalid JSON', async () => { + const route = new AiProxyRoute(services, options, aiConfigurations); + + const context = createMockContext({ + customProperties: { + params: { route: 'ai-query' }, + }, + requestBody: { messages: [] }, + headers: { 'x-mcp-oauth-tokens': '{ invalid json }' }, + }); + context.query = {}; + + await expect((route as any).handleAiProxy(context)).rejects.toThrow(BadRequestError); + await expect((route as any).handleAiProxy(context)).rejects.toThrow( + 'Invalid JSON in x-mcp-oauth-tokens header', + ); + }); + describe('error handling', () => { test('should convert AINotConfiguredError to UnprocessableError', async () => { const route = new AiProxyRoute(services, options, aiConfigurations); diff --git a/packages/ai-proxy/src/mcp-client.ts b/packages/ai-proxy/src/mcp-client.ts index eaf3a4981..981d7d303 100644 --- a/packages/ai-proxy/src/mcp-client.ts +++ b/packages/ai-proxy/src/mcp-client.ts @@ -46,7 +46,7 @@ export default class McpClient { if (!token) return serverConfig; // Only inject token for HTTP-based transports (sse, http) - if (serverConfig.type === 'http') { + if (serverConfig.type === 'http' || serverConfig.type === 'sse') { const { oauthConfig, ...headers } = serverConfig.headers || {}; return { diff --git a/packages/ai-proxy/test/mcp-client.test.ts b/packages/ai-proxy/test/mcp-client.test.ts index 5652af533..68f494437 100644 --- a/packages/ai-proxy/test/mcp-client.test.ts +++ b/packages/ai-proxy/test/mcp-client.test.ts @@ -273,6 +273,34 @@ describe('McpClient', () => { ); }); + it('should inject OAuth token as Authorization header into SSE type transport', () => { + const sseConfig: McpConfiguration = { + configs: { + remote: { + type: 'sse', + url: 'https://example.com/mcp', + }, + }, + }; + + // eslint-disable-next-line no-new + new McpClient(sseConfig, undefined, { remote: 'my-oauth-token' }); + + expect(MockedMultiServerMCPClient).toHaveBeenCalledWith( + expect.objectContaining({ + mcpServers: { + remote: { + type: 'sse', + url: 'https://example.com/mcp', + headers: { + Authorization: 'my-oauth-token', + }, + }, + }, + }), + ); + }); + it('should merge Authorization header with existing headers and strip oauthConfig', () => { const httpConfig: McpConfiguration = { configs: { From 3febacddd8f11e55461e9401386cc30ec204159a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Vincent=20Molini=C3=A9?= Date: Fri, 30 Jan 2026 10:42:10 +0100 Subject: [PATCH 5/5] refactor: move the injection of oauthTokens out of the ai-proxy package --- packages/agent/src/routes/ai/ai-proxy.ts | 7 +- .../agent/test/routes/ai/ai-proxy.test.ts | 29 +- packages/ai-proxy/src/index.ts | 7 +- packages/ai-proxy/src/mcp-client.ts | 81 +++--- packages/ai-proxy/src/router.ts | 16 +- .../ai-proxy/src/types/mcp-config-checker.ts | 4 +- packages/ai-proxy/test/mcp-client.test.ts | 260 ++++++++---------- packages/ai-proxy/test/router.test.ts | 21 -- 8 files changed, 207 insertions(+), 218 deletions(-) diff --git a/packages/agent/src/routes/ai/ai-proxy.ts b/packages/agent/src/routes/ai/ai-proxy.ts index a933bc800..581216a29 100644 --- a/packages/agent/src/routes/ai/ai-proxy.ts +++ b/packages/agent/src/routes/ai/ai-proxy.ts @@ -8,6 +8,7 @@ import { AIError, AINotFoundError, Router as AiProxyRouter, + injectOauthTokens, } from '@forestadmin/ai-proxy'; import { BadRequestError, @@ -51,12 +52,14 @@ export default class AiProxyRoute extends BaseRoute { } } + const mcpConfigs = + await this.options.forestAdminClient.mcpServerConfigService.getConfiguration(); + context.response.body = await this.aiProxyRouter.route({ route: context.params.route, body: context.request.body, query: context.query, - mcpConfigs: await this.options.forestAdminClient.mcpServerConfigService.getConfiguration(), - mcpOAuthTokens, + mcpConfigs: injectOauthTokens(mcpConfigs, mcpOAuthTokens), }); context.response.status = HttpCode.Ok; } catch (error) { diff --git a/packages/agent/test/routes/ai/ai-proxy.test.ts b/packages/agent/test/routes/ai/ai-proxy.test.ts index d86975bd6..0a7747e8b 100644 --- a/packages/agent/test/routes/ai/ai-proxy.test.ts +++ b/packages/agent/test/routes/ai/ai-proxy.test.ts @@ -105,14 +105,24 @@ describe('AiProxyRoute', () => { body: { messages: [{ role: 'user', content: 'Hello' }] }, query: { 'ai-name': 'gpt4' }, mcpConfigs: undefined, // mcpServerConfigService.getConfiguration returns undefined in test - mcpOAuthTokens: undefined, }); }); - test('should parse x-mcp-oauth-tokens header as JSON and pass to router', async () => { + test('should inject oauth tokens into mcpConfigs when header is provided', async () => { const route = new AiProxyRoute(services, options, aiConfigurations); mockRoute.mockResolvedValueOnce({}); + // Mock mcpServerConfigService to return actual configs + const mcpConfigs = { + configs: { + server1: { type: 'http' as const, url: 'https://server1.com' }, + server2: { type: 'http' as const, url: 'https://server2.com' }, + }, + }; + options.forestAdminClient.mcpServerConfigService.getConfiguration = jest + .fn() + .mockResolvedValue(mcpConfigs); + const tokens = { server1: 'token1', server2: 'token2' }; const context = createMockContext({ customProperties: { @@ -127,7 +137,20 @@ describe('AiProxyRoute', () => { expect(mockRoute).toHaveBeenCalledWith( expect.objectContaining({ - mcpOAuthTokens: tokens, + mcpConfigs: { + configs: { + server1: { + type: 'http', + url: 'https://server1.com', + headers: { Authorization: 'token1' }, + }, + server2: { + type: 'http', + url: 'https://server2.com', + headers: { Authorization: 'token2' }, + }, + }, + }, }), ); }); diff --git a/packages/ai-proxy/src/index.ts b/packages/ai-proxy/src/index.ts index d5915732d..2cc9c9892 100644 --- a/packages/ai-proxy/src/index.ts +++ b/packages/ai-proxy/src/index.ts @@ -9,9 +9,6 @@ export * from './mcp-client'; export * from './types/errors'; -export function validMcpConfigurationOrThrow( - mcpConfig: McpConfiguration, - mcpOauthTokens?: Record, -) { - return McpConfigChecker.check(mcpConfig, mcpOauthTokens); +export function validMcpConfigurationOrThrow(mcpConfig: McpConfiguration) { + return McpConfigChecker.check(mcpConfig); } diff --git a/packages/ai-proxy/src/mcp-client.ts b/packages/ai-proxy/src/mcp-client.ts index 981d7d303..b0211eea7 100644 --- a/packages/ai-proxy/src/mcp-client.ts +++ b/packages/ai-proxy/src/mcp-client.ts @@ -7,60 +7,75 @@ import McpServerRemoteTool from './types/mcp-server-remote-tool'; export type McpAuthenticationType = 'none' | 'bearer' | 'oauth2'; +export type McpServerConfig = MultiServerMCPClient['config']['mcpServers'][string]; + export type McpConfiguration = { configs: MultiServerMCPClient['config']['mcpServers']; authenticationTypes?: Record; } & Omit; +/** + * Injects the OAuth token as Authorization header into HTTP-based transport configurations. + * For stdio transports, returns the config unchanged. + */ +export function injectOauthToken(serverConfig: McpServerConfig, token?: string): McpServerConfig { + if (!token) return serverConfig; + + // Only inject token for HTTP-based transports (sse, http) + if (serverConfig.type === 'http' || serverConfig.type === 'sse') { + const { oauthConfig, ...headers } = serverConfig.headers || {}; + + return { + ...serverConfig, + headers: { + ...headers, + Authorization: token, + }, + }; + } + + return serverConfig; +} + +/** + * Injects OAuth tokens into all server configurations. + * Returns a new McpConfiguration with tokens injected, or undefined if no configs provided. + */ +export function injectOauthTokens( + mcpConfigs: McpConfiguration | undefined, + mcpOAuthTokens: Record | undefined, +): McpConfiguration | undefined { + if (!mcpConfigs) return undefined; + if (!mcpOAuthTokens) return mcpConfigs; + + const configsWithTokens = Object.fromEntries( + Object.entries(mcpConfigs.configs).map(([name, serverConfig]) => [ + name, + injectOauthToken(serverConfig, mcpOAuthTokens[name]), + ]), + ); + + return { ...mcpConfigs, configs: configsWithTokens }; +} + export default class McpClient { private readonly mcpClients: Record = {}; - private readonly authenticationTypes: Record; private readonly logger?: Logger; readonly tools: McpServerRemoteTool[] = []; - constructor(config: McpConfiguration, logger?: Logger, mcpOauthTokens?: Record) { + constructor(config: McpConfiguration, logger?: Logger) { this.logger = logger; - this.authenticationTypes = config.authenticationTypes ?? {}; // 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 token = mcpOauthTokens?.[name]; - const configWithToken = this.injectOauthToken(serverConfig, token); this.mcpClients[name] = new MultiServerMCPClient({ - mcpServers: { [name]: configWithToken }, + mcpServers: { [name]: serverConfig }, ...config, }); }); } - /** - * Injects the OAuth token as Authorization header into HTTP-based transport configurations. - * The frontend sends 'x-mcp-oauth-tokens' with tokens keyed by server name. - * For stdio transports, returns the config unchanged. - */ - private injectOauthToken( - serverConfig: MultiServerMCPClient['config']['mcpServers'][string], - token?: string, - ): MultiServerMCPClient['config']['mcpServers'][string] { - if (!token) return serverConfig; - - // Only inject token for HTTP-based transports (sse, http) - if (serverConfig.type === 'http' || serverConfig.type === 'sse') { - const { oauthConfig, ...headers } = serverConfig.headers || {}; - - return { - ...serverConfig, - headers: { - ...headers, - Authorization: token, - }, - }; - } - - return serverConfig; - } - async loadTools(): Promise { const errors: Array<{ server: string; error: Error }> = []; diff --git a/packages/ai-proxy/src/router.ts b/packages/ai-proxy/src/router.ts index 8d16ba9d0..049d760a2 100644 --- a/packages/ai-proxy/src/router.ts +++ b/packages/ai-proxy/src/router.ts @@ -16,12 +16,6 @@ export type Query = { }; export type ApiKeys = RemoteToolsApiKeys; -/** Headers that can be passed to the route method */ -export type RouteHeaders = { - /** OAuth tokens for MCP server authentication, keyed by server name */ - mcpOauthTokens?: Record; -}; - export class Router { private readonly localToolsApiKeys?: ApiKeys; private readonly aiConfigurations: AiConfiguration[]; @@ -67,18 +61,12 @@ export class Router { * - invoke-remote-tool: Execute a remote tool by name with the provided inputs * - remote-tools: Return the list of available remote tools definitions */ - async route(args: { - body?: Body; - route: Route; - query?: Query; - mcpConfigs?: McpConfiguration; - mcpOAuthTokens?: Record; - }) { + async route(args: { body?: Body; route: Route; query?: Query; mcpConfigs?: McpConfiguration }) { let mcpClient: McpClient | undefined; try { if (args.mcpConfigs) { - mcpClient = new McpClient(args.mcpConfigs, this.logger, args.mcpOAuthTokens); + mcpClient = new McpClient(args.mcpConfigs, this.logger); } const remoteTools = new RemoteTools( diff --git a/packages/ai-proxy/src/types/mcp-config-checker.ts b/packages/ai-proxy/src/types/mcp-config-checker.ts index 7ce26bf7b..1474354ac 100644 --- a/packages/ai-proxy/src/types/mcp-config-checker.ts +++ b/packages/ai-proxy/src/types/mcp-config-checker.ts @@ -3,7 +3,7 @@ import type { McpConfiguration } from '../mcp-client'; import McpClient from '../mcp-client'; export default class McpConfigChecker { - static check(mcpConfig: McpConfiguration, mcpOauthTokens?: Record) { - return new McpClient(mcpConfig, undefined, mcpOauthTokens).testConnections(); + static check(mcpConfig: McpConfiguration) { + return new McpClient(mcpConfig).testConnections(); } } diff --git a/packages/ai-proxy/test/mcp-client.test.ts b/packages/ai-proxy/test/mcp-client.test.ts index 68f494437..bc70fbb73 100644 --- a/packages/ai-proxy/test/mcp-client.test.ts +++ b/packages/ai-proxy/test/mcp-client.test.ts @@ -3,7 +3,7 @@ import type { McpConfiguration } from '../src'; import { tool } from '@langchain/core/tools'; import { McpConnectionError } from '../src'; -import McpClient from '../src/mcp-client'; +import McpClient, { injectOauthToken, injectOauthTokens } from '../src/mcp-client'; import McpServerRemoteTool from '../src/types/mcp-server-remote-tool'; const getToolsMock = jest.fn(); @@ -244,182 +244,166 @@ describe('McpClient', () => { }); }); - describe('OAuth token injection', () => { + describe('injectOauthToken', () => { it('should inject OAuth token as Authorization header into HTTP type transport', () => { - const httpConfig: McpConfiguration = { - configs: { - remote: { - type: 'http', - url: 'https://example.com/mcp', - }, - }, + const serverConfig = { + type: 'http' as const, + url: 'https://example.com/mcp', }; - // eslint-disable-next-line no-new - new McpClient(httpConfig, undefined, { remote: 'my-oauth-token' }); - - expect(MockedMultiServerMCPClient).toHaveBeenCalledWith( - expect.objectContaining({ - mcpServers: { - remote: { - type: 'http', - url: 'https://example.com/mcp', - headers: { - Authorization: 'my-oauth-token', - }, - }, - }, - }), - ); + const result = injectOauthToken(serverConfig, 'my-oauth-token'); + + expect(result).toEqual({ + type: 'http', + url: 'https://example.com/mcp', + headers: { + Authorization: 'my-oauth-token', + }, + }); }); it('should inject OAuth token as Authorization header into SSE type transport', () => { - const sseConfig: McpConfiguration = { - configs: { - remote: { - type: 'sse', - url: 'https://example.com/mcp', - }, - }, + const serverConfig = { + type: 'sse' as const, + url: 'https://example.com/mcp', }; - // eslint-disable-next-line no-new - new McpClient(sseConfig, undefined, { remote: 'my-oauth-token' }); - - expect(MockedMultiServerMCPClient).toHaveBeenCalledWith( - expect.objectContaining({ - mcpServers: { - remote: { - type: 'sse', - url: 'https://example.com/mcp', - headers: { - Authorization: 'my-oauth-token', - }, - }, - }, - }), - ); + const result = injectOauthToken(serverConfig, 'my-oauth-token'); + + expect(result).toEqual({ + type: 'sse', + url: 'https://example.com/mcp', + headers: { + Authorization: 'my-oauth-token', + }, + }); }); it('should merge Authorization header with existing headers and strip oauthConfig', () => { - const httpConfig: McpConfiguration = { - configs: { - remote: { - type: 'http', - url: 'https://example.com/mcp', - headers: { - 'x-custom-header': 'custom-value', - oauthConfig: { clientId: 'test' } as unknown as string, - }, - }, + const serverConfig = { + type: 'http' as const, + url: 'https://example.com/mcp', + headers: { + 'x-custom-header': 'custom-value', + oauthConfig: { clientId: 'test' } as unknown as string, }, }; - // eslint-disable-next-line no-new - new McpClient(httpConfig, undefined, { remote: 'my-oauth-token' }); - - expect(MockedMultiServerMCPClient).toHaveBeenCalledWith( - expect.objectContaining({ - mcpServers: { - remote: { - type: 'http', - url: 'https://example.com/mcp', - headers: { - 'x-custom-header': 'custom-value', - Authorization: 'my-oauth-token', - }, - }, - }, - }), - ); + const result = injectOauthToken(serverConfig, 'my-oauth-token'); + + expect(result).toEqual({ + type: 'http', + url: 'https://example.com/mcp', + headers: { + 'x-custom-header': 'custom-value', + Authorization: 'my-oauth-token', + }, + }); }); it('should not inject OAuth token into stdio transport even if token provided', () => { - // eslint-disable-next-line no-new - new McpClient(aConfig, undefined, { slack: 'my-oauth-token' }); + const serverConfig = { + transport: 'stdio' as const, + command: 'npx', + args: ['-y', '@modelcontextprotocol/server-slack'], + env: {}, + }; - expect(MockedMultiServerMCPClient).toHaveBeenCalledWith( - expect.objectContaining({ - mcpServers: { - slack: { - transport: 'stdio', - command: 'npx', - args: ['-y', '@modelcontextprotocol/server-slack'], - env: {}, - }, - }, - }), - ); + const result = injectOauthToken(serverConfig, 'my-oauth-token'); + + expect(result).toEqual(serverConfig); + }); + + it('should not modify config when no token is provided', () => { + const serverConfig = { + type: 'http' as const, + url: 'https://example.com/mcp', + }; + + const result = injectOauthToken(serverConfig); + + expect(result).toEqual(serverConfig); + }); + + it('should return same reference when no token is provided', () => { + const serverConfig = { + type: 'http' as const, + url: 'https://example.com/mcp', + }; + + const result = injectOauthToken(serverConfig); + + expect(result).toBe(serverConfig); }); + }); + + describe('injectOauthTokens', () => { + it('should inject tokens into all matching server configs', () => { + const mcpConfigs = { + configs: { + server1: { type: 'http' as const, url: 'https://server1.com' }, + server2: { type: 'http' as const, url: 'https://server2.com' }, + }, + }; + const tokens = { server1: 'token1', server2: 'token2' }; + + const result = injectOauthTokens(mcpConfigs, tokens); - it('should only inject token for servers that have a matching token', () => { - const multiServerConfig: McpConfiguration = { + expect(result).toEqual({ configs: { server1: { type: 'http', - url: 'https://server1.com/mcp', + url: 'https://server1.com', + headers: { Authorization: 'token1' }, }, server2: { type: 'http', - url: 'https://server2.com/mcp', + url: 'https://server2.com', + headers: { Authorization: 'token2' }, }, }, + }); + }); + + it('should only inject tokens for servers that have matching tokens', () => { + const mcpConfigs = { + configs: { + server1: { type: 'http' as const, url: 'https://server1.com' }, + server2: { type: 'http' as const, url: 'https://server2.com' }, + }, }; + const tokens = { server1: 'token1' }; - // eslint-disable-next-line no-new - new McpClient(multiServerConfig, undefined, { server1: 'token-for-server1' }); - - // server1 should have the token - expect(MockedMultiServerMCPClient).toHaveBeenCalledWith( - expect.objectContaining({ - mcpServers: { - server1: { - type: 'http', - url: 'https://server1.com/mcp', - headers: { - Authorization: 'token-for-server1', - }, - }, - }, - }), - ); + const result = injectOauthTokens(mcpConfigs, tokens); - // server2 should not have any headers injected - expect(MockedMultiServerMCPClient).toHaveBeenCalledWith( - expect.objectContaining({ - mcpServers: { - server2: { - type: 'http', - url: 'https://server2.com/mcp', - }, + expect(result).toEqual({ + configs: { + server1: { + type: 'http', + url: 'https://server1.com', + headers: { Authorization: 'token1' }, }, - }), - ); + server2: { type: 'http', url: 'https://server2.com' }, + }, + }); + }); + + it('should return undefined when mcpConfigs is undefined', () => { + const result = injectOauthTokens(undefined, { server1: 'token1' }); + + expect(result).toBeUndefined(); }); - it('should not modify config when no OAuth tokens are provided', () => { - const httpConfig: McpConfiguration = { + it('should return mcpConfigs unchanged when tokens is undefined', () => { + const mcpConfigs = { configs: { - remote: { - type: 'http', - url: 'https://example.com/mcp', - }, + server1: { type: 'http' as const, url: 'https://server1.com' }, }, }; - // eslint-disable-next-line no-new - new McpClient(httpConfig); + const result = injectOauthTokens(mcpConfigs, undefined); - expect(MockedMultiServerMCPClient).toHaveBeenCalledWith( - expect.objectContaining({ - mcpServers: { - remote: { - type: 'http', - url: 'https://example.com/mcp', - }, - }, - }), - ); + expect(result).toBe(mcpConfigs); }); }); }); diff --git a/packages/ai-proxy/test/router.test.ts b/packages/ai-proxy/test/router.test.ts index 2606d3aaa..a5800e3a7 100644 --- a/packages/ai-proxy/test/router.test.ts +++ b/packages/ai-proxy/test/router.test.ts @@ -362,27 +362,6 @@ describe('route', () => { expect(MockedMcpClient).toHaveBeenCalledWith( { configs: { server1: { command: 'test', args: [] } } }, customLogger, - undefined, - ); - }); - - it('passes mcpOAuthTokens to McpClient when provided', async () => { - const customLogger: Logger = jest.fn(); - const router = new Router({ - logger: customLogger, - }); - - const tokens = { server1: 'token-for-server1' }; - await router.route({ - route: 'remote-tools', - mcpConfigs: { configs: { server1: { command: 'test', args: [] } } }, - mcpOAuthTokens: tokens, - }); - - expect(MockedMcpClient).toHaveBeenCalledWith( - { configs: { server1: { command: 'test', args: [] } } }, - customLogger, - tokens, ); }); });