mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-10-22 10:02:41 +00:00
🧺 fix: Funnel MCP Catalog Work Through a Shared Gate (#15429)
* 🧯 fix: Bound MCP Catalog Work Across Requests * 🛑 fix: Reject Saturated MCP Discovery * fix: fairly bound MCP catalog work * test: type catalog snapshot mock arguments * test: assert snapshot deadline forwarding * fix: harden MCP catalog scheduling
This commit is contained in:
parent
2e2c29e10d
commit
49e1e32f3a
10 changed files with 357 additions and 82 deletions
|
|
@ -13,6 +13,7 @@ const {
|
|||
createAuthIdentityContext,
|
||||
MCPConnection,
|
||||
MCPErrorCodes,
|
||||
MCPCatalogCapacityError,
|
||||
splitMCPToolKey,
|
||||
normalizeServerName,
|
||||
findShadowedServerNames,
|
||||
|
|
@ -201,21 +202,35 @@ const getMCPTools = async (req, res) => {
|
|||
user: req.user,
|
||||
tenantId: getTenantId(),
|
||||
});
|
||||
const { serverTools: serverToolsMap, serversWithoutTools } = await loadMCPServerCatalogs({
|
||||
user: req.user,
|
||||
servers: configuredServers.map((serverName) => ({
|
||||
serverName,
|
||||
serverConfig: mcpConfig[serverName],
|
||||
})),
|
||||
upstreamTokenProvider: createOpenIDSessionTokenProvider({
|
||||
req,
|
||||
res,
|
||||
const catalogAbortController = new AbortController();
|
||||
const abortCatalogLoad = () => {
|
||||
if (!res.writableEnded) {
|
||||
catalogAbortController.abort();
|
||||
}
|
||||
};
|
||||
res.once('close', abortCatalogLoad);
|
||||
let catalogResult;
|
||||
try {
|
||||
catalogResult = await loadMCPServerCatalogs({
|
||||
user: req.user,
|
||||
identityContext: oboIdentityContext,
|
||||
tokenPreference: 'access_token',
|
||||
}),
|
||||
oboIdentityContext,
|
||||
});
|
||||
servers: configuredServers.map((serverName) => ({
|
||||
serverName,
|
||||
serverConfig: mcpConfig[serverName],
|
||||
})),
|
||||
upstreamTokenProvider: createOpenIDSessionTokenProvider({
|
||||
req,
|
||||
res,
|
||||
user: req.user,
|
||||
identityContext: oboIdentityContext,
|
||||
tokenPreference: 'access_token',
|
||||
}),
|
||||
oboIdentityContext,
|
||||
signal: catalogAbortController.signal,
|
||||
});
|
||||
} finally {
|
||||
res.off('close', abortCatalogLoad);
|
||||
}
|
||||
const { serverTools: serverToolsMap, serversWithoutTools } = catalogResult;
|
||||
if (serversWithoutTools.length > 0) {
|
||||
logger.debug(
|
||||
`[getMCPTools] No tools (${serversWithoutTools.length}): ${serversWithoutTools.join(', ')}`,
|
||||
|
|
@ -286,7 +301,11 @@ const getMCPTools = async (req, res) => {
|
|||
res.status(200).json({ servers: mcpServers });
|
||||
} catch (error) {
|
||||
logger.error('[getMCPTools]', error);
|
||||
res.status(500).json({ message: error.message });
|
||||
if (res.destroyed || res.headersSent) {
|
||||
return;
|
||||
}
|
||||
const status = error instanceof MCPCatalogCapacityError ? 503 : 500;
|
||||
res.status(status).json({ message: error.message });
|
||||
}
|
||||
};
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -3324,6 +3324,7 @@ describe('MCP Routes', () => {
|
|||
servers: [{ serverName: 'user-server', serverConfig }],
|
||||
upstreamTokenProvider: expect.any(Function),
|
||||
oboIdentityContext: expect.any(Object),
|
||||
signal: expect.any(AbortSignal),
|
||||
});
|
||||
});
|
||||
|
||||
|
|
@ -3376,6 +3377,7 @@ describe('MCP Routes', () => {
|
|||
servers: [{ serverName: 'connected-server', serverConfig }],
|
||||
upstreamTokenProvider: expect.any(Function),
|
||||
oboIdentityContext: expect.any(Object),
|
||||
signal: expect.any(AbortSignal),
|
||||
});
|
||||
});
|
||||
|
||||
|
|
@ -3474,8 +3476,22 @@ describe('MCP Routes', () => {
|
|||
})),
|
||||
upstreamTokenProvider: expect.any(Function),
|
||||
oboIdentityContext: expect.any(Object),
|
||||
signal: expect.any(AbortSignal),
|
||||
});
|
||||
});
|
||||
|
||||
it('returns a retryable failure when process-wide catalog capacity is saturated', async () => {
|
||||
const { MCPCatalogCapacityError } = require('@librechat/api');
|
||||
mockResolveAllMcpConfigs.mockResolvedValueOnce({
|
||||
saturated: { type: 'sse', url: 'https://saturated.example.com/sse' },
|
||||
});
|
||||
mockLoadMCPServerCatalogs.mockRejectedValueOnce(new MCPCatalogCapacityError());
|
||||
|
||||
const response = await request(app).get('/api/mcp/tools');
|
||||
|
||||
expect(response.status).toBe(503);
|
||||
expect(response.body).toEqual({ message: 'MCP catalog capacity is temporarily exhausted' });
|
||||
});
|
||||
});
|
||||
|
||||
describe('GET /servers', () => {
|
||||
|
|
|
|||
|
|
@ -399,11 +399,12 @@ async function getAssistantToolDefinitions({ req, res, tools }) {
|
|||
getAllServerConfigs: (userId, configServers, role) =>
|
||||
registry.getAllServerConfigs(userId, configServers, role),
|
||||
getMCPServerTools,
|
||||
getServerToolFunctionsSnapshot: async (userId, serverName, serverConfig) =>
|
||||
getServerToolFunctionsSnapshot: async (userId, serverName, serverConfig, options) =>
|
||||
(await getMCPManager()?.getServerToolFunctionsSnapshot(
|
||||
userId,
|
||||
serverName,
|
||||
serverConfig,
|
||||
options,
|
||||
)) ?? {
|
||||
tools: null,
|
||||
},
|
||||
|
|
|
|||
|
|
@ -40,13 +40,20 @@ const MCP_REINITIALIZE_FAILURE_REASONS = {
|
|||
* @param {Array<{ serverName: string, serverConfig: object }>} params.servers
|
||||
* @param {import('@librechat/api').UpstreamTokenProvider} [params.upstreamTokenProvider] - Live upstream-token closure for OBO discovery, built at the request boundary so this layer never receives the raw Express request.
|
||||
* @param {import('@librechat/api').AuthIdentityContext} [params.oboIdentityContext] - Non-template-visible OBO identity context built from the real request user.
|
||||
* @param {AbortSignal} [params.signal] - Cancels queued and in-flight catalog reads when the request ends.
|
||||
*/
|
||||
async function loadMCPServerCatalogs({ user, servers, upstreamTokenProvider, oboIdentityContext }) {
|
||||
async function loadMCPServerCatalogs({
|
||||
user,
|
||||
servers,
|
||||
upstreamTokenProvider,
|
||||
oboIdentityContext,
|
||||
signal,
|
||||
}) {
|
||||
const flowManager = getFlowStateManager(getLogStores(CacheKeys.FLOWS));
|
||||
const tokenMethods = { findToken, updateToken, createToken, deleteTokens };
|
||||
const mcpManager = getMCPManager();
|
||||
return loadCatalogs(
|
||||
{ user, servers },
|
||||
{ user, servers, signal },
|
||||
{
|
||||
loadUserMCPAuthMap: (userId, serverNames) =>
|
||||
getUserMCPAuthMap({
|
||||
|
|
@ -67,8 +74,8 @@ async function loadMCPServerCatalogs({ user, servers, upstreamTokenProvider, obo
|
|||
}),
|
||||
formatServerTools: formatMCPServerTools,
|
||||
getCachedServerTools: getMCPServerTools,
|
||||
getServerToolFunctionsSnapshot: (userId, serverName, serverConfig) =>
|
||||
mcpManager.getServerToolFunctionsSnapshot(userId, serverName, serverConfig),
|
||||
getServerToolFunctionsSnapshot: (userId, serverName, serverConfig, options) =>
|
||||
mcpManager.getServerToolFunctionsSnapshot(userId, serverName, serverConfig, options),
|
||||
cacheServerTools: cacheMCPServerTools,
|
||||
},
|
||||
);
|
||||
|
|
|
|||
|
|
@ -85,7 +85,9 @@ describe('loadMCPServerCatalogs', () => {
|
|||
});
|
||||
deps.formatServerTools('config-only', []);
|
||||
await deps.getCachedServerTools(user.id, 'config-only', servers[0].serverConfig);
|
||||
await deps.getServerToolFunctionsSnapshot(user.id, 'config-only', servers[0].serverConfig);
|
||||
await deps.getServerToolFunctionsSnapshot(user.id, 'config-only', servers[0].serverConfig, {
|
||||
deadlineMs: 123,
|
||||
});
|
||||
await deps.cacheServerTools({ serverName: 'config-only' });
|
||||
return { serverTools: new Map([['config-only', {}]]), serversWithoutTools: [] };
|
||||
});
|
||||
|
|
@ -126,6 +128,7 @@ describe('loadMCPServerCatalogs', () => {
|
|||
user.id,
|
||||
'config-only',
|
||||
servers[0].serverConfig,
|
||||
{ deadlineMs: 123 },
|
||||
);
|
||||
expect(mockCacheMCPServerTools).toHaveBeenCalledWith({ serverName: 'config-only' });
|
||||
expect(result).toEqual({
|
||||
|
|
|
|||
|
|
@ -450,12 +450,18 @@ export class MCPManager extends UserConnectionManager {
|
|||
userId: string,
|
||||
serverName: string,
|
||||
serverConfig?: t.ParsedServerConfig,
|
||||
options?: { deadlineMs?: number; signal?: AbortSignal },
|
||||
): Promise<{
|
||||
tools: t.LCAvailableTools | null;
|
||||
publicationGeneration?: string;
|
||||
publicationRevision?: string;
|
||||
}> {
|
||||
try {
|
||||
const signal = createDeadlineAbortSignal(options?.deadlineMs, options?.signal);
|
||||
const readToolCatalog = (connection: MCPConnection) =>
|
||||
options == null
|
||||
? MCPServerInspector.getToolCatalog(serverName, connection)
|
||||
: MCPServerInspector.getToolCatalog(serverName, connection, options.deadlineMs, signal);
|
||||
const registry = MCPServersRegistry.getInstance();
|
||||
const effectiveConfig = serverConfig ?? (await registry.getServerConfig(serverName, userId));
|
||||
const useAppConnection =
|
||||
|
|
@ -466,7 +472,7 @@ export class MCPManager extends UserConnectionManager {
|
|||
? await this.appConnections?.get(serverName)
|
||||
: null;
|
||||
if (existingAppConnection != null) {
|
||||
return MCPServerInspector.getToolCatalog(serverName, existingAppConnection);
|
||||
return readToolCatalog(existingAppConnection);
|
||||
}
|
||||
|
||||
let awaitedRecovery: Promise<void> | undefined;
|
||||
|
|
@ -507,12 +513,12 @@ export class MCPManager extends UserConnectionManager {
|
|||
if (recovery && recovery !== awaitedRecovery) {
|
||||
awaitedRecovery = recovery;
|
||||
await this.releaseConnection(connection);
|
||||
await this.waitForConnectionRecovery(recovery);
|
||||
await this.waitForConnectionRecovery(recovery, signal);
|
||||
continue;
|
||||
}
|
||||
|
||||
try {
|
||||
const { tools } = await MCPServerInspector.getToolCatalog(serverName, connection);
|
||||
const { tools } = await readToolCatalog(connection);
|
||||
const generationAfterFetch = await getMCPToolsChangedGeneration({ userId, serverName });
|
||||
if (
|
||||
publicationGeneration != null &&
|
||||
|
|
|
|||
|
|
@ -3878,13 +3878,27 @@ describe('MCPManager', () => {
|
|||
)?.[1];
|
||||
const tools: t.MCPTool[] = [{ name: 'current', inputSchema: { type: 'object' } }];
|
||||
listener(tools);
|
||||
const snapshot = await manager.getServerToolFunctionsSnapshot(userId, serverName);
|
||||
const deadlineMs = Date.now() + 3000;
|
||||
const snapshot = await manager.getServerToolFunctionsSnapshot(
|
||||
userId,
|
||||
serverName,
|
||||
undefined,
|
||||
{
|
||||
deadlineMs,
|
||||
},
|
||||
);
|
||||
|
||||
expect(generationSpy).toHaveBeenCalledWith({ userId, serverName });
|
||||
expect(snapshot).toEqual({
|
||||
tools: expectedToolFunctions,
|
||||
publicationGeneration: 'generation-a',
|
||||
});
|
||||
expect(MCPServerInspector.getToolCatalog).toHaveBeenCalledWith(
|
||||
serverName,
|
||||
mockConnection,
|
||||
deadlineMs,
|
||||
expect.any(AbortSignal),
|
||||
);
|
||||
expect(notifySpy).toHaveBeenCalledWith({
|
||||
tools,
|
||||
userId,
|
||||
|
|
|
|||
|
|
@ -3,7 +3,11 @@ import { Constants } from 'librechat-data-provider';
|
|||
import { ErrorCode, McpError } from '@modelcontextprotocol/sdk/types.js';
|
||||
import type { IUser } from '@librechat/data-schemas';
|
||||
import type { LCAvailableTools, ParsedServerConfig, ToolDiscoveryOptions } from '../types';
|
||||
import { loadMCPServerCatalogs, recoverMCPServerCatalogs } from './recovery';
|
||||
import {
|
||||
MCPCatalogCapacityError,
|
||||
loadMCPServerCatalogs,
|
||||
recoverMCPServerCatalogs,
|
||||
} from './recovery';
|
||||
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: { info: jest.fn(), warn: jest.fn(), error: jest.fn(), debug: jest.fn() },
|
||||
|
|
@ -71,7 +75,7 @@ describe('recoverMCPServerCatalogs', () => {
|
|||
expect(result.size).toBe(2);
|
||||
});
|
||||
|
||||
it('limits passive discovery to three concurrent servers', async () => {
|
||||
it('limits passive discovery to three concurrent servers across simultaneous catalog loads', async () => {
|
||||
const servers = Array.from({ length: 7 }, (_, index) => ({
|
||||
serverName: `server-${index}`,
|
||||
serverConfig: serverConfig(`server-${index}`),
|
||||
|
|
@ -85,19 +89,39 @@ describe('recoverMCPServerCatalogs', () => {
|
|||
active -= 1;
|
||||
return { tools: [] };
|
||||
});
|
||||
const getServerToolFunctionsSnapshot = jest.fn(async () => {
|
||||
active += 1;
|
||||
maxActive = Math.max(maxActive, active);
|
||||
await new Promise((resolve) => setTimeout(resolve, 5));
|
||||
active -= 1;
|
||||
return { tools: null };
|
||||
});
|
||||
|
||||
const result = await recoverMCPServerCatalogs(
|
||||
{ user, servers },
|
||||
{
|
||||
loadUserMCPAuthMap: jest.fn().mockResolvedValue({}),
|
||||
discoverServerTools,
|
||||
formatServerTools: jest.fn().mockReturnValue({}),
|
||||
},
|
||||
const deps = {
|
||||
getCachedServerTools: jest.fn().mockResolvedValue(null),
|
||||
getServerToolFunctionsSnapshot,
|
||||
cacheServerTools: jest.fn(),
|
||||
loadUserMCPAuthMap: jest.fn().mockResolvedValue({}),
|
||||
discoverServerTools,
|
||||
formatServerTools: jest.fn().mockReturnValue({}),
|
||||
};
|
||||
const results = await Promise.allSettled(
|
||||
Array.from({ length: 4 }, () => loadMCPServerCatalogs({ user, servers }, deps)),
|
||||
);
|
||||
|
||||
expect(discoverServerTools).toHaveBeenCalledTimes(7);
|
||||
expect(getServerToolFunctionsSnapshot).toHaveBeenCalledTimes(21);
|
||||
expect(discoverServerTools).toHaveBeenCalledTimes(21);
|
||||
expect(maxActive).toBe(3);
|
||||
expect(result.size).toBe(7);
|
||||
expect(results).toHaveLength(4);
|
||||
expect(results.filter(({ status }) => status === 'fulfilled')).toHaveLength(3);
|
||||
expect(results.filter(({ status }) => status === 'rejected')).toEqual([
|
||||
{ status: 'rejected', reason: expect.any(MCPCatalogCapacityError) },
|
||||
]);
|
||||
expect(
|
||||
results.flatMap((result) =>
|
||||
result.status === 'fulfilled' ? [...result.value.serverTools.keys()] : [],
|
||||
),
|
||||
).toHaveLength(21);
|
||||
});
|
||||
|
||||
it('logs configuration-impossible discovery failures at debug, keeping error for the unexpected', async () => {
|
||||
|
|
@ -500,22 +524,31 @@ describe('recoverMCPServerCatalogs — bounded, skippable discovery', () => {
|
|||
it('bounds snapshot refreshes, which each issue a real tools/list', async () => {
|
||||
let active = 0;
|
||||
let maxActive = 0;
|
||||
const before = Date.now();
|
||||
const servers = Array.from({ length: 9 }, (_, index) => ({
|
||||
serverName: `server-${index}`,
|
||||
serverConfig: serverConfig(`server-${index}`),
|
||||
}));
|
||||
const getServerToolFunctionsSnapshot = jest.fn(
|
||||
async (
|
||||
_userId: string,
|
||||
_serverName: string,
|
||||
_serverConfig: ParsedServerConfig,
|
||||
_options?: { deadlineMs?: number; signal?: AbortSignal },
|
||||
) => {
|
||||
active += 1;
|
||||
maxActive = Math.max(maxActive, active);
|
||||
await new Promise((resolve) => setTimeout(resolve, 20));
|
||||
active -= 1;
|
||||
return { tools: null };
|
||||
},
|
||||
);
|
||||
|
||||
await loadMCPServerCatalogs(
|
||||
{ user, servers },
|
||||
{
|
||||
getCachedServerTools: jest.fn().mockResolvedValue(null),
|
||||
getServerToolFunctionsSnapshot: jest.fn(async () => {
|
||||
active += 1;
|
||||
maxActive = Math.max(maxActive, active);
|
||||
await new Promise((resolve) => setTimeout(resolve, 20));
|
||||
active -= 1;
|
||||
return { tools: null };
|
||||
}),
|
||||
getServerToolFunctionsSnapshot,
|
||||
cacheServerTools: jest.fn(),
|
||||
loadUserMCPAuthMap: jest.fn().mockResolvedValue({}),
|
||||
discoverServerTools: jest.fn().mockResolvedValue({ tools: null }),
|
||||
|
|
@ -524,5 +557,49 @@ describe('recoverMCPServerCatalogs — bounded, skippable discovery', () => {
|
|||
);
|
||||
|
||||
expect(maxActive).toBe(3);
|
||||
const options = getServerToolFunctionsSnapshot.mock.calls[0]?.[3];
|
||||
expect(options?.deadlineMs).toBeGreaterThanOrEqual(before + 30_000);
|
||||
expect(options?.deadlineMs).toBeLessThanOrEqual(Date.now() + 30_000);
|
||||
});
|
||||
|
||||
it('cancels queued catalog work when the originating request ends', async () => {
|
||||
const controller = new AbortController();
|
||||
const releases: Array<() => void> = [];
|
||||
const servers = Array.from({ length: 9 }, (_, index) => ({
|
||||
serverName: `server-${index}`,
|
||||
serverConfig: serverConfig(`server-${index}`),
|
||||
}));
|
||||
const getServerToolFunctionsSnapshot = jest.fn(
|
||||
() =>
|
||||
new Promise<{ tools: null }>((resolve) => {
|
||||
releases.push(() => resolve({ tools: null }));
|
||||
}),
|
||||
);
|
||||
|
||||
const loading = loadMCPServerCatalogs(
|
||||
{ user, servers, signal: controller.signal },
|
||||
{
|
||||
getCachedServerTools: jest.fn().mockResolvedValue(null),
|
||||
getServerToolFunctionsSnapshot,
|
||||
cacheServerTools: jest.fn(),
|
||||
loadUserMCPAuthMap: jest.fn().mockResolvedValue({}),
|
||||
discoverServerTools: jest.fn().mockResolvedValue({ tools: null }),
|
||||
formatServerTools: jest.fn().mockReturnValue({}),
|
||||
},
|
||||
);
|
||||
await new Promise((resolve) => setImmediate(resolve));
|
||||
expect(getServerToolFunctionsSnapshot).toHaveBeenCalledTimes(3);
|
||||
|
||||
controller.abort();
|
||||
releases.splice(0).forEach((release) => release());
|
||||
await loading;
|
||||
|
||||
expect(getServerToolFunctionsSnapshot).toHaveBeenCalledTimes(3);
|
||||
expect(getServerToolFunctionsSnapshot).toHaveBeenCalledWith(
|
||||
user.id,
|
||||
'server-0',
|
||||
servers[0].serverConfig,
|
||||
expect.objectContaining({ signal: controller.signal }),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -4,12 +4,14 @@ import type { Tool } from '@modelcontextprotocol/sdk/types.js';
|
|||
import type { IUser } from '@librechat/data-schemas';
|
||||
import type { LCAvailableTools, ParsedServerConfig, ToolDiscoveryOptions } from '../types';
|
||||
import { hasCustomUserVars, getMissingCustomUserVars } from '../utils';
|
||||
import { createConcurrencyLimiter } from '~/utils/promise';
|
||||
import { getServerCustomUserVars } from '../auth';
|
||||
import { mcpConfig } from '../mcpConfig';
|
||||
|
||||
/** Bounds outbound MCP fan-out from one catalog request: snapshot refreshes issue a real
|
||||
* `tools/list`, so they burst as readily as passive discovery does. */
|
||||
/** Bounds all outbound catalog reads across this runtime. */
|
||||
const CATALOG_FANOUT_CONCURRENCY = 3;
|
||||
const activeCatalogLanes = new Set<CatalogWorkLane>();
|
||||
const pendingCatalogLanes: CatalogWorkLane[] = [];
|
||||
let activeCatalogWork = 0;
|
||||
/**
|
||||
* Bounds one server's discovery end to end — connect, `tools/list` pagination, and the
|
||||
* unauthenticated fallback all draw down this single budget, so a slot is held for at most this
|
||||
|
|
@ -48,6 +50,7 @@ export interface MCPServerCatalogLoaderDeps extends MCPServerCatalogRecoveryDeps
|
|||
userId: string,
|
||||
serverName: string,
|
||||
serverConfig: ParsedServerConfig,
|
||||
options?: { deadlineMs?: number; signal?: AbortSignal },
|
||||
) => Promise<MCPServerCatalogSnapshot>;
|
||||
cacheServerTools: (params: {
|
||||
userId: string;
|
||||
|
|
@ -64,10 +67,34 @@ export interface MCPServerCatalogLoaderResult {
|
|||
serversWithoutTools: string[];
|
||||
}
|
||||
|
||||
interface MCPServerCatalogEntry extends MCPServerCatalogSnapshot {
|
||||
serverName: string;
|
||||
serverConfig: ParsedServerConfig;
|
||||
source: 'cache' | 'snapshot';
|
||||
}
|
||||
|
||||
interface RecoveryCandidate extends MCPServerCatalogRecoveryInput {
|
||||
customUserVars?: Record<string, string>;
|
||||
}
|
||||
|
||||
interface CatalogWorkLane {
|
||||
tasks: ReadonlyArray<() => Promise<void>>;
|
||||
resolve: (admitted: boolean) => void;
|
||||
nextIndex: number;
|
||||
inFlight: number;
|
||||
signal?: AbortSignal;
|
||||
onAbort?: () => void;
|
||||
}
|
||||
|
||||
export class MCPCatalogCapacityError extends Error {
|
||||
readonly code = 'MCP_CATALOG_CAPACITY';
|
||||
|
||||
constructor() {
|
||||
super('MCP catalog capacity is temporarily exhausted');
|
||||
this.name = 'MCPCatalogCapacityError';
|
||||
}
|
||||
}
|
||||
|
||||
/** Bounds one server's discovery, honouring a shorter operator `initTimeout`. */
|
||||
function resolveBudget(serverConfig: ParsedServerConfig): number {
|
||||
const { initTimeout } = serverConfig;
|
||||
|
|
@ -81,6 +108,7 @@ async function discoverCandidate(
|
|||
user: IUser,
|
||||
{ serverName, serverConfig, customUserVars }: RecoveryCandidate,
|
||||
deps: MCPServerCatalogRecoveryDeps,
|
||||
signal?: AbortSignal,
|
||||
): Promise<[string, LCAvailableTools | null]> {
|
||||
try {
|
||||
const result = await deps.discoverServerTools({
|
||||
|
|
@ -89,6 +117,7 @@ async function discoverCandidate(
|
|||
configServers: { [serverName]: serverConfig },
|
||||
customUserVars,
|
||||
deadlineMs: Date.now() + resolveBudget(serverConfig),
|
||||
signal,
|
||||
});
|
||||
return [
|
||||
serverName,
|
||||
|
|
@ -110,6 +139,82 @@ async function discoverCandidate(
|
|||
}
|
||||
}
|
||||
|
||||
function finishCatalogLane(lane: CatalogWorkLane): void {
|
||||
if (lane.nextIndex < lane.tasks.length || lane.inFlight > 0) {
|
||||
return;
|
||||
}
|
||||
activeCatalogLanes.delete(lane);
|
||||
if (lane.signal != null && lane.onAbort != null) {
|
||||
lane.signal.removeEventListener('abort', lane.onAbort);
|
||||
}
|
||||
lane.resolve(true);
|
||||
}
|
||||
|
||||
/** Shares the fixed process budget round-robin across at most three request lanes. */
|
||||
function scheduleCatalogWork(): void {
|
||||
while (activeCatalogWork < CATALOG_FANOUT_CONCURRENCY && pendingCatalogLanes.length > 0) {
|
||||
const lane = pendingCatalogLanes.shift();
|
||||
if (lane == null) {
|
||||
return;
|
||||
}
|
||||
const task = lane.tasks[lane.nextIndex];
|
||||
lane.nextIndex += 1;
|
||||
lane.inFlight += 1;
|
||||
activeCatalogWork += 1;
|
||||
if (lane.nextIndex < lane.tasks.length) {
|
||||
pendingCatalogLanes.push(lane);
|
||||
}
|
||||
void task()
|
||||
.catch((error) => {
|
||||
logger.error('[MCP catalog] Scheduled outbound work failed:', error);
|
||||
})
|
||||
.finally(() => {
|
||||
lane.inFlight -= 1;
|
||||
activeCatalogWork -= 1;
|
||||
finishCatalogLane(lane);
|
||||
scheduleCatalogWork();
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
function runCatalogWork(
|
||||
tasks: ReadonlyArray<() => Promise<void>>,
|
||||
signal?: AbortSignal,
|
||||
): Promise<boolean> {
|
||||
if (tasks.length === 0) {
|
||||
return Promise.resolve(true);
|
||||
}
|
||||
if (signal?.aborted) {
|
||||
return Promise.resolve(true);
|
||||
}
|
||||
if (activeCatalogLanes.size >= CATALOG_FANOUT_CONCURRENCY) {
|
||||
logger.debug('[MCP catalog] Skipping request: outbound capacity reached');
|
||||
return Promise.resolve(false);
|
||||
}
|
||||
return new Promise((resolve) => {
|
||||
const lane: CatalogWorkLane = {
|
||||
tasks,
|
||||
resolve,
|
||||
nextIndex: 0,
|
||||
inFlight: 0,
|
||||
signal,
|
||||
};
|
||||
lane.onAbort = () => {
|
||||
lane.nextIndex = lane.tasks.length;
|
||||
for (let index = pendingCatalogLanes.length - 1; index >= 0; index -= 1) {
|
||||
if (pendingCatalogLanes[index] === lane) {
|
||||
pendingCatalogLanes.splice(index, 1);
|
||||
}
|
||||
}
|
||||
finishCatalogLane(lane);
|
||||
};
|
||||
signal?.addEventListener('abort', lane.onAbort, { once: true });
|
||||
activeCatalogLanes.add(lane);
|
||||
pendingCatalogLanes.push(lane);
|
||||
scheduleCatalogWork();
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Passively discovers cold MCP catalogs for one request.
|
||||
*
|
||||
|
|
@ -120,10 +225,14 @@ async function discoverCandidate(
|
|||
* proves pointless, and bounds everything else by a per-server budget.
|
||||
*/
|
||||
export async function recoverMCPServerCatalogs(
|
||||
params: { user: IUser; servers: readonly MCPServerCatalogRecoveryInput[] },
|
||||
params: {
|
||||
user: IUser;
|
||||
servers: readonly MCPServerCatalogRecoveryInput[];
|
||||
signal?: AbortSignal;
|
||||
},
|
||||
deps: MCPServerCatalogRecoveryDeps,
|
||||
): Promise<Map<string, LCAvailableTools>> {
|
||||
const { user, servers } = params;
|
||||
const { user, servers, signal } = params;
|
||||
/** Only the config tier retries a failed stub on its own clock. A `yaml`- or `user`-sourced
|
||||
* stub has no such timer, so skipping it unconditionally would hide the server for good —
|
||||
* exactly the state this recovery exists to escape. */
|
||||
|
|
@ -168,21 +277,35 @@ export async function recoverMCPServerCatalogs(
|
|||
return new Map();
|
||||
}
|
||||
|
||||
const recover = createConcurrencyLimiter(CATALOG_FANOUT_CONCURRENCY);
|
||||
const results = await Promise.all(
|
||||
authorized.map((candidate) => recover(() => discoverCandidate(user, candidate, deps))),
|
||||
const results: Array<[string, LCAvailableTools | null] | undefined> = [];
|
||||
const admitted = await runCatalogWork(
|
||||
authorized.map((candidate, index) => async () => {
|
||||
results[index] = await discoverCandidate(user, candidate, deps, signal);
|
||||
}),
|
||||
signal,
|
||||
);
|
||||
if (!admitted) {
|
||||
throw new MCPCatalogCapacityError();
|
||||
}
|
||||
|
||||
return new Map(results.filter((entry): entry is [string, LCAvailableTools] => entry[1] != null));
|
||||
return new Map(
|
||||
results.filter(
|
||||
(entry): entry is [string, LCAvailableTools] => entry != null && entry[1] != null,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
/** Loads cached, connected, then passive MCP catalogs for a marketplace-style list request. */
|
||||
export async function loadMCPServerCatalogs(
|
||||
params: { user: IUser; servers: readonly MCPServerCatalogRecoveryInput[] },
|
||||
params: {
|
||||
user: IUser;
|
||||
servers: readonly MCPServerCatalogRecoveryInput[];
|
||||
signal?: AbortSignal;
|
||||
},
|
||||
deps: MCPServerCatalogLoaderDeps,
|
||||
): Promise<MCPServerCatalogLoaderResult> {
|
||||
const { user, servers } = params;
|
||||
const cached = await Promise.all(
|
||||
const { user, servers, signal } = params;
|
||||
const cached: MCPServerCatalogEntry[] = await Promise.all(
|
||||
servers.map(async ({ serverName, serverConfig }) => {
|
||||
try {
|
||||
const tools = await deps.getCachedServerTools(user.id, serverName, serverConfig);
|
||||
|
|
@ -194,32 +317,36 @@ export async function loadMCPServerCatalogs(
|
|||
}),
|
||||
);
|
||||
|
||||
/** A snapshot is not a local read — both connection paths issue a fresh `tools/list` — so a
|
||||
* cache reset across many servers would otherwise burst unbounded outbound requests. */
|
||||
const refresh = createConcurrencyLimiter(CATALOG_FANOUT_CONCURRENCY);
|
||||
const snapshots = await Promise.all(
|
||||
cached.map((entry) => {
|
||||
if (entry.tools != null) {
|
||||
return entry;
|
||||
}
|
||||
return refresh(async () => {
|
||||
try {
|
||||
const snapshot = await deps.getServerToolFunctionsSnapshot(
|
||||
user.id,
|
||||
entry.serverName,
|
||||
entry.serverConfig,
|
||||
);
|
||||
return { ...entry, ...snapshot, source: 'snapshot' as const };
|
||||
} catch (error) {
|
||||
logger.error(
|
||||
`[MCP catalog loader] Failed to read connected tools for ${entry.serverName}:`,
|
||||
error,
|
||||
);
|
||||
return { ...entry, tools: null, source: 'snapshot' as const };
|
||||
}
|
||||
});
|
||||
}),
|
||||
const snapshots = [...cached];
|
||||
const snapshotsAdmitted = await runCatalogWork(
|
||||
cached.flatMap((entry, index) =>
|
||||
entry.tools != null
|
||||
? []
|
||||
: [
|
||||
async () => {
|
||||
try {
|
||||
const snapshot = await deps.getServerToolFunctionsSnapshot(
|
||||
user.id,
|
||||
entry.serverName,
|
||||
entry.serverConfig,
|
||||
{ deadlineMs: Date.now() + mcpConfig.TOOLS_LIST_TIMEOUT_MS, signal },
|
||||
);
|
||||
snapshots[index] = { ...entry, ...snapshot, source: 'snapshot' as const };
|
||||
} catch (error) {
|
||||
logger.error(
|
||||
`[MCP catalog loader] Failed to read connected tools for ${entry.serverName}:`,
|
||||
error,
|
||||
);
|
||||
snapshots[index] = { ...entry, tools: null, source: 'snapshot' as const };
|
||||
}
|
||||
},
|
||||
],
|
||||
),
|
||||
signal,
|
||||
);
|
||||
if (!snapshotsAdmitted) {
|
||||
throw new MCPCatalogCapacityError();
|
||||
}
|
||||
|
||||
const coldServers = snapshots
|
||||
.filter(({ tools }) => tools == null)
|
||||
|
|
@ -227,8 +354,11 @@ export async function loadMCPServerCatalogs(
|
|||
let recovered = new Map<string, LCAvailableTools>();
|
||||
if (coldServers.length > 0) {
|
||||
try {
|
||||
recovered = await recoverMCPServerCatalogs({ user, servers: coldServers }, deps);
|
||||
recovered = await recoverMCPServerCatalogs({ user, servers: coldServers, signal }, deps);
|
||||
} catch (error) {
|
||||
if (error instanceof MCPCatalogCapacityError) {
|
||||
throw error;
|
||||
}
|
||||
logger.error('[MCP catalog loader] Failed to recover cold server catalogs:', error);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -177,8 +177,10 @@ export class MCPServerInspector {
|
|||
public static async getToolCatalog(
|
||||
serverName: string,
|
||||
connection: MCPConnection,
|
||||
deadlineMs?: number,
|
||||
signal?: AbortSignal,
|
||||
): Promise<{ tools: t.LCAvailableTools; publicationRevision?: string }> {
|
||||
const snapshot = await connection.fetchOrderedToolsSnapshot();
|
||||
const snapshot = await connection.fetchOrderedToolsSnapshot(deadlineMs, signal);
|
||||
if (!snapshot.complete) {
|
||||
throw new Error(`Incomplete tools/list snapshot for MCP server ${serverName}`);
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue