fix: mcp teggCtxLifecycleMiddleware (#312)

akitaSummer · web-flow · commit 53043840c3aa · 2025-05-07T13:59:49.000+08:00
diff --git a/core/types/controller-decorator/MCPToolParams.ts b/core/types/controller-decorator/MCPToolParams.ts
@@ -3,7 +3,7 @@ import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js';
 import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
 
 export type ToolArgs<T extends Parameters<McpServer['tool']>['2']> = objectOutputType<T, ZodTypeAny>;
-export type ToolExtra = Parameters<Parameters<McpServer['tool']>['3']>['1'];
+export type ToolExtra = Parameters<Parameters<McpServer['tool']>['4']>['1'];
 
 export type MCPToolResponse = CallToolResult;
 
diff --git a/plugin/controller/lib/impl/mcp/MCPControllerRegister.ts b/plugin/controller/lib/impl/mcp/MCPControllerRegister.ts
@@ -122,6 +122,7 @@ export class MCPControllerRegister implements ControllerRegister {
       sessionIdGenerator: undefined,
     });
     self.statelessTransport = transport;
+    const mw = self.app.middleware.teggCtxLifecycleMiddleware();
     const initHandler = async (ctx: Context) => {
       if (MCPControllerRegister.hooks.length > 0) {
         for (const hook of MCPControllerRegister.hooks) {
@@ -135,7 +136,9 @@ export class MCPControllerRegister implements ControllerRegister {
       });
 
       await ctx.app.ctxStorage.run(ctx, async () => {
-        await self.statelessTransport.handleRequest(ctx.req, ctx.res);
+        await mw(ctx, async () => {
+          await self.statelessTransport.handleRequest(ctx.req, ctx.res);
+        });
       });
       return;
     };
@@ -176,6 +179,7 @@ export class MCPControllerRegister implements ControllerRegister {
   mcpStreamServerInit() {
     const allRouterFunc = this.router.all;
     const self = this;
+    const mw = self.app.middleware.teggCtxLifecycleMiddleware();
     const initHandler = async (ctx: Context) => {
       ctx.respond = false;
       if (MCPControllerRegister.hooks.length > 0) {
@@ -233,7 +237,9 @@ export class MCPControllerRegister implements ControllerRegister {
           await self.mcpServer.connect(transport);
 
           await ctx.app.ctxStorage.run(ctx, async () => {
-            await transport.handleRequest(ctx.req, ctx.res, body);
+            await mw(ctx, async () => {
+              await transport.handleRequest(ctx.req, ctx.res, body);
+            });
           });
         } else {
           ctx.status = 400;
@@ -261,7 +267,9 @@ export class MCPControllerRegister implements ControllerRegister {
             'transfer-encoding': 'chunked',
           });
           await ctx.app.ctxStorage.run(ctx, async () => {
-            await transport.handleRequest(ctx.req, ctx.res);
+            await mw(ctx, async () => {
+              await transport.handleRequest(ctx.req, ctx.res);
+            });
           });
           return;
         }
@@ -327,6 +335,7 @@ export class MCPControllerRegister implements ControllerRegister {
   }
 
   sseCtxStorageRun(ctx: Context, transport: SSEServerTransport | StreamableHTTPServerTransport) {
+    const mw = this.app.middleware.teggCtxLifecycleMiddleware();
     const closeFunc = transport.onclose;
     transport.onclose = (...args) => {
       closeFunc?.(...args);
@@ -337,7 +346,9 @@ export class MCPControllerRegister implements ControllerRegister {
     const messageFunc = transport.onmessage;
     transport.onmessage = async (...args: [ JSONRPCMessage ]) => {
       await ctx.app.ctxStorage.run(ctx, async () => {
-        await messageFunc!(...args);
+        await mw(ctx, async () => {
+          await messageFunc!(...args);
+        });
       });
     };
   }
diff --git a/plugin/controller/test/fixtures/apps/mcp-app/app/controller/AppController.ts b/plugin/controller/test/fixtures/apps/mcp-app/app/controller/AppController.ts
@@ -6,6 +6,8 @@ import {
   ToolExtra,
   ToolArgsSchema,
   Extra,
+  Logger,
+  Inject,
 } from '@eggjs/tegg';
 import z from 'zod';
 
@@ -22,6 +24,9 @@ export const NotificationType = {
 
 @MCPController()
 export class AppController {
+  @Inject()
+  logger: Logger;
+
   @MCPTool({
     name: 'start-notification-stream',
     description:
@@ -49,6 +54,8 @@ export class AppController {
       await sleep(interval);
     }
 
+    this.logger.info('startNotificationStream finish');
+
     return {
       content: [
         {
diff --git a/plugin/controller/test/fixtures/apps/mcp-app/package.json b/plugin/controller/test/fixtures/apps/mcp-app/package.json
@@ -1,3 +1,3 @@
 {
-  "name": "controller-app"
+  "name": "mcp-app"
 }
diff --git a/plugin/controller/test/mcp/mcp.test.ts b/plugin/controller/test/mcp/mcp.test.ts
@@ -1,5 +1,6 @@
 import mm from 'egg-mock';
 import path from 'path';
+import fs from 'fs/promises';
 import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js';
 import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
 import { Client } from '@modelcontextprotocol/sdk/client/index.js';
@@ -274,6 +275,10 @@ describe('plugin/controller/test/mcp/mcp.test.ts', () => {
 
       await streamableTransport.terminateSession();
       await streamableClient.close();
+
+      const logContent = await fs.readFile(path.join(__dirname, '../fixtures/apps/mcp-app/logs/mcp-app/mcp-app-web.log'));
+
+      assert.ok(logContent.includes('startNotificationStream finish'));
     });
 
     it('stateless streamable should work', async () => {
diff --git a/plugin/mcp-proxy/index.ts b/plugin/mcp-proxy/index.ts
@@ -13,7 +13,6 @@ import cluster from 'node:cluster';
 import { MCPControllerRegister, MCPControllerHook } from '@eggjs/tegg-controller-plugin/lib/impl/mcp/MCPControllerRegister';
 import querystring from 'node:querystring';
 import url from 'node:url';
-import { Context as EggContext, Request as EggRequest, Response as EggResponse } from '@eggjs/core';
 
 const MAXIMUM_MESSAGE_SIZE = '4mb';
 
@@ -59,13 +58,7 @@ export const MCPProxyHook: MCPControllerHook = {
     await self.app.mcpProxy.registerClient(id, process.pid);
     self.app.mcpProxy.setProxyHandler(MCPProtocols.SSE, async (req, res) => {
       const sessionId = req.query?.sessionId ?? querystring.parse(url.parse(req.url).query ?? '').sessionId as string;
-      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-      // @ts-ignore
-      self.app.RequestClass = EggRequest;
-      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-      // @ts-ignore
-      self.app.ResponseClass = EggResponse;
-      const ctx = new EggContext(self.app as any, req, res) as any;
+      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);
@@ -129,13 +122,8 @@ export const MCPProxyHook: MCPControllerHook = {
     const sessionId = transport.sessionId!;
     self.streamTransports[sessionId] = transport;
     self.app.mcpProxy.setProxyHandler(MCPProtocols.STREAM, async (req: http.IncomingMessage, res: http.ServerResponse) => {
-      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-      // @ts-ignore
-      self.app.RequestClass = EggRequest;
-      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-      // @ts-ignore
-      self.app.ResponseClass = EggResponse;
-      const ctx = new EggContext(self.app as any, req, res) as any;
+      const mw = self.app.middleware.teggCtxLifecycleMiddleware();
+      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);
@@ -171,7 +159,9 @@ export const MCPProxyHook: MCPControllerHook = {
         }
         if (transport) {
           await self.app.ctxStorage.run(ctx, async () => {
-            await transport.handleRequest(ctx.req, ctx.res);
+            await mw(ctx, async () => {
+              await transport.handleRequest(ctx.req, ctx.res);
+            });
           });
         } else {
           res.writeHead(400, { 'content-type': 'application/json' });

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,3 @@`
`1`	`1`	`{`
`2`		`- "name": "controller-app"`
	`2`	`+ "name": "mcp-app"`
`3`	`3`	`}`