diff --git a/plugin/controller/lib/impl/mcp/MCPControllerRegister.ts b/plugin/controller/lib/impl/mcp/MCPControllerRegister.ts index ea62564ea..b18645e91 100644 --- a/plugin/controller/lib/impl/mcp/MCPControllerRegister.ts +++ b/plugin/controller/lib/impl/mcp/MCPControllerRegister.ts @@ -97,6 +97,7 @@ export class MCPControllerRegister implements ControllerRegister { mcpConfig: MCPConfig; streamTransports: Record = {}; private streamServers: Partial }>> = {}; + private streamServerNames: Partial> = {}; private streamCleanupTimers: Partial> = {}; private streamClosePromises: Partial>> = {}; // eslint-disable-next-line no-spaced-func @@ -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 = {}; } @@ -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?.( @@ -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); @@ -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; } @@ -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]; diff --git a/plugin/controller/test/mcp/mcp.test.ts b/plugin/controller/test/mcp/mcp.test.ts index c250e2243..a01f164ad 100644 --- a/plugin/controller/test/mcp/mcp.test.ts +++ b/plugin/controller/test/mcp/mcp.test.ts @@ -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', diff --git a/plugin/mcp-proxy/index.ts b/plugin/mcp-proxy/index.ts index c36fc68ac..82aa8e73f 100644 --- a/plugin/mcp-proxy/index.ts +++ b/plugin/mcp-proxy/index.ts @@ -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'; @@ -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; @@ -64,27 +65,105 @@ function buildProxyRequestHeaders(headers: http.IncomingHttpHeaders, action: Pro return proxyHeaders; } +const streamProxyHandlerMap = new WeakMap(); +const sseProxyHandlerMap = new WeakMap(); -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, + })); + } + }; + 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, @@ -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); } } else { - res.writeHead(404).end(JSON.stringify({ + res.writeHead(400, { 'content-type': 'application/json' }); + res.end(JSON.stringify({ jsonrpc: '2.0', error: { code: -32602, @@ -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)); @@ -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 () => {