diff --git a/api/server/services/Tools/mcp.spec.js b/api/server/services/Tools/mcp.spec.js index 31d8a02195..8c0d942216 100644 --- a/api/server/services/Tools/mcp.spec.js +++ b/api/server/services/Tools/mcp.spec.js @@ -139,12 +139,13 @@ describe('loadMCPServerCatalogs', () => { await loadMCPServerCatalogs({ user, servers }); await loadMCPServerCatalogs({ user, servers }); + const identity = `${user.id}:offline:generation`; expect(cooldowns[0]).toBeDefined(); expect(cooldowns[1]).toBe(cooldowns[0]); - expect(cooldowns[0].isCoolingDown(user.id, 'offline')).toBe(false); - cooldowns[0].recordFailure(user.id, 'offline'); - expect(cooldowns[0].isCoolingDown(user.id, 'offline')).toBe(true); - cooldowns[0].recordSuccess(user.id, 'offline'); + expect(cooldowns[0].isCoolingDown(identity)).toBe(false); + cooldowns[0].recordFailure(identity); + expect(cooldowns[0].isCoolingDown(identity)).toBe(true); + cooldowns[0].recordSuccess(identity); }); }); diff --git a/packages/api/src/mcp/catalog/recovery.spec.ts b/packages/api/src/mcp/catalog/recovery.spec.ts index 1844e88b01..6d3e4f5563 100644 --- a/packages/api/src/mcp/catalog/recovery.spec.ts +++ b/packages/api/src/mcp/catalog/recovery.spec.ts @@ -232,7 +232,7 @@ describe('recoverMCPServerCatalogs — bounded, skippable discovery', () => { formatServerTools: jest.fn().mockReturnValue({}), }); - it('caps the discovery timeout so an unreachable server cannot hold the request', async () => { + it('bounds each connection attempt the factory makes', async () => { const discoverServerTools = jest.fn().mockResolvedValue({ tools: [] }); await recoverMCPServerCatalogs( @@ -379,27 +379,28 @@ describe('createMCPCatalogRecoveryCooldown', () => { jest.useFakeTimers(); try { const cooldown = createMCPCatalogRecoveryCooldown(60_000); - cooldown.recordFailure('user-1', 'alpha'); + cooldown.recordFailure('user-1:alpha:gen-1'); - expect(cooldown.isCoolingDown('user-1', 'alpha')).toBe(true); + expect(cooldown.isCoolingDown('user-1:alpha:gen-1')).toBe(true); jest.advanceTimersByTime(59_000); - expect(cooldown.isCoolingDown('user-1', 'alpha')).toBe(true); + expect(cooldown.isCoolingDown('user-1:alpha:gen-1')).toBe(true); jest.advanceTimersByTime(2_000); - expect(cooldown.isCoolingDown('user-1', 'alpha')).toBe(false); + expect(cooldown.isCoolingDown('user-1:alpha:gen-1')).toBe(false); } finally { jest.useRealTimers(); } }); - it('scopes failures per user and per server, and clears them on success', () => { + it('keeps distinct identities apart and clears one on success', () => { const cooldown = createMCPCatalogRecoveryCooldown(60_000); - cooldown.recordFailure('user-1', 'alpha'); + cooldown.recordFailure('user-1:alpha:gen-1'); - expect(cooldown.isCoolingDown('user-1', 'beta')).toBe(false); - expect(cooldown.isCoolingDown('user-2', 'alpha')).toBe(false); + expect(cooldown.isCoolingDown('user-1:beta:gen-1')).toBe(false); + expect(cooldown.isCoolingDown('user-2:alpha:gen-1')).toBe(false); + expect(cooldown.isCoolingDown('user-1:alpha:gen-2')).toBe(false); - cooldown.recordSuccess('user-1', 'alpha'); - expect(cooldown.isCoolingDown('user-1', 'alpha')).toBe(false); + cooldown.recordSuccess('user-1:alpha:gen-1'); + expect(cooldown.isCoolingDown('user-1:alpha:gen-1')).toBe(false); }); it('sweeps expired entries so the failure map cannot grow without bound', () => { @@ -407,17 +408,102 @@ describe('createMCPCatalogRecoveryCooldown', () => { try { const cooldown = createMCPCatalogRecoveryCooldown(60_000); for (let i = 0; i < 50; i++) { - cooldown.recordFailure(`user-${i}`, 'alpha'); + cooldown.recordFailure(`user-${i}:alpha:gen-1`); } jest.advanceTimersByTime(61_000); - cooldown.recordFailure('user-current', 'alpha'); + cooldown.recordFailure('user-current:alpha:gen-1'); - expect(cooldown.isCoolingDown('user-current', 'alpha')).toBe(true); + expect(cooldown.isCoolingDown('user-current:alpha:gen-1')).toBe(true); for (let i = 0; i < 50; i++) { - expect(cooldown.isCoolingDown(`user-${i}`, 'alpha')).toBe(false); + expect(cooldown.isCoolingDown(`user-${i}:alpha:gen-1`)).toBe(false); } } finally { jest.useRealTimers(); } }); }); + +describe('recoverMCPServerCatalogs — request-level bounds', () => { + const hangingDeps = (discoverServerTools: jest.Mock) => ({ + loadUserMCPAuthMap: jest.fn().mockResolvedValue({}), + discoverServerTools, + formatServerTools: jest.fn().mockReturnValue({}), + }); + + it('gives up on a server that outlives the deadline, however long the factory waits', async () => { + jest.useFakeTimers(); + try { + /** Mirrors MCPConnectionFactory spending `connectionTimeout` on the authenticated attempt + * and again on the unauthenticated one: the deadline must bound the pair, not each. */ + const discoverServerTools = jest.fn( + ({ connectionTimeout }: ToolDiscoveryOptions) => + new Promise<{ tools: null }>((resolve) => + setTimeout(() => resolve({ tools: null }), (connectionTimeout ?? 0) * 2), + ), + ); + const servers = [{ serverName: 'unreachable', serverConfig: serverConfig('unreachable') }]; + + const pending = recoverMCPServerCatalogs({ user, servers }, hangingDeps(discoverServerTools)); + /** Only the 5s deadline can settle this; the discovery itself resolves at 10s. */ + await jest.advanceTimersByTimeAsync(5_000); + const result = await pending; + + expect(discoverServerTools).toHaveBeenCalledWith( + expect.objectContaining({ connectionTimeout: 5000 }), + ); + expect(result.size).toBe(0); + await jest.advanceTimersByTimeAsync(10_000); + } finally { + jest.useRealTimers(); + } + }); + + it('stops dialing once the request budget is spent and records no cooldown for the untried', async () => { + jest.useFakeTimers(); + try { + const recoveryCooldown = createMCPCatalogRecoveryCooldown(60_000); + const recordFailure = jest.spyOn(recoveryCooldown, 'recordFailure'); + /** Each discovery consumes the full per-server deadline, so the 10s budget funds two waves. */ + const discoverServerTools = jest.fn( + () => + new Promise<{ tools: null }>((resolve) => + setTimeout(() => resolve({ tools: null }), 5000), + ), + ); + const servers = Array.from({ length: 9 }, (_, index) => ({ + serverName: `server-${index}`, + serverConfig: serverConfig(`server-${index}`), + })); + + const pending = recoverMCPServerCatalogs( + { user, servers }, + { ...hangingDeps(discoverServerTools), recoveryCooldown }, + ); + await jest.advanceTimersByTimeAsync(30_000); + const result = await pending; + + expect(discoverServerTools).toHaveBeenCalledTimes(6); + expect(recordFailure).toHaveBeenCalledTimes(6); + expect(result.size).toBe(0); + } finally { + jest.useRealTimers(); + } + }); + + it('retries immediately once a failed server’s configuration changes', async () => { + const recoveryCooldown = createMCPCatalogRecoveryCooldown(60_000); + const discoverServerTools = jest.fn().mockResolvedValue({ tools: null }); + const deps = { ...hangingDeps(discoverServerTools), recoveryCooldown }; + const broken = { serverName: 'edited', serverConfig: serverConfig('typo') }; + const corrected = { serverName: 'edited', serverConfig: serverConfig('fixed') }; + + await recoverMCPServerCatalogs({ user, servers: [broken] }, deps); + await recoverMCPServerCatalogs({ user, servers: [broken] }, deps); + await recoverMCPServerCatalogs({ user, servers: [corrected] }, deps); + + expect(discoverServerTools).toHaveBeenCalledTimes(2); + expect(discoverServerTools).toHaveBeenLastCalledWith( + expect.objectContaining({ configServers: { edited: corrected.serverConfig } }), + ); + }); +}); diff --git a/packages/api/src/mcp/catalog/recovery.ts b/packages/api/src/mcp/catalog/recovery.ts index f076a29a20..19842f326b 100644 --- a/packages/api/src/mcp/catalog/recovery.ts +++ b/packages/api/src/mcp/catalog/recovery.ts @@ -2,17 +2,21 @@ import { logger } from '@librechat/data-schemas'; import type { Tool } from '@modelcontextprotocol/sdk/types.js'; import type { IUser } from '@librechat/data-schemas'; import type { LCAvailableTools, ParsedServerConfig, ToolDiscoveryOptions } from '../types'; -import { createConcurrencyLimiter } from '~/utils/promise'; +import { createConcurrencyLimiter, withTimeout } from '~/utils/promise'; +import { getMCPAppToolsPublicationGeneration } from '../toolsChanged'; import { getMissingCustomUserVars } from '../utils'; import { getServerCustomUserVars } from '../auth'; const RECOVERY_CONCURRENCY = 3; /** - * Passive recovery runs inline on a catalog list request, so an unreachable server must not - * hold the response for the connection default (`initTimeout ?? 30s`). A server configured to - * connect faster keeps its own shorter limit. + * Wall-clock deadline for one server's discovery, enforced here rather than by tuning + * `connectionTimeout`: the factory spends that per connection attempt and makes several + * (authenticated, then unauthenticated), so a per-attempt value bounds no total this layer + * can reason about. The deadline holds however many attempts the factory grows. */ -const RECOVERY_TIMEOUT_MS = 5000; +const RECOVERY_SERVER_DEADLINE_MS = 5000; +/** Ceiling on the recovery a single list request performs, whatever the server count. */ +const RECOVERY_REQUEST_BUDGET_MS = 10_000; /** How long a server that failed passive recovery is skipped before it is dialed again. */ const RECOVERY_COOLDOWN_MS = 60 * 1000; @@ -22,13 +26,14 @@ export interface MCPServerCatalogRecoveryInput { } /** - * Per-process record of servers whose passive recovery just failed. Recovered catalogs are - * request-local, so without it every list request re-dials the same unreachable servers. + * Per-process record of recoveries that just failed, keyed by an opaque identity the caller + * builds. Recovered catalogs are request-local, so without this every list request re-dials the + * same unreachable servers. */ export interface MCPCatalogRecoveryCooldown { - isCoolingDown: (userId: string, serverName: string) => boolean; - recordFailure: (userId: string, serverName: string) => void; - recordSuccess: (userId: string, serverName: string) => void; + isCoolingDown: (identity: string) => boolean; + recordFailure: (identity: string) => void; + recordSuccess: (identity: string) => void; } export interface MCPServerCatalogRecoveryDeps { @@ -81,7 +86,6 @@ export function createMCPCatalogRecoveryCooldown( cooldownMs: number = RECOVERY_COOLDOWN_MS, ): MCPCatalogRecoveryCooldown { const failedAt = new Map(); - const cooldownKey = (userId: string, serverName: string): string => `${userId}:${serverName}`; let lastSweptAt = 0; const sweep = (now: number): void => { @@ -89,35 +93,61 @@ export function createMCPCatalogRecoveryCooldown( return; } lastSweptAt = now; - for (const [key, timestamp] of failedAt) { + for (const [identity, timestamp] of failedAt) { if (now - timestamp >= cooldownMs) { - failedAt.delete(key); + failedAt.delete(identity); } } }; return { - isCoolingDown: (userId, serverName) => { + isCoolingDown: (identity) => { const now = Date.now(); sweep(now); - const timestamp = failedAt.get(cooldownKey(userId, serverName)); + const timestamp = failedAt.get(identity); return timestamp != null && now - timestamp < cooldownMs; }, - recordFailure: (userId, serverName) => { - failedAt.set(cooldownKey(userId, serverName), Date.now()); + recordFailure: (identity) => { + failedAt.set(identity, Date.now()); }, - recordSuccess: (userId, serverName) => { - failedAt.delete(cooldownKey(userId, serverName)); + recordSuccess: (identity) => { + failedAt.delete(identity); }, }; } -function resolveRecoveryTimeout(serverConfig: ParsedServerConfig): number { +interface RecoveryCandidate extends MCPServerCatalogRecoveryInput { + cooldownIdentity: string; + customUserVars?: Record; +} + +/** + * A cooldown must not outlive the configuration that failed: correcting a server's URL or + * transport has to be retryable at once, and the client refetches this catalog as soon as the + * server is updated. Keying by the publication generation — the same effective-config identity + * the tool caches fence on — means an edited server simply keys a new entry. + */ +function cooldownIdentity( + userId: string, + { serverName, serverConfig }: MCPServerCatalogRecoveryInput, +): string { + try { + return `${userId}:${serverName}:${getMCPAppToolsPublicationGeneration(serverConfig)}`; + } catch { + logger.debug( + `[MCP catalog recovery] ${serverName}: cooldown falls back to config-agnostic key`, + ); + return `${userId}:${serverName}`; + } +} + +/** Bounds one connection attempt: honours a shorter operator `initTimeout`, never the deadline. */ +function resolveAttemptTimeout(serverConfig: ParsedServerConfig): number { const { initTimeout } = serverConfig; if (typeof initTimeout === 'number') { - return Math.min(initTimeout, RECOVERY_TIMEOUT_MS); + return Math.min(initTimeout, RECOVERY_SERVER_DEADLINE_MS); } - return RECOVERY_TIMEOUT_MS; + return RECOVERY_SERVER_DEADLINE_MS; } /** @@ -126,21 +156,62 @@ function resolveRecoveryTimeout(serverConfig: ParsedServerConfig): number { * only repeats a known failure. */ function isRecoverable( - userId: string, - { serverName, serverConfig }: MCPServerCatalogRecoveryInput, + { serverName, serverConfig, cooldownIdentity: identity }: RecoveryCandidate, cooldown?: MCPCatalogRecoveryCooldown, ): boolean { if (serverConfig.inspectionFailed) { logger.debug(`[MCP catalog recovery] Skipping ${serverName}: awaiting config-tier retry`); return false; } - if (cooldown?.isCoolingDown(userId, serverName)) { + if (cooldown?.isCoolingDown(identity)) { logger.debug(`[MCP catalog recovery] Skipping ${serverName}: recent discovery failure`); return false; } return true; } +async function discoverCandidate( + user: IUser, + candidate: RecoveryCandidate, + budgetExpiresAt: number, + deps: MCPServerCatalogRecoveryDeps, +): Promise<[string, LCAvailableTools | null]> { + const { serverName, serverConfig, customUserVars, cooldownIdentity: identity } = candidate; + /** Never dialing a server is not evidence against it, so an exhausted budget records no + * cooldown; a later request reaches it once the servers ahead are cached or cooling down. */ + if (Date.now() + RECOVERY_SERVER_DEADLINE_MS > budgetExpiresAt) { + logger.debug(`[MCP catalog recovery] Skipping ${serverName}: request recovery budget spent`); + return [serverName, null]; + } + + try { + /** The deadline is what bounds this server's share of the request. `connectionTimeout` only + * bounds each attempt the factory makes inside it, and an attempt abandoned by the deadline + * still disposes its own connection when it eventually settles. */ + const result = await withTimeout( + deps.discoverServerTools({ + user, + serverName, + configServers: { [serverName]: serverConfig }, + customUserVars, + connectionTimeout: resolveAttemptTimeout(serverConfig), + }), + RECOVERY_SERVER_DEADLINE_MS, + `Discovery for ${serverName} exceeded ${RECOVERY_SERVER_DEADLINE_MS}ms`, + ); + if (result.tools == null) { + deps.recoveryCooldown?.recordFailure(identity); + return [serverName, null]; + } + deps.recoveryCooldown?.recordSuccess(identity); + return [serverName, deps.formatServerTools(serverName, result.tools)]; + } catch (error) { + deps.recoveryCooldown?.recordFailure(identity); + logger.error(`[MCP catalog recovery] Failed to discover tools for ${serverName}:`, error); + return [serverName, null]; + } +} + /** * Passively discovers cold MCP catalogs for one request. Results are intentionally not cached: * discovery connections do not own a publication generation, so publishing them could overwrite @@ -151,9 +222,9 @@ export async function recoverMCPServerCatalogs( deps: MCPServerCatalogRecoveryDeps, ): Promise> { const { user, servers } = params; - const recoverable = servers.filter((server) => - isRecoverable(user.id, server, deps.recoveryCooldown), - ); + const recoverable = servers + .map((server) => ({ ...server, cooldownIdentity: cooldownIdentity(user.id, server) })) + .filter((candidate) => isRecoverable(candidate, deps.recoveryCooldown)); if (recoverable.length === 0) { return new Map(); } @@ -165,48 +236,27 @@ export async function recoverMCPServerCatalogs( /** A server missing its user-provided credentials fails auth on connect (see issue #10969), * so discovering it would spend a doomed connection on every request. */ - const authorized: Array< - MCPServerCatalogRecoveryInput & { customUserVars?: Record } - > = []; - for (const server of recoverable) { - const customUserVars = getServerCustomUserVars(userMCPAuthMap, server.serverName); - const missingUserVars = getMissingCustomUserVars(server.serverConfig, customUserVars); + const authorized: RecoveryCandidate[] = []; + for (const candidate of recoverable) { + const customUserVars = getServerCustomUserVars(userMCPAuthMap, candidate.serverName); + const missingUserVars = getMissingCustomUserVars(candidate.serverConfig, customUserVars); if (missingUserVars.length > 0) { logger.debug( - `[MCP catalog recovery] Skipping ${server.serverName}: ${missingUserVars.length} user-provided variable(s) unset`, + `[MCP catalog recovery] Skipping ${candidate.serverName}: ${missingUserVars.length} user-provided variable(s) unset`, ); continue; } - authorized.push({ ...server, customUserVars }); + authorized.push({ ...candidate, customUserVars }); } if (authorized.length === 0) { return new Map(); } + const budgetExpiresAt = Date.now() + RECOVERY_REQUEST_BUDGET_MS; const recover = createConcurrencyLimiter(RECOVERY_CONCURRENCY); const results = await Promise.all( - authorized.map(({ serverName, serverConfig, customUserVars }) => - recover(async (): Promise<[string, LCAvailableTools | null]> => { - try { - const result = await deps.discoverServerTools({ - user, - serverName, - configServers: { [serverName]: serverConfig }, - customUserVars, - connectionTimeout: resolveRecoveryTimeout(serverConfig), - }); - if (result.tools == null) { - deps.recoveryCooldown?.recordFailure(user.id, serverName); - return [serverName, null]; - } - deps.recoveryCooldown?.recordSuccess(user.id, serverName); - return [serverName, deps.formatServerTools(serverName, result.tools)]; - } catch (error) { - deps.recoveryCooldown?.recordFailure(user.id, serverName); - logger.error(`[MCP catalog recovery] Failed to discover tools for ${serverName}:`, error); - return [serverName, null]; - } - }), + authorized.map((candidate) => + recover(() => discoverCandidate(user, candidate, budgetExpiresAt, deps)), ), );