Skip to content

Commit 055585f

Browse files
danny-aviladanny-avila
authored
🪢 fix: Tie MCP Cleanup To Resumable Runs (danny-avila#13769)
* fix: Clean up request-scoped MCP connections * test: Format MCP request context spec * refactor: Move MCP request context to API package
1 parent 0537930 commit 055585f

6 files changed

Lines changed: 395 additions & 70 deletions

File tree

‎api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
const { EventEmitter } = require('events');
2+
13
const mockLogger = {
24
debug: jest.fn(),
35
warn: jest.fn(),
@@ -21,6 +23,58 @@ const mockFilterPersistableAbortContent = jest.fn((content) =>
2123
const mockGetConvo = jest.fn();
2224
const mockGetMessages = jest.fn();
2325
const mockSaveMessage = jest.fn();
26+
let mockMCPContexts = new WeakMap();
27+
28+
const mockCreateMCPRequestContext = jest.fn(() => ({
29+
connections: new Map(),
30+
pending: new Map(),
31+
cleanupStarted: false,
32+
cleanupOnResponse: false,
33+
responseCleanupAttached: false,
34+
}));
35+
const mockGetMCPRequestContext = jest.fn((req) => {
36+
if (!req) {
37+
return undefined;
38+
}
39+
40+
let context = mockMCPContexts.get(req);
41+
if (!context) {
42+
context = mockCreateMCPRequestContext();
43+
mockMCPContexts.set(req, context);
44+
}
45+
46+
return context.cleanupStarted ? undefined : context;
47+
});
48+
const mockCleanupMCPRequestContext = jest.fn(async (context) => {
49+
if (!context || context.cleanupStarted) {
50+
return;
51+
}
52+
53+
context.cleanupStarted = true;
54+
const connections = new Set(context.connections.values());
55+
const settled = await Promise.allSettled(context.pending.values());
56+
for (const result of settled) {
57+
if (result.status === 'fulfilled' && result.value) {
58+
connections.add(result.value);
59+
}
60+
}
61+
62+
await Promise.allSettled(Array.from(connections).map((connection) => connection.disconnect?.()));
63+
context.connections.clear();
64+
context.pending.clear();
65+
});
66+
const mockCleanupMCPRequestContextForReq = jest.fn(async (req) => {
67+
const context = mockMCPContexts.get(req);
68+
if (!context) {
69+
return;
70+
}
71+
72+
try {
73+
await mockCleanupMCPRequestContext(context);
74+
} finally {
75+
mockMCPContexts.delete(req);
76+
}
77+
});
2478

2579
jest.mock('@librechat/data-schemas', () => ({
2680
logger: mockLogger,
@@ -32,7 +86,11 @@ jest.mock('@librechat/api', () => ({
3286
buildMessageFiles: jest.fn(() => []),
3387
resolveTitleTiming: jest.fn(() => 'immediate'),
3488
GenerationJobManager: mockGenerationJobManager,
89+
cleanupMCPRequestContext: (...args) => mockCleanupMCPRequestContext(...args),
90+
createMCPRequestContext: (...args) => mockCreateMCPRequestContext(...args),
91+
getMCPRequestContext: (...args) => mockGetMCPRequestContext(...args),
3592
filterPersistableAbortContent: (...args) => mockFilterPersistableAbortContent(...args),
93+
cleanupMCPRequestContextForReq: (...args) => mockCleanupMCPRequestContextForReq(...args),
3694
decrementPendingRequest: (...args) => mockDecrementPendingRequest(...args),
3795
sanitizeMessageForTransmit: jest.fn((message) => message),
3896
checkAndIncrementPendingRequest: (...args) => mockCheckAndIncrementPendingRequest(...args),
@@ -79,10 +137,33 @@ jest.mock('~/models', () => ({
79137
}));
80138

81139
const AgentController = require('../request');
140+
const { getMCPRequestContext } = require('~/server/services/MCPRequestContext');
141+
142+
function createResumableResponse() {
143+
const res = new EventEmitter();
144+
res.headersSent = false;
145+
res.writableEnded = false;
146+
res.finished = false;
147+
res.destroyed = false;
148+
res.json = jest.fn(() => {
149+
res.headersSent = true;
150+
res.writableEnded = true;
151+
res.finished = true;
152+
res.emit('finish');
153+
return res;
154+
});
155+
res.status = jest.fn(() => res);
156+
return res;
157+
}
158+
159+
function nextTick() {
160+
return new Promise((resolve) => setImmediate(resolve));
161+
}
82162

83163
describe('ResumableAgentController resume metadata', () => {
84164
beforeEach(() => {
85165
jest.clearAllMocks();
166+
mockMCPContexts = new WeakMap();
86167
mockCheckAndIncrementPendingRequest.mockResolvedValue({ allowed: true });
87168
mockDecrementPendingRequest.mockResolvedValue(undefined);
88169
mockGetConvo.mockResolvedValue({ createdAt: '2026-06-07T00:00:00.000Z' });
@@ -228,6 +309,47 @@ describe('ResumableAgentController resume metadata', () => {
228309
);
229310
});
230311

312+
it('keeps request-scoped MCP connections until resumable initialization finishes', async () => {
313+
const conversationId = 'conversation-123';
314+
const disconnect = jest.fn().mockResolvedValue(undefined);
315+
const initializeClient = jest.fn(async ({ req, res }) => {
316+
const context = getMCPRequestContext(req, res);
317+
context.connections.set('mcp-server', { disconnect });
318+
319+
await nextTick();
320+
expect(disconnect).not.toHaveBeenCalled();
321+
322+
throw new Error('stop after request-scoped MCP connection');
323+
});
324+
const req = {
325+
user: { id: 'user-123' },
326+
body: {
327+
text: 'Use a BODY-scoped MCP server.',
328+
messageId: 'user-message',
329+
parentMessageId: 'parent-message',
330+
conversationId,
331+
endpointOption: {
332+
endpoint: 'agents',
333+
modelOptions: { model: 'gpt-4.1' },
334+
},
335+
},
336+
config: {},
337+
};
338+
const res = createResumableResponse();
339+
340+
await AgentController(req, res, jest.fn(), initializeClient, null);
341+
342+
expect(res.json).toHaveBeenCalledWith({
343+
streamId: conversationId,
344+
conversationId,
345+
status: 'started',
346+
});
347+
expect(disconnect).toHaveBeenCalledTimes(1);
348+
expect(disconnect.mock.invocationCallOrder[0]).toBeLessThan(
349+
mockDecrementPendingRequest.mock.invocationCallOrder[0],
350+
);
351+
});
352+
231353
it('stores model spec icon fallbacks and agent ids in early resume metadata', async () => {
232354
const conversationId = 'conversation-123';
233355
const initializeClient = jest.fn().mockRejectedValue(new Error('stop before tool loading'));

‎api/server/controllers/agents/request.js‎

Lines changed: 20 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,10 @@ const {
1313
isUnpersistedPreliminaryParent,
1414
} = require('@librechat/api');
1515
const { disposeClient, clientRegistry, requestDataMap } = require('~/server/cleanup');
16+
const {
17+
getMCPRequestContext,
18+
cleanupMCPRequestContextForReq,
19+
} = require('~/server/services/MCPRequestContext');
1620
const { handleAbortError } = require('~/server/middleware');
1721
const { logViolation } = require('~/cache');
1822
const { saveMessage, getMessages, getConvo } = require('~/models');
@@ -139,6 +143,14 @@ function getAgentResponseModel(req, endpointOption) {
139143
return getEndpointResponseModel(endpointOption);
140144
}
141145

146+
async function finishResumableRequest(req, userId) {
147+
try {
148+
await cleanupMCPRequestContextForReq(req);
149+
} finally {
150+
await decrementPendingRequest(userId);
151+
}
152+
}
153+
142154
function rejectPreliminaryParentMessageId(res) {
143155
return res.status(409).json({
144156
error:
@@ -209,6 +221,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
209221
const job = await GenerationJobManager.createJob(streamId, userId, conversationId);
210222
const jobCreatedAt = job.createdAt; // Capture creation time to detect job replacement
211223
req._resumableStreamId = streamId;
224+
getMCPRequestContext(req, undefined, { cleanupOnResponse: false });
212225

213226
// Send JSON response IMMEDIATELY so client can connect to SSE stream
214227
// This is critical: tool loading (MCP OAuth) may emit events that the client needs to receive
@@ -316,7 +329,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
316329

317330
if (job.abortController.signal.aborted) {
318331
GenerationJobManager.completeJob(streamId, 'Request aborted during initialization');
319-
await decrementPendingRequest(userId);
332+
await finishResumableRequest(req, userId);
320333
return;
321334
}
322335

@@ -552,7 +565,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
552565
acceptsTitleEvents = false;
553566
resolveConvoReady();
554567
// Still decrement pending request since we incremented at start
555-
await decrementPendingRequest(userId);
568+
await finishResumableRequest(req, userId);
556569
if (immediateTitlePromise) {
557570
immediateTitlePromise.finally(() => {
558571
if (client) {
@@ -602,7 +615,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
602615

603616
await GenerationJobManager.emitDone(streamId, finalEvent);
604617
GenerationJobManager.completeJob(streamId);
605-
await decrementPendingRequest(userId);
618+
await finishResumableRequest(req, userId);
606619
} else {
607620
const finalEvent = {
608621
final: true,
@@ -622,7 +635,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
622635

623636
await GenerationJobManager.emitDone(streamId, finalEvent);
624637
GenerationJobManager.completeJob(streamId, 'Request aborted');
625-
await decrementPendingRequest(userId);
638+
await finishResumableRequest(req, userId);
626639
}
627640

628641
if (titleTiming === 'immediate') {
@@ -680,7 +693,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
680693
GenerationJobManager.completeJob(streamId, error.message);
681694
}
682695

683-
await decrementPendingRequest(userId);
696+
await finishResumableRequest(req, userId);
684697

685698
// Defer disposal until any immediate title settles (it holds the run/req).
686699
if (immediateTitlePromise) {
@@ -704,7 +717,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
704717
`[ResumableAgentController] Unhandled error in background generation: ${err.message}`,
705718
);
706719
GenerationJobManager.completeJob(streamId, err.message);
707-
await decrementPendingRequest(userId);
720+
await finishResumableRequest(req, userId);
708721
});
709722
} catch (error) {
710723
logger.error('[ResumableAgentController] Initialization error:', error);
@@ -715,7 +728,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
715728
await GenerationJobManager.emitError(streamId, error.message || 'Failed to start generation');
716729
}
717730
GenerationJobManager.completeJob(streamId, error.message);
718-
await decrementPendingRequest(userId);
731+
await finishResumableRequest(req, userId);
719732
if (client) {
720733
disposeClient(client);
721734
}

‎api/server/services/MCPRequestContext.js‎

Lines changed: 7 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -1,68 +1,12 @@
1-
const { logger } = require('@librechat/data-schemas');
2-
3-
const MCP_REQUEST_CONTEXT = Symbol.for('librechat.mcpRequestContext');
4-
5-
function createMCPRequestContext() {
6-
return {
7-
connections: new Map(),
8-
pending: new Map(),
9-
cleanupStarted: false,
10-
};
11-
}
12-
13-
async function cleanupMCPRequestContext(context) {
14-
if (!context || context.cleanupStarted) {
15-
return;
16-
}
17-
context.cleanupStarted = true;
18-
19-
const connections = new Set(context.connections.values());
20-
const pending = Array.from(context.pending.values());
21-
if (pending.length > 0) {
22-
const settled = await Promise.allSettled(pending);
23-
for (const result of settled) {
24-
if (result.status === 'fulfilled' && result.value) {
25-
connections.add(result.value);
26-
}
27-
}
28-
}
29-
30-
await Promise.allSettled(
31-
Array.from(connections).map(async (connection) => {
32-
try {
33-
await connection.disconnect();
34-
} catch (error) {
35-
logger.warn('[MCP Request Context] Failed to disconnect request-scoped connection', error);
36-
}
37-
}),
38-
);
39-
40-
context.connections.clear();
41-
context.pending.clear();
42-
}
43-
44-
function getMCPRequestContext(req, res) {
45-
if (!req) {
46-
return undefined;
47-
}
48-
49-
if (!req[MCP_REQUEST_CONTEXT]) {
50-
const context = createMCPRequestContext();
51-
req[MCP_REQUEST_CONTEXT] = context;
52-
53-
const cleanup = () => {
54-
cleanupMCPRequestContext(context).catch((error) => {
55-
logger.warn('[MCP Request Context] Cleanup failed', error);
56-
});
57-
};
58-
res?.once?.('finish', cleanup);
59-
res?.once?.('close', cleanup);
60-
}
61-
62-
return req[MCP_REQUEST_CONTEXT];
63-
}
1+
const {
2+
cleanupMCPRequestContextForReq,
3+
cleanupMCPRequestContext,
4+
createMCPRequestContext,
5+
getMCPRequestContext,
6+
} = require('@librechat/api');
647

658
module.exports = {
9+
cleanupMCPRequestContextForReq,
6610
cleanupMCPRequestContext,
6711
createMCPRequestContext,
6812
getMCPRequestContext,

‎packages/api/src/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ export * from './mcp/zod';
1919
export * from './mcp/errors';
2020
export * from './mcp/cache';
2121
export * from './mcp/tools';
22+
export * from './mcp/request';
2223
/* Utilities */
2324
export * from './mcp/utils';
2425
export * from './utils';

0 commit comments

Comments
 (0)