mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-09-07 15:09:41 +00:00
⏱️ fix: Bound MCP Tool Discovery With a Caller Deadline (#15346)
* fix: bound MCP tool discovery end to end with a caller deadline connectionTimeout bounds a single connect() only. Discovery then spends it again on the unauthenticated fallback and hands tools/list its own 30s budget, so a caller working to a deadline had no way to cap the whole operation. Thread an optional absolute deadline through discovery into both connect() and the tools/list walk, and dispose a timed-out authenticated connection before the fallback opens its own socket — withTimeout does not cancel the connect it abandoned, so the two were briefly concurrent. Passive catalog recovery now sets one 3s per-server budget instead of a per-attempt timeout it could spend several times over. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Xr1Dabvdn1mzzyYgpJgU5B * fix: close deadline gaps in the app path, refresh wait, and disposal race Codex review of the deadline threading found three real gaps: - discoverServerTools returns through an app-connection fast path that never reached the factory, so its tools/list kept the 30s default. - fetchOrderedToolsSnapshot checked the deadline before awaiting a refresh but the refresh runs on the connection's own budget, so an in-flight one could still hold a budgeted caller for that budget. Stop waiting on it rather than adopting it. - connectClient never rechecked isDisposed after awaiting constructTransport, so a connect abandoned by its caller could assign a transport and connect after dispose() had already found nothing to close. Disposing before the fallback widened that window, so bound the attempt to its own disposal. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Xr1Dabvdn1mzzyYgpJgU5B * fix: close the transport itself when disposal beats connect Rewriting the disposal-race test against a real SDK server and transport exposed a defect the mocked version could not see: client.close() only closes a transport the client has adopted, and it has not adopted one when disposal lands before client.connect(). The abandoned attempt's session therefore stayed open and the reference was merely dropped. Close the transport directly before closing the client. The test now asserts the server observed the close, which fails against the previous fix.
This commit is contained in:
Danny Avila
GitHub
parent
70f735336d
commit
41c92dfd06
9 changed files with 408 additions and 62 deletions
|
|
@ -73,6 +73,7 @@ export class MCPConnectionFactory {
|
|||
protected oauthEnd?: () => Promise<void>;
|
||||
protected returnOnOAuth?: boolean;
|
||||
protected readonly connectionTimeout?: number;
|
||||
protected readonly deadlineMs?: number;
|
||||
protected readonly oboTokenResolver?: OboTokenResolver;
|
||||
protected readonly oboTrustChecker?: OboTrustChecker;
|
||||
protected readonly upstreamTokenProvider?: UpstreamTokenProvider;
|
||||
|
|
@ -219,7 +220,7 @@ export class MCPConnectionFactory {
|
|||
connection.once('oauthRequired', oauthHandler);
|
||||
|
||||
try {
|
||||
const connectTimeout = this.connectionTimeout ?? this.serverConfig.initTimeout ?? 30000;
|
||||
const connectTimeout = this.resolveConnectTimeout(30000);
|
||||
await withTimeout(
|
||||
connection.connect(),
|
||||
connectTimeout,
|
||||
|
|
@ -227,7 +228,7 @@ export class MCPConnectionFactory {
|
|||
);
|
||||
|
||||
if (await connection.isConnected()) {
|
||||
const snapshot = await connection.fetchOrderedToolsSnapshot();
|
||||
const snapshot = await connection.fetchOrderedToolsSnapshot(this.deadlineMs);
|
||||
connection.removeListener('oauthRequired', oauthHandler);
|
||||
return {
|
||||
tools: snapshot.complete ? snapshot.tools : null,
|
||||
|
|
@ -242,24 +243,29 @@ export class MCPConnectionFactory {
|
|||
`${this.logPrefix} [Discovery] Connection failed, attempting unauthenticated tool listing`,
|
||||
);
|
||||
}
|
||||
|
||||
/** The authenticated attempt is done with, but `withTimeout` does not cancel the `connect()`
|
||||
* it gave up on — that socket stays open. Dispose before the fallback opens a second one so
|
||||
* a single discovery never holds two concurrent connects to the same server. */
|
||||
connection.removeListener('oauthRequired', oauthHandler);
|
||||
await this.disposeQuietly(connection);
|
||||
connection = null;
|
||||
oauthHandler = null;
|
||||
}
|
||||
|
||||
if (this.isPastDeadline()) {
|
||||
logger.debug(
|
||||
`${this.logPrefix} [Discovery] Budget exhausted; skipping unauthenticated tool listing`,
|
||||
);
|
||||
return { tools: null, connection: null, oauthRequired, oauthUrl };
|
||||
}
|
||||
|
||||
try {
|
||||
const tools = await this.attemptUnauthenticatedToolListing();
|
||||
if (connection && oauthHandler) {
|
||||
connection.removeListener('oauthRequired', oauthHandler);
|
||||
}
|
||||
if (tools && tools.length > 0) {
|
||||
logger.info(
|
||||
`${this.logPrefix} [Discovery] Successfully discovered ${tools.length} tools without auth`,
|
||||
);
|
||||
if (connection) {
|
||||
try {
|
||||
await connection.dispose();
|
||||
} catch {
|
||||
// Ignore cleanup errors
|
||||
}
|
||||
}
|
||||
return { tools, connection: null, oauthRequired, oauthUrl };
|
||||
}
|
||||
MCPConnection.decrementCycleCount(this.serverName);
|
||||
|
|
@ -268,21 +274,30 @@ export class MCPConnectionFactory {
|
|||
logger.debug(`${this.logPrefix} [Discovery] Unauthenticated tool listing failed`);
|
||||
}
|
||||
|
||||
if (connection && oauthHandler) {
|
||||
connection.removeListener('oauthRequired', oauthHandler);
|
||||
}
|
||||
|
||||
if (connection) {
|
||||
try {
|
||||
await connection.dispose();
|
||||
} catch {
|
||||
// Ignore cleanup errors
|
||||
}
|
||||
}
|
||||
|
||||
return { tools: null, connection: null, oauthRequired, oauthUrl };
|
||||
}
|
||||
|
||||
/** Clamps a single `connect()` to whatever remains of the caller's overall discovery budget. */
|
||||
private resolveConnectTimeout(fallback: number): number {
|
||||
const configured = this.connectionTimeout ?? this.serverConfig.initTimeout ?? fallback;
|
||||
if (this.deadlineMs == null) {
|
||||
return configured;
|
||||
}
|
||||
return Math.max(1, Math.min(configured, this.deadlineMs - Date.now()));
|
||||
}
|
||||
|
||||
private isPastDeadline(): boolean {
|
||||
return this.deadlineMs != null && Date.now() >= this.deadlineMs;
|
||||
}
|
||||
|
||||
private async disposeQuietly(connection: MCPConnection): Promise<void> {
|
||||
try {
|
||||
await connection.dispose();
|
||||
} catch {
|
||||
// Ignore cleanup errors
|
||||
}
|
||||
}
|
||||
|
||||
protected async attemptUnauthenticatedToolListing(): Promise<Tool[] | null> {
|
||||
const unauthConnection = new MCPConnection({
|
||||
serverName: this.serverName,
|
||||
|
|
@ -305,23 +320,19 @@ export class MCPConnectionFactory {
|
|||
});
|
||||
|
||||
try {
|
||||
const connectTimeout = this.connectionTimeout ?? this.serverConfig.initTimeout ?? 15000;
|
||||
const connectTimeout = this.resolveConnectTimeout(15000);
|
||||
await withTimeout(unauthConnection.connect(), connectTimeout, `Unauth connection timeout`);
|
||||
|
||||
if (await unauthConnection.isConnected()) {
|
||||
const snapshot = await unauthConnection.fetchOrderedToolsSnapshot();
|
||||
await unauthConnection.dispose();
|
||||
const snapshot = await unauthConnection.fetchOrderedToolsSnapshot(this.deadlineMs);
|
||||
await this.disposeQuietly(unauthConnection);
|
||||
return snapshot.complete ? snapshot.tools : null;
|
||||
}
|
||||
} catch {
|
||||
logger.debug(`${this.logPrefix} [Discovery] Unauthenticated connection attempt failed`);
|
||||
}
|
||||
|
||||
try {
|
||||
await unauthConnection.dispose();
|
||||
} catch {
|
||||
// Ignore cleanup errors
|
||||
}
|
||||
await this.disposeQuietly(unauthConnection);
|
||||
|
||||
return null;
|
||||
}
|
||||
|
|
@ -345,6 +356,7 @@ export class MCPConnectionFactory {
|
|||
this.allowedAddresses = basic.allowedAddresses;
|
||||
this.ephemeralConnection = basic.ephemeralConnection === true;
|
||||
this.connectionTimeout = options?.connectionTimeout;
|
||||
this.deadlineMs = options?.deadlineMs;
|
||||
this.tenantContext = tenantStorage?.getStore?.();
|
||||
this.tenantId = this.tenantContext?.tenantId ?? getTenantId();
|
||||
this.logPrefix = options?.user ? `[MCP][User: ${options.user.id}]` : '[MCP]';
|
||||
|
|
|
|||
|
|
@ -280,7 +280,7 @@ export class MCPManager extends UserConnectionManager {
|
|||
? await this.appConnections?.get(serverName)
|
||||
: null;
|
||||
if (existingAppConnection && (await existingAppConnection.isConnected())) {
|
||||
const snapshot = await existingAppConnection.fetchOrderedToolsSnapshot();
|
||||
const snapshot = await existingAppConnection.fetchOrderedToolsSnapshot(args.deadlineMs);
|
||||
return {
|
||||
tools: snapshot.complete ? snapshot.tools : null,
|
||||
oauthRequired: false,
|
||||
|
|
@ -350,6 +350,7 @@ export class MCPManager extends UserConnectionManager {
|
|||
requestBody: args.requestBody,
|
||||
graphTokenResolver: args.graphTokenResolver,
|
||||
connectionTimeout: args.connectionTimeout,
|
||||
deadlineMs: args.deadlineMs,
|
||||
});
|
||||
return finalizeDiscoveryResult(result);
|
||||
}
|
||||
|
|
@ -370,6 +371,7 @@ export class MCPManager extends UserConnectionManager {
|
|||
requestBody: args.requestBody,
|
||||
graphTokenResolver: args.graphTokenResolver,
|
||||
connectionTimeout: args.connectionTimeout,
|
||||
deadlineMs: args.deadlineMs,
|
||||
oboTokenResolver: args.oboTokenResolver,
|
||||
oboTrustChecker: args.oboTrustChecker,
|
||||
upstreamTokenProvider: args.upstreamTokenProvider,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,88 @@
|
|||
/**
|
||||
* Real-SDK coverage for disposal landing mid-connect.
|
||||
*
|
||||
* A discovery caller that gives up on a slow `connect()` disposes immediately, so it can dispose
|
||||
* while `constructTransport()` is still pending. `dispose()` then finds no transport to close, and
|
||||
* without a post-await check the abandoned attempt goes on to connect, leaving a live session
|
||||
* nobody owns. Only transport construction is delayed here; the client, transport, and server are
|
||||
* real SDK objects, so the assertions are about a genuinely open or closed session rather than
|
||||
* about which mock was called.
|
||||
*/
|
||||
import { Server } from '@modelcontextprotocol/sdk/server/index.js';
|
||||
import { InMemoryTransport } from '@modelcontextprotocol/sdk/inMemory.js';
|
||||
import { ListToolsRequestSchema } from '@modelcontextprotocol/sdk/types.js';
|
||||
import { MCPConnection } from '~/mcp/connection';
|
||||
|
||||
jest.setTimeout(10_000);
|
||||
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: {
|
||||
info: jest.fn(),
|
||||
warn: jest.fn(),
|
||||
error: jest.fn(),
|
||||
debug: jest.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
jest.mock('~/auth', () => ({
|
||||
createSSRFSafeUndiciConnect: jest.fn(() => undefined),
|
||||
isOAuthUrlAllowed: jest.fn(() => false),
|
||||
isSSRFTarget: jest.fn(() => false),
|
||||
resolveHostnameSSRF: jest.fn(async () => false),
|
||||
}));
|
||||
|
||||
describe('MCPConnection disposal during connect', () => {
|
||||
let server: Server | undefined;
|
||||
|
||||
afterEach(async () => {
|
||||
await server?.close().catch(() => undefined);
|
||||
server = undefined;
|
||||
});
|
||||
|
||||
it('leaves no live session when disposal lands while the transport is being constructed', async () => {
|
||||
server = new Server(
|
||||
{ name: 'dispose-race-server', version: '1.0.0' },
|
||||
{ capabilities: { tools: {} } },
|
||||
);
|
||||
server.setRequestHandler(ListToolsRequestSchema, async () => ({ tools: [] }));
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await server.connect(serverTransport);
|
||||
|
||||
const connection = new MCPConnection({
|
||||
serverName: 'dispose-race',
|
||||
serverConfig: { type: 'streamable-http', url: 'http://localhost/mcp' },
|
||||
useSSRFProtection: false,
|
||||
});
|
||||
|
||||
let serverSawClose = false;
|
||||
serverTransport.onclose = () => {
|
||||
serverSawClose = true;
|
||||
};
|
||||
|
||||
/** The only stub: hold construction open so disposal can land inside this window. */
|
||||
let releaseTransport: (() => void) | undefined;
|
||||
const transportPending = new Promise<void>((resolve) => {
|
||||
releaseTransport = resolve;
|
||||
});
|
||||
jest
|
||||
.spyOn(
|
||||
connection as unknown as { constructTransport: () => Promise<unknown> },
|
||||
'constructTransport',
|
||||
)
|
||||
.mockImplementation(async () => {
|
||||
await transportPending;
|
||||
return clientTransport;
|
||||
});
|
||||
|
||||
const connecting = connection.connectClient();
|
||||
await connection.dispose();
|
||||
releaseTransport?.();
|
||||
await connecting;
|
||||
|
||||
/** A real `tools/list` is the honest probe: it succeeds over any session left open. */
|
||||
await expect(connection.client.listTools()).rejects.toThrow();
|
||||
expect(await connection.isConnected()).toBe(false);
|
||||
expect(serverSawClose).toBe(true);
|
||||
});
|
||||
});
|
||||
|
|
@ -4190,6 +4190,100 @@ describe('MCPConnectionFactory', () => {
|
|||
expect(result.tools).toBeNull();
|
||||
expect(mockLogger.debug).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
describe('deadline bounding', () => {
|
||||
/** Builds distinct per-construction connection mocks so the authenticated attempt and the
|
||||
* unauthenticated fallback can be told apart, recording lifecycle order across both. */
|
||||
const trackConnections = (events: string[]) => {
|
||||
let index = 0;
|
||||
mockMCPConnection.mockImplementation(() => {
|
||||
const label = index++ === 0 ? 'auth' : 'unauth';
|
||||
const instance = {
|
||||
connect: jest.fn(async () => {
|
||||
events.push(`${label}:connect`);
|
||||
throw new Error('Connection failed');
|
||||
}),
|
||||
isConnected: jest.fn().mockResolvedValue(false),
|
||||
setOAuthTokens: jest.fn(),
|
||||
on: jest.fn(),
|
||||
once: jest.fn(),
|
||||
off: jest.fn(),
|
||||
removeListener: jest.fn(),
|
||||
emit: jest.fn(),
|
||||
dispose: jest.fn(async () => {
|
||||
events.push(`${label}:dispose`);
|
||||
}),
|
||||
} as unknown as jest.Mocked<MCPConnection>;
|
||||
return instance;
|
||||
});
|
||||
};
|
||||
|
||||
it('disposes the authenticated attempt before the fallback opens a second connection', async () => {
|
||||
const events: string[] = [];
|
||||
trackConnections(events);
|
||||
|
||||
const result = await MCPConnectionFactory.discoverTools({
|
||||
serverName: 'test-server',
|
||||
serverConfig: mockServerConfig,
|
||||
});
|
||||
|
||||
expect(result.tools).toBeNull();
|
||||
expect(events).toEqual([
|
||||
'auth:connect',
|
||||
'auth:dispose',
|
||||
'unauth:connect',
|
||||
'unauth:dispose',
|
||||
]);
|
||||
});
|
||||
|
||||
it('skips the unauthenticated fallback once the deadline has passed', async () => {
|
||||
const events: string[] = [];
|
||||
trackConnections(events);
|
||||
|
||||
const result = await MCPConnectionFactory.discoverTools(
|
||||
{ serverName: 'test-server', serverConfig: mockServerConfig },
|
||||
{ deadlineMs: Date.now() - 1 },
|
||||
);
|
||||
|
||||
expect(result.tools).toBeNull();
|
||||
expect(events).toEqual(['auth:connect', 'auth:dispose']);
|
||||
expect(mockMCPConnection).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('forwards the deadline to the tools/list snapshot', async () => {
|
||||
const deadlineMs = Date.now() + 5000;
|
||||
mockConnectionInstance.connect.mockResolvedValue(undefined);
|
||||
mockConnectionInstance.isConnected.mockResolvedValue(true);
|
||||
mockConnectionInstance.fetchOrderedToolsSnapshot = jest
|
||||
.fn()
|
||||
.mockResolvedValue({ tools: [], complete: true });
|
||||
|
||||
await MCPConnectionFactory.discoverTools(
|
||||
{ serverName: 'test-server', serverConfig: mockServerConfig },
|
||||
{ deadlineMs },
|
||||
);
|
||||
|
||||
expect(mockConnectionInstance.fetchOrderedToolsSnapshot).toHaveBeenCalledWith(deadlineMs);
|
||||
});
|
||||
|
||||
it('clamps a long initTimeout to the remaining budget instead of waiting it out', async () => {
|
||||
mockProcessMCPEnv.mockReturnValue({
|
||||
...mockServerConfig,
|
||||
initTimeout: 30000,
|
||||
} as t.MCPOptions);
|
||||
mockConnectionInstance.connect.mockImplementation(() => new Promise(() => {}));
|
||||
mockConnectionInstance.isConnected.mockResolvedValue(false);
|
||||
|
||||
const start = Date.now();
|
||||
const result = await MCPConnectionFactory.discoverTools(
|
||||
{ serverName: 'test-server', serverConfig: mockServerConfig },
|
||||
{ deadlineMs: Date.now() + 50 },
|
||||
);
|
||||
|
||||
expect(result.tools).toBeNull();
|
||||
expect(Date.now() - start).toBeLessThan(2000);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('proactive OAuth flow', () => {
|
||||
|
|
|
|||
|
|
@ -263,6 +263,84 @@ describe('MCPConnection.fetchTools pagination', () => {
|
|||
expect(JSON.stringify(mockLogger.error.mock.calls)).not.toContain('Request timed out');
|
||||
});
|
||||
|
||||
it('caps the request timeout by a caller deadline shorter than the global budget', async () => {
|
||||
mcpConfig.TOOLS_LIST_TIMEOUT_MS = 30000;
|
||||
const listTools = jest.fn().mockResolvedValue({ tools: [makeTool('a')] });
|
||||
const conn = createConnectionWithListTools(listTools);
|
||||
|
||||
const snapshot = await conn.fetchToolsSnapshot(Date.now() + 40);
|
||||
|
||||
expect(snapshot.complete).toBe(true);
|
||||
const options = listTools.mock.calls[0][1]!;
|
||||
expect(options.timeout).toBeGreaterThan(0);
|
||||
expect(options.timeout).toBeLessThanOrEqual(40);
|
||||
});
|
||||
|
||||
it('keeps the global budget when the caller deadline is further out', async () => {
|
||||
mcpConfig.TOOLS_LIST_TIMEOUT_MS = 25;
|
||||
const listTools = jest.fn().mockResolvedValue({ tools: [makeTool('a')] });
|
||||
const conn = createConnectionWithListTools(listTools);
|
||||
|
||||
await conn.fetchToolsSnapshot(Date.now() + 10_000);
|
||||
|
||||
const options = listTools.mock.calls[0][1]!;
|
||||
expect(options.timeout).toBeLessThanOrEqual(25);
|
||||
});
|
||||
|
||||
it('stops paginating and reports incomplete when the caller deadline expires mid-walk', async () => {
|
||||
mcpConfig.TOOLS_LIST_TIMEOUT_MS = 30000;
|
||||
const listTools = jest.fn(async () => ({ tools: [makeTool('a')], nextCursor: 'c1' }));
|
||||
const conn = createConnectionWithListTools(listTools);
|
||||
const deadline = Date.now() + 20;
|
||||
const dateNow = jest
|
||||
.spyOn(Date, 'now')
|
||||
.mockReturnValueOnce(deadline - 20)
|
||||
.mockReturnValueOnce(deadline - 20)
|
||||
.mockReturnValue(deadline + 1);
|
||||
|
||||
const snapshot = await conn.fetchToolsSnapshot(deadline);
|
||||
|
||||
expect(snapshot.tools.map((t) => t.name)).toEqual(['a']);
|
||||
expect(snapshot.complete).toBe(false);
|
||||
expect(listTools).toHaveBeenCalledTimes(1);
|
||||
dateNow.mockRestore();
|
||||
});
|
||||
|
||||
it('returns an incomplete empty snapshot without a request when the deadline has passed', async () => {
|
||||
const listTools = jest.fn();
|
||||
const conn = createConnectionWithListTools(listTools);
|
||||
|
||||
const snapshot = await conn.fetchToolsSnapshot(Date.now() - 1);
|
||||
|
||||
expect(snapshot.tools).toEqual([]);
|
||||
expect(snapshot.complete).toBe(false);
|
||||
expect(listTools).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('stops waiting on an in-flight refresh that outlasts the caller deadline', async () => {
|
||||
mcpConfig.TOOLS_LIST_TIMEOUT_MS = 30000;
|
||||
const conn = createConnectionWithListTools(jest.fn());
|
||||
const mutable = conn as unknown as {
|
||||
toolListChangeGeneration: number;
|
||||
toolListRefreshPromise: Promise<void> | null;
|
||||
};
|
||||
/** A `list_changed` lands mid-fetch, so the ordered read must wait on a refresh... */
|
||||
const listTools = jest.fn(async () => {
|
||||
mutable.toolListChangeGeneration = 1;
|
||||
return { tools: [makeTool('a')] };
|
||||
});
|
||||
conn.client.listTools = listTools;
|
||||
/** ...and that refresh never settles, standing in for one on the connection's own budget. */
|
||||
mutable.toolListRefreshPromise = new Promise<void>(() => {});
|
||||
jest.spyOn(conn.client, 'getServerCapabilities').mockReturnValue({ tools: {} });
|
||||
|
||||
const start = Date.now();
|
||||
const snapshot = await conn.fetchOrderedToolsSnapshot(Date.now() + 30);
|
||||
|
||||
expect(snapshot.complete).toBe(false);
|
||||
expect(Date.now() - start).toBeLessThan(2000);
|
||||
});
|
||||
|
||||
it('stops and warns when the server repeats a cursor instead of looping forever', async () => {
|
||||
const listTools = jest.fn().mockResolvedValue({ tools: [makeTool('x')], nextCursor: 'same' });
|
||||
const conn = createConnectionWithListTools(listTools);
|
||||
|
|
|
|||
|
|
@ -234,31 +234,35 @@ describe('recoverMCPServerCatalogs — bounded, skippable discovery', () => {
|
|||
formatServerTools: jest.fn().mockReturnValue({}),
|
||||
});
|
||||
|
||||
it('bounds each connection attempt the factory makes', async () => {
|
||||
it('bounds each server discovery end to end rather than per attempt', async () => {
|
||||
const discoverServerTools = jest.fn().mockResolvedValue({ tools: [] });
|
||||
const before = Date.now();
|
||||
|
||||
await recoverMCPServerCatalogs(
|
||||
{ user, servers: [{ serverName: 'slow', serverConfig: serverConfig('slow') }] },
|
||||
recoveryDeps(discoverServerTools),
|
||||
);
|
||||
|
||||
expect(discoverServerTools).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ serverName: 'slow', connectionTimeout: 1500 }),
|
||||
);
|
||||
const [options] = discoverServerTools.mock.calls[0];
|
||||
expect(options.serverName).toBe('slow');
|
||||
expect(options.connectionTimeout).toBeUndefined();
|
||||
expect(options.deadlineMs).toBeGreaterThanOrEqual(before + 3000);
|
||||
expect(options.deadlineMs).toBeLessThanOrEqual(Date.now() + 3000);
|
||||
});
|
||||
|
||||
it('keeps a shorter configured initTimeout instead of raising it to the cap', async () => {
|
||||
const discoverServerTools = jest.fn().mockResolvedValue({ tools: [] });
|
||||
const impatient = { ...serverConfig('impatient'), initTimeout: 900 } as ParsedServerConfig;
|
||||
const before = Date.now();
|
||||
|
||||
await recoverMCPServerCatalogs(
|
||||
{ user, servers: [{ serverName: 'impatient', serverConfig: impatient }] },
|
||||
recoveryDeps(discoverServerTools),
|
||||
);
|
||||
|
||||
expect(discoverServerTools).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ connectionTimeout: 900 }),
|
||||
);
|
||||
const [options] = discoverServerTools.mock.calls[0];
|
||||
expect(options.deadlineMs).toBeGreaterThanOrEqual(before + 900);
|
||||
expect(options.deadlineMs).toBeLessThanOrEqual(Date.now() + 900);
|
||||
});
|
||||
|
||||
it('leaves a server the config tier marked unreachable to that tier’s retry window', async () => {
|
||||
|
|
|
|||
|
|
@ -10,18 +10,12 @@ import { getServerCustomUserVars } from '../auth';
|
|||
* `tools/list`, so they burst as readily as passive discovery does. */
|
||||
const CATALOG_FANOUT_CONCURRENCY = 3;
|
||||
/**
|
||||
* Bounds `connection.connect()` only — it is the sole segment of discovery a caller can bound
|
||||
* today. It does NOT bound the whole operation: `discoverToolsInternal` spends this value once
|
||||
* per connection attempt (authenticated, then unauthenticated), and `fetchToolsSnapshot` then
|
||||
* applies its own `TOOLS_LIST_TIMEOUT_MS` (30s) to `tools/list`. A server that connects quickly
|
||||
* and stalls while listing therefore still holds its slot for that longer window.
|
||||
*
|
||||
* Bounding discovery end to end needs a deadline threaded through `MCPConnectionFactory` into
|
||||
* both `connect()` and `fetchToolsSnapshot()`; until that exists, this keeps the common
|
||||
* unreachable case cheap, because recovery targets a server that is reachable and authorized but
|
||||
* whose catalog cache expired, and such a server connects well inside this window.
|
||||
* 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
|
||||
* long regardless of where the server stalls. Recovery targets a server that is reachable and
|
||||
* authorized but whose catalog cache expired, and such a server answers well inside this window.
|
||||
*/
|
||||
const RECOVERY_ATTEMPT_TIMEOUT_MS = 1500;
|
||||
const RECOVERY_BUDGET_MS = 3000;
|
||||
|
||||
export interface MCPServerCatalogRecoveryInput {
|
||||
serverName: string;
|
||||
|
|
@ -73,13 +67,13 @@ interface RecoveryCandidate extends MCPServerCatalogRecoveryInput {
|
|||
customUserVars?: Record<string, string>;
|
||||
}
|
||||
|
||||
/** Bounds one connection attempt, honouring a shorter operator `initTimeout`. */
|
||||
function resolveAttemptTimeout(serverConfig: ParsedServerConfig): number {
|
||||
/** Bounds one server's discovery, honouring a shorter operator `initTimeout`. */
|
||||
function resolveBudget(serverConfig: ParsedServerConfig): number {
|
||||
const { initTimeout } = serverConfig;
|
||||
if (typeof initTimeout === 'number') {
|
||||
return Math.min(initTimeout, RECOVERY_ATTEMPT_TIMEOUT_MS);
|
||||
return Math.min(initTimeout, RECOVERY_BUDGET_MS);
|
||||
}
|
||||
return RECOVERY_ATTEMPT_TIMEOUT_MS;
|
||||
return RECOVERY_BUDGET_MS;
|
||||
}
|
||||
|
||||
async function discoverCandidate(
|
||||
|
|
@ -93,7 +87,7 @@ async function discoverCandidate(
|
|||
serverName,
|
||||
configServers: { [serverName]: serverConfig },
|
||||
customUserVars,
|
||||
connectionTimeout: resolveAttemptTimeout(serverConfig),
|
||||
deadlineMs: Date.now() + resolveBudget(serverConfig),
|
||||
});
|
||||
return [
|
||||
serverName,
|
||||
|
|
@ -112,7 +106,7 @@ async function discoverCandidate(
|
|||
* and is disposed, and the tool cache refuses unfenced writes, so the result is served only to
|
||||
* the requesting user. Recovery therefore stays stateless and individually cheap rather than
|
||||
* scheduling around a result it is not allowed to keep — it skips only what configuration alone
|
||||
* proves pointless, and fails fast on everything else.
|
||||
* proves pointless, and bounds everything else by a per-server budget.
|
||||
*/
|
||||
export async function recoverMCPServerCatalogs(
|
||||
params: { user: IUser; servers: readonly MCPServerCatalogRecoveryInput[] },
|
||||
|
|
|
|||
|
|
@ -1926,6 +1926,12 @@ export class MCPConnection extends EventEmitter {
|
|||
}
|
||||
|
||||
this.transport = await runOutsideTracing(() => this.constructTransport(this.options));
|
||||
/** `dispose()` can land while the transport is still being constructed — it finds nothing
|
||||
* to close and returns, so without this check the attempt would go on to connect and
|
||||
* leave a live connection on a disposed object. Ownership of teardown is ours here. */
|
||||
if (await this.abandonIfDisposed()) {
|
||||
return;
|
||||
}
|
||||
this.patchTransportSend();
|
||||
|
||||
const connectTimeout = this.options.initTimeout ?? DEFAULT_INIT_TIMEOUT;
|
||||
|
|
@ -1937,6 +1943,9 @@ export class MCPConnection extends EventEmitter {
|
|||
),
|
||||
);
|
||||
|
||||
if (await this.abandonIfDisposed()) {
|
||||
return;
|
||||
}
|
||||
this.setupTransportOnMessageHandler();
|
||||
this.connectionState = 'connected';
|
||||
this.emit('connectionChange', 'connected');
|
||||
|
|
@ -2059,6 +2068,33 @@ export class MCPConnection extends EventEmitter {
|
|||
return this.connectPromise;
|
||||
}
|
||||
|
||||
/**
|
||||
* Tears down a transport this attempt created after the connection was already disposed.
|
||||
* Returns whether the caller should abandon the rest of the connect sequence.
|
||||
*/
|
||||
private async abandonIfDisposed(): Promise<boolean> {
|
||||
if (!this.isDisposed) {
|
||||
return false;
|
||||
}
|
||||
logger.debug(`${this.getLogPrefix()} Disposed mid-connect; discarding the transport it opened`);
|
||||
const transport = this.transport;
|
||||
this.transport = null;
|
||||
/** Closing the client only closes a transport the client has already adopted, which it has
|
||||
* not when disposal beat `client.connect()`. Close the transport itself first, or the
|
||||
* session this attempt opened outlives the connection that owned it. */
|
||||
try {
|
||||
await transport?.close();
|
||||
} catch {
|
||||
// Ignore cleanup errors
|
||||
}
|
||||
try {
|
||||
await this.client.close();
|
||||
} catch {
|
||||
// Ignore cleanup errors
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
private patchTransportSend(): void {
|
||||
if (!this.transport) {
|
||||
return;
|
||||
|
|
@ -2348,8 +2384,12 @@ export class MCPConnection extends EventEmitter {
|
|||
* Fetches a bounded tool snapshot while preserving whether every requested page succeeded.
|
||||
* Notification refreshes use `complete` to avoid replacing a known-good cache with an empty or
|
||||
* partial list after a transient `tools/list` failure.
|
||||
*
|
||||
* @param deadlineMs Absolute epoch-ms cap for a caller working to a fixed budget. Pagination
|
||||
* stops at whichever comes first, this or `TOOLS_LIST_TIMEOUT_MS`, and the partial result is
|
||||
* returned as incomplete so it is never published as an authoritative catalog.
|
||||
*/
|
||||
public async fetchToolsSnapshot(): Promise<MCPToolsSnapshot> {
|
||||
public async fetchToolsSnapshot(deadlineMs?: number): Promise<MCPToolsSnapshot> {
|
||||
const maxPages = mcpConfig.TOOLS_LIST_MAX_PAGES;
|
||||
const maxTools = mcpConfig.TOOLS_LIST_MAX_TOOLS;
|
||||
const maxBytes = mcpConfig.TOOLS_LIST_MAX_BYTES;
|
||||
|
|
@ -2357,7 +2397,8 @@ export class MCPConnection extends EventEmitter {
|
|||
* from a `tools/list` that started later. Every app-level publisher reads its ordering off
|
||||
* the snapshot it received, which is the only way to know when the data was actually read. */
|
||||
const ordering = await this.reserveToolsPublicationRevision();
|
||||
const deadline = Date.now() + mcpConfig.TOOLS_LIST_TIMEOUT_MS;
|
||||
const budgetDeadline = Date.now() + mcpConfig.TOOLS_LIST_TIMEOUT_MS;
|
||||
const deadline = deadlineMs != null ? Math.min(budgetDeadline, deadlineMs) : budgetDeadline;
|
||||
const allTools: MCPListToolsResult['tools'] = [];
|
||||
const seenCursors = new Set<string>();
|
||||
let cursor: string | undefined;
|
||||
|
|
@ -2467,11 +2508,14 @@ export class MCPConnection extends EventEmitter {
|
|||
* Returns a complete snapshot that cannot precede a concurrent `list_changed` refresh.
|
||||
* If the notification refresh cannot complete, the caller receives an incomplete snapshot
|
||||
* instead of publishing a request result that may already be stale.
|
||||
*
|
||||
* @param deadlineMs Absolute epoch-ms cap; see `fetchToolsSnapshot`. It also bounds the wait
|
||||
* for a concurrent refresh, so a caller on a budget is never held by another caller's fetch.
|
||||
*/
|
||||
public async fetchOrderedToolsSnapshot(): Promise<MCPToolsSnapshot> {
|
||||
public async fetchOrderedToolsSnapshot(deadlineMs?: number): Promise<MCPToolsSnapshot> {
|
||||
const startEpoch = this.toolListRefreshEpoch;
|
||||
const startGeneration = this.toolListChangeGeneration;
|
||||
const snapshot = await this.fetchToolsSnapshot();
|
||||
const snapshot = await this.fetchToolsSnapshot(deadlineMs);
|
||||
|
||||
if (
|
||||
startEpoch === this.toolListRefreshEpoch &&
|
||||
|
|
@ -2484,12 +2528,22 @@ export class MCPConnection extends EventEmitter {
|
|||
startEpoch === this.toolListRefreshEpoch &&
|
||||
this.handledToolListChangeGeneration < this.toolListChangeGeneration
|
||||
) {
|
||||
if (deadlineMs != null && Date.now() >= deadlineMs) {
|
||||
break;
|
||||
}
|
||||
this.startToolListRefresh();
|
||||
const refresh = this.toolListRefreshPromise;
|
||||
if (!refresh) {
|
||||
break;
|
||||
}
|
||||
await refresh;
|
||||
/** The refresh runs on the connection's own budget, not the caller's, so a refresh already
|
||||
* in flight can outlast this deadline. Stop waiting on it rather than adopting its budget;
|
||||
* it keeps running for whoever else wants it and this caller reports an incomplete read. */
|
||||
if (deadlineMs == null) {
|
||||
await refresh;
|
||||
} else if (!(await this.settlesBefore(refresh, deadlineMs))) {
|
||||
break;
|
||||
}
|
||||
if (this.toolListRefreshRetryTimer) {
|
||||
break;
|
||||
}
|
||||
|
|
@ -2513,6 +2567,20 @@ export class MCPConnection extends EventEmitter {
|
|||
return { tools: [], complete: false };
|
||||
}
|
||||
|
||||
/** Waits for `promise` only until `deadlineMs`, reporting whether it settled in time. */
|
||||
private async settlesBefore(promise: Promise<unknown>, deadlineMs: number): Promise<boolean> {
|
||||
let timer: NodeJS.Timeout | undefined;
|
||||
const expiry = new Promise<false>((resolve) => {
|
||||
timer = setTimeout(() => resolve(false), Math.max(0, deadlineMs - Date.now()));
|
||||
timer.unref?.();
|
||||
});
|
||||
try {
|
||||
return await Promise.race([promise.then(() => true), expiry]);
|
||||
} finally {
|
||||
clearTimeout(timer);
|
||||
}
|
||||
}
|
||||
|
||||
private warnToolsListBudgetExceeded(reason: string, toolCount: number): void {
|
||||
logger.warn(
|
||||
`${this.getLogPrefix()} Stopping tools/list pagination because the ${reason} budget was reached after ${toolCount} tool(s).`,
|
||||
|
|
|
|||
|
|
@ -227,6 +227,10 @@ export interface UserConnectionContext {
|
|||
requestScopedConnections?: RequestScopedMCPConnectionStore;
|
||||
graphTokenResolver?: GraphTokenResolver;
|
||||
connectionTimeout?: number;
|
||||
/** Absolute epoch-ms bound on the whole connect-and-list operation. `connectionTimeout` bounds
|
||||
* only a single `connect()`, so a caller that must return within a fixed budget sets this to
|
||||
* cap every segment, including `tools/list` pagination and the unauthenticated fallback. */
|
||||
deadlineMs?: number;
|
||||
}
|
||||
|
||||
export interface RequestScopedMCPConnectionStore {
|
||||
|
|
@ -284,6 +288,8 @@ export interface ToolDiscoveryOptions {
|
|||
requestBody?: RequestBody;
|
||||
graphTokenResolver?: GraphTokenResolver;
|
||||
connectionTimeout?: number;
|
||||
/** Absolute epoch-ms bound on the whole discovery operation; see `UserConnectionContext`. */
|
||||
deadlineMs?: number;
|
||||
/** Pre-resolved config-source servers for tenant-scoped lookup */
|
||||
configServers?: Record<string, ParsedServerConfig>;
|
||||
oboTokenResolver?: OboTokenResolver;
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue