Skip to content

fix: clean idle mcp stream proxy sessions #461

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
akitaSummer merged 1 commit into from
Jun 26, 2026
Merged

fix: clean idle mcp stream proxy sessions #461

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 10 additions & 3 deletions plugin/controller/lib/impl/mcp/MCPControllerRegister.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,7 @@ export class MCPControllerRegister implements ControllerRegister {
mcpConfig: MCPConfig;
streamTransports: Record<string, StreamableHTTPServerTransport> = {};
private streamServers: Partial<Record<string, { close(): Promise<void> }>> = {};
private streamServerNames: Partial<Record<string, string>> = {};
private streamCleanupTimers: Partial<Record<string, NodeJS.Timeout>> = {};
private streamClosePromises: Partial<Record<string, Promise<void>>> = {};
// eslint-disable-next-line no-spaced-func
Expand Down Expand Up @@ -163,6 +164,7 @@ export class MCPControllerRegister implements ControllerRegister {
this.instance.controllerProtos = [];
this.instance.streamTransports = {};
this.instance.streamServers = {};
this.instance.streamServerNames = {};
this.instance.streamCleanupTimers = {};
this.instance.streamClosePromises = {};
}
Expand Down Expand Up @@ -352,6 +354,10 @@ export class MCPControllerRegister implements ControllerRegister {
self.clearStreamCleanupTimer(sessionId);
self.streamTransports[sessionId] = transport;
self.streamServers[sessionId] = mcpServerHelper.server;
if (name) {
self.streamServerNames[sessionId] = name;
}
self.scheduleStreamMcpServerCleanup(sessionId, name);
if (MCPControllerRegister.hooks.length > 0) {
for (const hook of MCPControllerRegister.hooks) {
await hook.onStreamSessionInitialized?.(
Expand Down Expand Up @@ -423,7 +429,7 @@ export class MCPControllerRegister implements ControllerRegister {
} else if (sessionId) {
const transport = self.streamTransports[sessionId];
if (transport) {
self.clearStreamCleanupTimer(sessionId);
self.scheduleStreamMcpServerCleanup(sessionId, name);
if (MCPControllerRegister.hooks.length > 0) {
for (const hook of MCPControllerRegister.hooks) {
await hook.preHandle?.(self.app.currentContext);
Expand Down Expand Up @@ -581,11 +587,11 @@ export class MCPControllerRegister implements ControllerRegister {
}
}

private scheduleStreamMcpServerCleanup(sessionId: string | undefined, name?: string) {
scheduleStreamMcpServerCleanup(sessionId: string | undefined, name?: string) {
if (!sessionId || !this.streamTransports[sessionId]) {
return;
}
const timeout = this.mcpConfig.getStreamSessionIdleTimeout(name);
const timeout = this.mcpConfig.getStreamSessionIdleTimeout(name ?? this.streamServerNames[sessionId]);
if (timeout <= 0) {
return;
}
Expand Down Expand Up @@ -641,6 +647,7 @@ export class MCPControllerRegister implements ControllerRegister {
this.clearStreamCleanupTimer(sessionId);
delete this.streamTransports[sessionId];
delete this.streamServers[sessionId];
delete this.streamServerNames[sessionId];
if (this.pingIntervals[sessionId]) {
clearInterval(this.pingIntervals[sessionId]);
delete this.pingIntervals[sessionId];
Expand Down
52 changes: 52 additions & 0 deletions plugin/controller/test/mcp/mcp.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -477,6 +477,58 @@ describe('plugin/controller/test/mcp/mcp.test.ts', () => {
}
});

it('streamable idle session should be cleaned while client stream is still open', async () => {
const streamableClient = new Client({
name: 'streamable-open-idle-client',
version: '1.0.0',
});
const baseUrl = await app.httpRequest()
.post('/mcp/stream').url;
const streamableTransport = new StreamableHTTPClientTransport(
new URL(baseUrl),
{
authProvider: {
get redirectUrl() { return 'http://localhost/callback'; },
get clientMetadata() { return { redirect_uris: [ 'http://localhost/callback' ] }; },
clientInformation: () => ({ client_id: 'test-client-id', client_secret: 'test-client-secret' }),
tokens: () => {
return {
access_token: Buffer.from('akita').toString('base64'),
token_type: 'Bearer',
};
},
// eslint-disable-next-line @typescript-eslint/no-empty-function
saveTokens: () => {},
// eslint-disable-next-line @typescript-eslint/no-empty-function
redirectToAuthorization: () => {},
// eslint-disable-next-line @typescript-eslint/no-empty-function
saveCodeVerifier: () => {},
codeVerifier: () => '',
},
requestInit: { headers: { 'custom-session-id': 'custom-session-id-open-idle' } },
},
);
const register = MCPControllerRegister.instance!;
const originalIdleTimeout = (register.mcpConfig as any)._streamSessionIdleTimeout;
(register.mcpConfig as any)._streamSessionIdleTimeout = 50;
try {
await streamableClient.connect(streamableTransport);
await listTools(streamableClient);
assert.equal(streamableTransport.sessionId, 'custom-session-id-open-idle');
assert.ok(register.streamTransports['custom-session-id-open-idle']);
assert.ok((register as any).streamCleanupTimers['custom-session-id-open-idle']);

await new Promise(resolve => setTimeout(resolve, 200));

assert.equal(register.streamTransports['custom-session-id-open-idle'], undefined);
assert.equal(register.pingIntervals['custom-session-id-open-idle'], undefined);
assert.equal((register as any).streamCleanupTimers['custom-session-id-open-idle'], undefined);
} finally {
(register.mcpConfig as any)._streamSessionIdleTimeout = originalIdleTimeout;
await streamableClient.close().catch(() => undefined);
}
});

it('stateless streamable should work', async () => {
const streamableClient = new Client({
name: 'streamable-demo-client',
Expand Down
213 changes: 121 additions & 92 deletions plugin/mcp-proxy/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import { APIClientBase } from 'cluster-client';
import awaitEvent from 'await-event';
import { MCPProxyDataClient } from './lib/MCPProxyDataClient';
import { Application, Context, EggLogger, Messenger } from 'egg';
import getRawBody from 'raw-body';
Expand All @@ -25,6 +24,8 @@ export interface MCPProxyPayload {
}

type ProxyAction = 'MCP_STDIO_PROXY' | 'MCP_SEE_PROXY' | 'MCP_STREAM_PROXY';
type StreamProxyHandler = StreamableHTTPServerTransport['handleRequest'];
type SseProxyHandler = SSEServerTransport['handlePostMessage'];

export interface ProxyMessageOptions {
detail: ClientDetail;
Expand Down Expand Up @@ -64,27 +65,105 @@ function buildProxyRequestHeaders(headers: http.IncomingHttpHeaders, action: Pro
return proxyHeaders;
}

const streamProxyHandlerMap = new WeakMap<MCPControllerRegister, StreamProxyHandler>();
const sseProxyHandlerMap = new WeakMap<MCPControllerRegister, SseProxyHandler>();

export const MCPProxyHook: MCPControllerHook = {
async preSSEInitHandle(ctx, transport, self) {
const id = transport.sessionId;
// cluster proxy
await self.app.mcpProxy.registerClient(id, process.pid);
self.app.mcpProxy.setProxyHandler(MCPProtocols.SSE, async (req, res) => {
const sessionId = querystring.parse(url.parse(req.url!).query ?? '').sessionId as string;
const ctx = self.app.createContext(req, res) as unknown as Context;
if (MCPControllerRegister.hooks.length > 0) {
for (const hook of MCPControllerRegister.hooks) {
await hook.preProxy?.(ctx, req, res);
function getSseProxyHandler(self: MCPControllerRegister): SseProxyHandler {
let handler = sseProxyHandlerMap.get(self);
if (handler) {
return handler;
}
handler = async (req, res) => {
const sessionId = querystring.parse(url.parse(req.url!).query ?? '').sessionId as string;
const ctx = self.app.createContext(req, res) as unknown as Context;
if (MCPControllerRegister.hooks.length > 0) {
for (const hook of MCPControllerRegister.hooks) {
await hook.preProxy?.(ctx, req, res);
}
}
let transport: SSEServerTransport;
const existingTransport = self.transports[sessionId];
if (existingTransport instanceof SSEServerTransport) {
transport = existingTransport;
} else {
// https://modelcontextprotocol.io/docs/concepts/architecture#error-handling
res.writeHead(400).end(JSON.stringify({
jsonrpc: '2.0',
error: {
code: -32000,
message: 'Bad Request: Session exists but uses a different transport protocol',
},
id: null,
}));
return;
}
if (transport) {
try {
await self.transports[sessionId].handlePostMessage(req, res);
} catch (error) {
self.app.logger.error('Error handling MCP message', error);
if (!ctx.res.headersSent) {
ctx.status = 500;
ctx.body = {
jsonrpc: '2.0',
error: {
code: -32603,
message: `Internal error: ${error.message}`,
},
id: null,
};
}
}
let transport: SSEServerTransport;
const existingTransport = self.transports[sessionId];
if (existingTransport instanceof SSEServerTransport) {
} else {
res.writeHead(404).end(JSON.stringify({
jsonrpc: '2.0',
error: {
code: -32602,
message: 'Bad Request: No transport found for sessionId',
},
id: null,
}));
}
Comment thread
akitaSummer marked this conversation as resolved.
Comment thread
akitaSummer marked this conversation as resolved.
};
sseProxyHandlerMap.set(self, handler);
return handler;
}

function getStreamProxyHandler(self: MCPControllerRegister): StreamProxyHandler {
let handler = streamProxyHandlerMap.get(self);
if (handler) {
return handler;
}
handler = async (req: http.IncomingMessage, res: http.ServerResponse) => {
let mw = self.app.middleware.teggCtxLifecycleMiddleware();
if (self.globalMiddlewares) {
mw = compose([ mw, self.globalMiddlewares ]);
}
const ctx = self.app.createContext(req, res) as unknown as Context;
if (MCPControllerRegister.hooks.length > 0) {
for (const hook of MCPControllerRegister.hooks) {
await hook.preProxy?.(ctx, req, res);
}
}
const sessionId = req.headers['mcp-session-id'] as string | undefined;
if (!sessionId) {
res.writeHead(500, { 'content-type': 'application/json' });
res.end(JSON.stringify({
jsonrpc: '2.0',
error: {
code: -32603,
message: 'session id not have and run in proxy',
},
id: null,
}));
} else {
let transport: StreamableHTTPServerTransport;
const existingTransport = self.streamTransports[sessionId];
if (existingTransport instanceof StreamableHTTPServerTransport) {
transport = existingTransport;
} else {
// https://modelcontextprotocol.io/docs/concepts/architecture#error-handling
res.writeHead(400).end(JSON.stringify({
res.writeHead(400, { 'content-type': 'application/json' });
res.end(JSON.stringify({
jsonrpc: '2.0',
error: {
code: -32000,
Expand All @@ -95,24 +174,20 @@ export const MCPProxyHook: MCPControllerHook = {
return;
}
if (transport) {
self.scheduleStreamMcpServerCleanup(sessionId);
try {
await self.transports[sessionId].handlePostMessage(req, res);
} catch (error) {
self.app.logger.error('Error handling MCP message', error);
if (!ctx.res.headersSent) {
ctx.status = 500;
ctx.body = {
jsonrpc: '2.0',
error: {
code: -32603,
message: `Internal error: ${error.message}`,
},
id: null,
};
}
await self.app.ctxStorage.run(ctx, async () => {
await mw(ctx, async () => {
await transport.handleRequest(ctx.req, ctx.res);
await self.waitResponseClosed(ctx.res);
});
});
} finally {
self.scheduleStreamMcpServerCleanup(sessionId);
Comment thread
akitaSummer marked this conversation as resolved.
}
} else {
res.writeHead(404).end(JSON.stringify({
res.writeHead(400, { 'content-type': 'application/json' });
res.end(JSON.stringify({
jsonrpc: '2.0',
error: {
code: -32602,
Expand All @@ -121,7 +196,19 @@ export const MCPProxyHook: MCPControllerHook = {
id: null,
}));
}
});
}
};
streamProxyHandlerMap.set(self, handler);
return handler;
}


export const MCPProxyHook: MCPControllerHook = {
async preSSEInitHandle(ctx, transport, self) {
const id = transport.sessionId;
// cluster proxy
await self.app.mcpProxy.registerClient(id, process.pid);
self.app.mcpProxy.setProxyHandler(MCPProtocols.SSE, getSseProxyHandler(self));
ctx.res.once('close', () => {
self.clearSseMcpServer(transport)
.catch(error => self.app.logger.error('[mcp-proxy] clear SSE MCP server failed: %s', error.message));
Expand All @@ -132,65 +219,7 @@ export const MCPProxyHook: MCPControllerHook = {
async onStreamSessionInitialized(_ctx, transport, _server, self) {
const sessionId = transport.sessionId!;
self.streamTransports[sessionId] = transport;
self.app.mcpProxy.setProxyHandler(MCPProtocols.STREAM, async (req: http.IncomingMessage, res: http.ServerResponse) => {
let mw = self.app.middleware.teggCtxLifecycleMiddleware();
if (self.globalMiddlewares) {
mw = compose([ mw, self.globalMiddlewares ]);
}
const ctx = self.app.createContext(req, res) as unknown as Context;
if (MCPControllerRegister.hooks.length > 0) {
for (const hook of MCPControllerRegister.hooks) {
await hook.preProxy?.(ctx, req, res);
}
}
const sessionId = req.headers['mcp-session-id'] as string | undefined;
if (!sessionId) {
res.writeHead(500, { 'content-type': 'application/json' });
res.end(JSON.stringify({
jsonrpc: '2.0',
error: {
code: -32603,
message: 'session id not have and run in proxy',
},
id: null,
}));
} else {
let transport: StreamableHTTPServerTransport;
const existingTransport = self.streamTransports[sessionId];
if (existingTransport instanceof StreamableHTTPServerTransport) {
transport = existingTransport;
} else {
res.writeHead(400, { 'content-type': 'application/json' });
res.end(JSON.stringify({
jsonrpc: '2.0',
error: {
code: -32000,
message: 'Bad Request: Session exists but uses a different transport protocol',
},
id: null,
}));
return;
}
if (transport) {
await self.app.ctxStorage.run(ctx, async () => {
await mw(ctx, async () => {
await transport.handleRequest(ctx.req, ctx.res);
await awaitEvent(ctx.res, 'close');
});
});
} else {
res.writeHead(400, { 'content-type': 'application/json' });
res.end(JSON.stringify({
jsonrpc: '2.0',
error: {
code: -32602,
message: 'Bad Request: No transport found for sessionId',
},
id: null,
}));
}
}
});
self.app.mcpProxy.setProxyHandler(MCPProtocols.STREAM, getStreamProxyHandler(self));
await self.app.mcpProxy.registerClient(sessionId, process.pid);
const closeFunc = transport.onclose;
transport.onclose = async () => {
Expand Down
Loading