From c0e876a2e6f6346b76604587b5d4ef1e74ca9ad8 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 10 Mar 2026 16:19:07 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=84=20refactor:=20OAuth=20Metadata=20D?= =?UTF-8?q?iscovery=20with=20Origin=20Fallback=20(#12170)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * ๐Ÿ”„ refactor: OAuth Metadata Discovery with Origin Fallback Updated the `discoverWithOriginFallback` method to improve the handling of OAuth authorization server metadata discovery. The method now retries with the origin URL when discovery fails for a path-based URL, ensuring consistent behavior across `discoverMetadata` and token refresh flows. This change reduces code duplication and enhances the reliability of the OAuth flow by providing a unified implementation for origin fallback logic. * ๐Ÿงช test: Add tests for OAuth Token Refresh with Origin Fallback Introduced new tests for the `refreshOAuthTokens` method in `MCPOAuthHandler` to validate the retry mechanism with the origin URL when path-based discovery fails. The tests cover scenarios where the first discovery attempt throws an error and the subsequent attempt succeeds, as well as cases where the discovery fails entirely. This enhances the reliability of the OAuth token refresh process by ensuring proper handling of discovery failures. * chore: imports order * fix: Improve Base URL Logging and Metadata Discovery in MCPOAuthHandler Updated the logging to use a consistent base URL object when handling discovery failures in the MCPOAuthHandler. This change enhances error reporting by ensuring that the base URL is logged correctly, and it refines the metadata discovery process by returning the result of the discovery attempt with the base URL, improving the reliability of the OAuth flow. --- .../api/src/mcp/__tests__/handler.test.ts | 141 +++++++++++++++++- packages/api/src/mcp/oauth/handler.ts | 22 ++- .../MCPReinitRecovery.integration.test.ts | 6 +- 3 files changed, 153 insertions(+), 16 deletions(-) diff --git a/packages/api/src/mcp/__tests__/handler.test.ts b/packages/api/src/mcp/__tests__/handler.test.ts index e5d94b23e3..3b68d88e9c 100644 --- a/packages/api/src/mcp/__tests__/handler.test.ts +++ b/packages/api/src/mcp/__tests__/handler.test.ts @@ -1498,20 +1498,19 @@ describe('MCPOAuthHandler - Configurable OAuth Metadata', () => { ); const firstDiscoveryUrl = mockDiscoverAuthorizationServerMetadata.mock.calls[0][0] as URL; const secondDiscoveryUrl = mockDiscoverAuthorizationServerMetadata.mock.calls[1][0] as URL; - expect(firstDiscoveryUrl).toBeInstanceOf(URL); expect(firstDiscoveryUrl.href).toBe('https://mcp.sentry.dev/mcp'); - expect(secondDiscoveryUrl).toBeInstanceOf(URL); expect(secondDiscoveryUrl.href).toBe('https://mcp.sentry.dev/'); - // Token endpoint from origin discovery metadata is used + // Token endpoint from origin discovery metadata is used (string in stored-clientInfo branch) expect(mockFetch).toHaveBeenCalled(); const [fetchUrl, fetchOptions] = mockFetch.mock.calls[0]; - expect(String(fetchUrl)).toBe('https://mcp.sentry.dev/oauth/token'); + expect(typeof fetchUrl).toBe('string'); + expect(fetchUrl).toBe('https://mcp.sentry.dev/oauth/token'); expect(fetchOptions).toEqual(expect.objectContaining({ method: 'POST' })); expect(result.access_token).toBe('new-access-token'); }); - it('retries with origin URL when path-based discovery fails (auto-discovered path)', async () => { + it('retries with origin URL when path-based discovery fails (no stored clientInfo)', async () => { // No clientInfo โ€” uses the auto-discovered branch const metadata = { serverName: 'sentry', @@ -1563,12 +1562,10 @@ describe('MCPOAuthHandler - Configurable OAuth Metadata', () => { ); const firstDiscoveryUrl = mockDiscoverAuthorizationServerMetadata.mock.calls[0][0] as URL; const secondDiscoveryUrl = mockDiscoverAuthorizationServerMetadata.mock.calls[1][0] as URL; - expect(firstDiscoveryUrl).toBeInstanceOf(URL); expect(firstDiscoveryUrl.href).toBe('https://mcp.sentry.dev/mcp'); - expect(secondDiscoveryUrl).toBeInstanceOf(URL); expect(secondDiscoveryUrl.href).toBe('https://mcp.sentry.dev/'); - // Token endpoint from origin discovery metadata is used + // Token endpoint from origin discovery metadata is used (URL object in auto-discovered branch) expect(mockFetch).toHaveBeenCalled(); const [fetchUrl, fetchOptions] = mockFetch.mock.calls[0]; expect(fetchUrl).toBeInstanceOf(URL); @@ -1577,6 +1574,46 @@ describe('MCPOAuthHandler - Configurable OAuth Metadata', () => { expect(result.access_token).toBe('new-access-token'); }); + it('falls back to /token when both path and origin discovery fail', async () => { + const metadata = { + serverName: 'sentry', + serverUrl: 'https://mcp.sentry.dev/mcp', + clientInfo: { + client_id: 'test-client-id', + client_secret: 'test-client-secret', + grant_types: ['authorization_code', 'refresh_token'], + }, + }; + + // Both path AND origin discovery return undefined + mockDiscoverAuthorizationServerMetadata + .mockResolvedValueOnce(undefined) + .mockResolvedValueOnce(undefined); + + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => ({ + access_token: 'new-access-token', + refresh_token: 'new-refresh-token', + expires_in: 3600, + }), + } as Response); + + const result = await MCPOAuthHandler.refreshOAuthTokens( + 'test-refresh-token', + metadata, + {}, + {}, + ); + + expect(mockDiscoverAuthorizationServerMetadata).toHaveBeenCalledTimes(2); + + // Falls back to /token relative to server URL origin + const [fetchUrl] = mockFetch.mock.calls[0]; + expect(String(fetchUrl)).toBe('https://mcp.sentry.dev/token'); + expect(result.access_token).toBe('new-access-token'); + }); + it('does not retry with origin when server URL has no path (root URL)', async () => { const metadata = { serverName: 'test-server', @@ -1600,6 +1637,94 @@ describe('MCPOAuthHandler - Configurable OAuth Metadata', () => { // Only one discovery attempt for a root URL expect(mockDiscoverAuthorizationServerMetadata).toHaveBeenCalledTimes(1); }); + + it('retries with origin when path-based discovery throws', async () => { + const metadata = { + serverName: 'sentry', + serverUrl: 'https://mcp.sentry.dev/mcp', + clientInfo: { + client_id: 'test-client-id', + client_secret: 'test-client-secret', + grant_types: ['authorization_code', 'refresh_token'], + }, + }; + + const originMetadata = { + issuer: 'https://mcp.sentry.dev/', + authorization_endpoint: 'https://mcp.sentry.dev/oauth/authorize', + token_endpoint: 'https://mcp.sentry.dev/oauth/token', + token_endpoint_auth_methods_supported: ['client_secret_post'], + response_types_supported: ['code'], + } as AuthorizationServerMetadata; + + // First call throws, second call succeeds + mockDiscoverAuthorizationServerMetadata + .mockRejectedValueOnce(new Error('Network error')) + .mockResolvedValueOnce(originMetadata); + + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => ({ + access_token: 'new-access-token', + refresh_token: 'new-refresh-token', + expires_in: 3600, + }), + } as Response); + + const result = await MCPOAuthHandler.refreshOAuthTokens( + 'test-refresh-token', + metadata, + {}, + {}, + ); + + expect(mockDiscoverAuthorizationServerMetadata).toHaveBeenCalledTimes(2); + const [fetchUrl] = mockFetch.mock.calls[0]; + expect(String(fetchUrl)).toBe('https://mcp.sentry.dev/oauth/token'); + expect(result.access_token).toBe('new-access-token'); + }); + + it('propagates the throw when root URL discovery throws', async () => { + const metadata = { + serverName: 'test-server', + serverUrl: 'https://auth.example.com/', + clientInfo: { + client_id: 'test-client-id', + client_secret: 'test-client-secret', + }, + }; + + mockDiscoverAuthorizationServerMetadata.mockRejectedValueOnce( + new Error('Discovery failed'), + ); + + await expect( + MCPOAuthHandler.refreshOAuthTokens('test-refresh-token', metadata, {}, {}), + ).rejects.toThrow('Discovery failed'); + + expect(mockDiscoverAuthorizationServerMetadata).toHaveBeenCalledTimes(1); + }); + + it('propagates the throw when both path and origin discovery throw', async () => { + const metadata = { + serverName: 'sentry', + serverUrl: 'https://mcp.sentry.dev/mcp', + clientInfo: { + client_id: 'test-client-id', + client_secret: 'test-client-secret', + }, + }; + + mockDiscoverAuthorizationServerMetadata + .mockRejectedValueOnce(new Error('Network error')) + .mockRejectedValueOnce(new Error('Origin also failed')); + + await expect( + MCPOAuthHandler.refreshOAuthTokens('test-refresh-token', metadata, {}, {}), + ).rejects.toThrow('Origin also failed'); + + expect(mockDiscoverAuthorizationServerMetadata).toHaveBeenCalledTimes(2); + }); }); }); }); diff --git a/packages/api/src/mcp/oauth/handler.ts b/packages/api/src/mcp/oauth/handler.ts index 6ef444bf47..83e855591e 100644 --- a/packages/api/src/mcp/oauth/handler.ts +++ b/packages/api/src/mcp/oauth/handler.ts @@ -209,16 +209,28 @@ export class MCPOAuthHandler { } /** - * Discovers OAuth authorization server metadata with origin-URL fallback. - * If discovery fails for a path-based URL, retries with just the origin. - * Mirrors the fallback behavior in `discoverMetadata` and `initiateOAuthFlow`. + * Discovers OAuth authorization server metadata, retrying with just the origin + * when discovery fails for a path-based URL. Shared implementation used by + * `discoverMetadata` and both `refreshOAuthTokens` branches. */ private static async discoverWithOriginFallback( serverUrl: URL, fetchFn: FetchLike, ): ReturnType { - const metadata = await discoverAuthorizationServerMetadata(serverUrl, { fetchFn }); - // If discovery failed and we're using a path-based URL, try the base URL + let metadata: Awaited>; + try { + metadata = await discoverAuthorizationServerMetadata(serverUrl, { fetchFn }); + } catch (err) { + if (serverUrl.pathname === '/') { + throw err; + } + const baseUrl = new URL(serverUrl.origin); + logger.debug( + `[MCPOAuth] Discovery threw for path URL, trying base URL: ${sanitizeUrlForLogging(baseUrl)}`, + { error: err }, + ); + return discoverAuthorizationServerMetadata(baseUrl, { fetchFn }); + } if (!metadata && serverUrl.pathname !== '/') { const baseUrl = new URL(serverUrl.origin); logger.debug( diff --git a/packages/api/src/mcp/registry/__tests__/MCPReinitRecovery.integration.test.ts b/packages/api/src/mcp/registry/__tests__/MCPReinitRecovery.integration.test.ts index b171e84d13..9545486fde 100644 --- a/packages/api/src/mcp/registry/__tests__/MCPReinitRecovery.integration.test.ts +++ b/packages/api/src/mcp/registry/__tests__/MCPReinitRecovery.integration.test.ts @@ -17,20 +17,20 @@ import * as net from 'net'; import * as http from 'http'; +import { Keyv } from 'keyv'; import { Agent } from 'undici'; +import { Types } from 'mongoose'; import { randomUUID } from 'crypto'; import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js'; -import { Keyv } from 'keyv'; -import { Types } from 'mongoose'; import type { IUser } from '@librechat/data-schemas'; import type { Socket } from 'net'; import type * as t from '~/mcp/types'; -import { MCPInspectionFailedError } from '~/mcp/errors'; import { registryStatusCache } from '~/mcp/registry/cache/RegistryStatusCache'; import { MCPServersInitializer } from '~/mcp/registry/MCPServersInitializer'; import { MCPServersRegistry } from '~/mcp/registry/MCPServersRegistry'; import { ConnectionsRepository } from '~/mcp/ConnectionsRepository'; +import { MCPInspectionFailedError } from '~/mcp/errors'; import { FlowStateManager } from '~/flow/manager'; import { MCPConnection } from '~/mcp/connection'; import { MCPManager } from '~/mcp/MCPManager';