feat: add mcp middleware hook

plugin/controller/app.ts

@@ -132,6 +132,8 @@ export default class ControllerAppBootHook {
     // The HTTPControllerRegister will collect all the methods
     // and register methods after collect is done.
     HTTPControllerRegister.instance?.doRegister(this.app.rootProtoManager);
+
+    this.app.config.mcp.hooks = MCPControllerRegister.hooks;
   }
 
   async willReady() {

plugin/controller/app/middleware/mcp_body_middleware.ts

@@ -18,7 +18,24 @@ export default () => {
     });
     if (res) {
       ctx.disableBodyParser = true;
+      try {
+        for (const hook of ctx.app.config.mcp.hooks) {
+          await hook.middlewareStart?.(ctx);
+        }
+        await next();
+        if (!ctx.mcpArg) {
+          ctx.mcpArg = JSON.parse(ctx.response.header['mcp-proxy-arg'] as string ?? '{}');
+        }
+        for (const hook of ctx.app.config.mcp.hooks) {
+          await hook.middlewareEnd?.(ctx);
+        }
+      } catch (e) {
+        for (const hook of ctx.app.config.mcp.hooks) {
+          await hook.middlewareError?.(ctx, e);
+        }
+      }
+    } else {
+      await next();
     }
-    await next();
   };
 };

plugin/controller/lib/impl/mcp/MCPControllerRegister.ts

@@ -19,7 +19,7 @@ import { ControllerRegister } from '../../ControllerRegister';
 import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js';
 import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
 import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
-import { isInitializeRequest, isJSONRPCRequest, JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js';
+import { isInitializeRequest, isJSONRPCRequest, JSONRPCMessage, MessageExtraInfo } from '@modelcontextprotocol/sdk/types.js';
 import { MCPProtocols } from '@eggjs/mcp-proxy/types';
 import awaitEvent from 'await-event';
 import compose from 'koa-compose';
@@ -43,6 +43,11 @@ export interface MCPControllerHook {
   preProxy?: (ctx: Context, proxyReq: http.IncomingMessage, proxyResp: http.ServerResponse) => Promise<void>
   schemaLoader?: (controllerMeta: MCPControllerMeta, meta: MCPPromptMeta | MCPToolMeta) => Promise<Parameters<McpServer['tool']>['2'] | Parameters<McpServer['prompt']>['2']>
   checkAndRunProxy?: (ctx: Context, type: MCPProtocols, sessionId: string) => Promise<boolean>;
+
+  // middleware
+  middlewareStart?: (ctx: Context) => Promise<void>
+  middlewareEnd?: (ctx: Context) => Promise<void>
+  middlewareError?: (ctx: Context, e: Error) => Promise<void>
 }
 
 class InnerSSEServerTransport extends SSEServerTransport {
@@ -183,6 +188,14 @@ export class MCPControllerRegister implements ControllerRegister {
     if (self.globalMiddlewares) {
       mw = compose([ mw, self.globalMiddlewares ]);
     }
+    const onmessage = transport.onmessage;
+
+    transport.onmessage = async (message: JSONRPCMessage, extra?: MessageExtraInfo) => {
+      if (self.app.currentContext) {
+        self.app.currentContext.mcpArg = message;
+      }
+      onmessage && await onmessage(message, extra);
+    };
     const initHandler = async (ctx: Context) => {
       if (MCPControllerRegister.hooks.length > 0) {
         for (const hook of MCPControllerRegister.hooks) {
@@ -194,7 +207,6 @@ export class MCPControllerRegister implements ControllerRegister {
         'content-type': 'text/event-stream',
         'transfer-encoding': 'chunked',
       });
-
       await ctx.app.ctxStorage.run(ctx, async () => {
         await mw(ctx, async () => {
           await transport.handleRequest(ctx.req, ctx.res);
@@ -253,7 +265,7 @@ export class MCPControllerRegister implements ControllerRegister {
       }
       const sessionId = ctx.req.headers['mcp-session-id'] as string | undefined;
       if (!sessionId) {
-        const ct = contentType.parse(ctx.req.headers['content-type'] ?? '');
+        const ct = contentType.parse(ctx.req.headers['content-type'] || 'application/json');
 
         let body;
 
@@ -307,6 +319,15 @@ export class MCPControllerRegister implements ControllerRegister {
             transport,
           );
 
+          const onmessage = transport.onmessage;
+
+          transport.onmessage = async (message: JSONRPCMessage, extra?: MessageExtraInfo) => {
+            if (self.app.currentContext) {
+              self.app.currentContext.mcpArg = message;
+            }
+            onmessage && await onmessage(message, extra);
+          };
+
           await ctx.app.ctxStorage.run(ctx, async () => {
             await mw(ctx, async () => {
               await transport.handleRequest(ctx.req, ctx.res, body);
@@ -338,6 +359,7 @@ export class MCPControllerRegister implements ControllerRegister {
             'content-type': 'text/event-stream',
             'transfer-encoding': 'chunked',
           });
+
           await ctx.app.ctxStorage.run(ctx, async () => {
             await mw(ctx, async () => {
               await transport.handleRequest(ctx.req, ctx.res);
@@ -442,7 +464,8 @@ export class MCPControllerRegister implements ControllerRegister {
     };
     const messageFunc = transport.onmessage;
     self.sseTransportsRequestMap.set(transport, {});
-    transport.onmessage = async (...args: [JSONRPCMessage]) => {
+    transport.onmessage = async (message: JSONRPCMessage, extra?: MessageExtraInfo) => {
+      const args = [ message, extra ];
       // 这里需要 new 一个新的 ctx，否则 ContextProto 会未被初始化
       const socket = new Socket();
       const req = new IncomingMessage(socket);
@@ -451,6 +474,7 @@ export class MCPControllerRegister implements ControllerRegister {
       req.url = self.mcpConfig.getSseInitPath(name);
       req.headers = {
         ...ctx.req.headers,
+        ...extra?.requestInfo?.headers,
         accept: 'application/json, text/event-stream',
         'content-type': 'application/json',
       };
@@ -462,12 +486,12 @@ export class MCPControllerRegister implements ControllerRegister {
           }
         }
         await mw(newCtx, async () => {
-          messageFunc!(...args);
+          messageFunc!(message, extra);
           if (isJSONRPCRequest(args[0])) {
             const map = self.sseTransportsRequestMap.get(transport)!;
             const wait = new Promise<null>((resolve, reject) => {
-              if ('id' in args[0]) {
-                map[args[0].id] = { resolve, reject };
+              if (extra && 'id' in extra) {
+                map[extra.id as string] = { resolve, reject };
               }
             });
             await wait;
@@ -500,7 +524,16 @@ export class MCPControllerRegister implements ControllerRegister {
         }
         self.app.logger.info('message coming', sessionId);
         try {
-          await self.transports[sessionId].handlePostMessage(ctx.req, ctx.res);
+          const ct = contentType.parse(ctx.req.headers['content-type'] ?? '');
+
+          const rawBody = await getRawBody(ctx.req, {
+            limit: '4mb',
+            encoding: ct.parameters.charset ?? 'utf-8',
+          });
+
+          const body = JSON.parse(rawBody);
+          ctx.mcpArg = body;
+          await self.transports[sessionId].handlePostMessage(ctx.req, ctx.res, body);
         } catch (error) {
           self.app.logger.error('Error handling MCP message', error);
           if (!ctx.res.headersSent) {

plugin/controller/test/fixtures/apps/mcp-app/app/controller/McpController.ts

@@ -16,7 +16,7 @@
   Extra,
   ToolExtra,
 } from '@eggjs/tegg';
 import z from 'zod';
 
 export const PromptType = {
   name: z.string(),
@@ -80,6 +80,12 @@
     };
   }
 
+  @MCPTool()
+  async mockError(): Promise<MCPToolResponse> {
+    throw new Error('mock error');
+  }
+
+
   @MCPTool()
   async echoUser(): Promise<MCPToolResponse> {
     return {

plugin/controller/test/fixtures/apps/mcp-app/config/config.default.js

@@ -1,8 +1,10 @@
 'use strict';
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const { randomUUID } = require('node:crypto');
+// eslint-disable-next-line @typescript-eslint/no-var-requires
+const path = require('node:path');
 
-module.exports = function() {
+module.exports = function(appInfo) {
   const config = {
     keys: 'test key',
     mcp: {
@@ -20,6 +22,17 @@ module.exports = function() {
         },
       },
     },
+    customLogger: {
+      mcpMiddewareStartLogger: {
+        file: path.join(appInfo.root, 'logs', 'tracelog', 'mcpMiddlewareStart.log'),
+      },
+      mcpMiddewareEndLogger: {
+        file: path.join(appInfo.root, 'logs', 'tracelog', 'mcpMiddlewareEnd.log'),
+      },
+      mcpMiddewareErrorLogger: {
+        file: path.join(appInfo.root, 'logs', 'tracelog', 'mcpMiddlewareError.log'),
+      },
+    },
   };
   return config;
 };

plugin/controller/test/fixtures/apps/mcp-app/hook-plugin/lib/MCPControllerHook.ts

@@ -32,6 +32,16 @@ export const GetAlipayTeggHook = (app: Application) => {
     async preProxy(ctx) {
       setUser(ctx);
     },
+    async middlewareStart(ctx) {
+      ctx.mcpStartTime = Date.now();
+      ctx.getLogger('mcpMiddewareStartLogger').info('mcp middleware start');
+    },
+    async middlewareEnd(ctx) {
+      ctx.getLogger('mcpMiddewareEndLogger').info('mcp middleware end, arg: ', JSON.stringify(ctx.mcpArg), `, time: ${Date.now() - ctx.mcpStartTime}`);
+    },
+    async middlewareError(ctx, e) {
+      ctx.getLogger('mcpMiddewareErrorLogger').info('mcp middleware error: ', e);
+    },
   };
 
   return AlipayTeggControllerHook;

plugin/controller/test/mcp/mcp.test.ts

@@ -130,6 +130,10 @@ describe('plugin/controller/test/mcp/mcp.test.ts', () => {
           description: undefined,
           name: 'bar',
         },
+        {
+          description: undefined,
+          name: 'mockError',
+        },
         {
           description: undefined,
           name: 'echoUser',
@@ -219,6 +223,13 @@ describe('plugin/controller/test/mcp/mcp.test.ts', () => {
         ],
       });
       await sseTransport.close();
+
+
+      const middlewareStartTracelog = await fs.readFile(path.join(__dirname, '../fixtures/apps/mcp-app/logs/tracelog/mcpMiddlewareStart.log'), 'utf-8');
+      const middlewareEndTracelog = await fs.readFile(path.join(__dirname, '../fixtures/apps/mcp-app/logs/tracelog/mcpMiddlewareEnd.log'), 'utf-8');
+
+      assert.ok(middlewareStartTracelog.includes('mcp middleware start'));
+      assert.ok(middlewareEndTracelog.includes('mcp middleware end, arg:  {'));
     });
 
     it('streamable should work', async () => {
@@ -269,6 +280,10 @@ describe('plugin/controller/test/mcp/mcp.test.ts', () => {
           description: undefined,
           name: 'bar',
         },
+        {
+          description: undefined,
+          name: 'mockError',
+        },
         {
           description: undefined,
           name: 'echoUser',
@@ -364,6 +379,12 @@ describe('plugin/controller/test/mcp/mcp.test.ts', () => {
       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'));
+
+      const middlewareStartTracelog = await fs.readFile(path.join(__dirname, '../fixtures/apps/mcp-app/logs/tracelog/mcpMiddlewareStart.log'), 'utf-8');
+      const middlewareEndTracelog = await fs.readFile(path.join(__dirname, '../fixtures/apps/mcp-app/logs/tracelog/mcpMiddlewareEnd.log'), 'utf-8');
+
+      assert.ok(middlewareStartTracelog.includes(' POST /mcp/stream] mcp middleware start'));
+      assert.ok(middlewareEndTracelog.includes(' POST /mcp/stream] mcp middleware end, arg: '));
     });
 
     it('stateless streamable should work', async () => {
@@ -412,6 +433,10 @@ describe('plugin/controller/test/mcp/mcp.test.ts', () => {
           description: undefined,
           name: 'bar',
         },
+        {
+          description: undefined,
+          name: 'mockError',
+        },
         {
           description: undefined,
           name: 'echoUser',
@@ -503,6 +528,13 @@ describe('plugin/controller/test/mcp/mcp.test.ts', () => {
 
       await streamableTransport.terminateSession();
       await streamableClient.close();
+
+
+      const middlewareStartTracelog = await fs.readFile(path.join(__dirname, '../fixtures/apps/mcp-app/logs/tracelog/mcpMiddlewareStart.log'), 'utf-8');
+      const middlewareEndTracelog = await fs.readFile(path.join(__dirname, '../fixtures/apps/mcp-app/logs/tracelog/mcpMiddlewareEnd.log'), 'utf-8');
+
+      assert.ok(middlewareStartTracelog.includes(' POST /mcp/stateless/stream] mcp middleware start'));
+      assert.ok(middlewareEndTracelog.includes(' POST /mcp/stateless/stream] mcp middleware end, arg:  {'));
     });
 
     it('multiple sse should work', async () => {
@@ -640,6 +672,13 @@ describe('plugin/controller/test/mcp/mcp.test.ts', () => {
         ],
       });
       await sseTransport.close();
+
+
+      const middlewareStartTracelog = await fs.readFile(path.join(__dirname, '../fixtures/apps/mcp-app/logs/tracelog/mcpMiddlewareStart.log'), 'utf-8');
+      const middlewareEndTracelog = await fs.readFile(path.join(__dirname, '../fixtures/apps/mcp-app/logs/tracelog/mcpMiddlewareEnd.log'), 'utf-8');
+
+      assert.ok(middlewareStartTracelog.includes('mcp middleware start'));
+      assert.ok(middlewareEndTracelog.includes('mcp middleware end'));
     });
 
     it('multiple streamable should work', async () => {
@@ -785,6 +824,13 @@ describe('plugin/controller/test/mcp/mcp.test.ts', () => {
       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'));
+
+
+      const middlewareStartTracelog = await fs.readFile(path.join(__dirname, '../fixtures/apps/mcp-app/logs/tracelog/mcpMiddlewareStart.log'), 'utf-8');
+      const middlewareEndTracelog = await fs.readFile(path.join(__dirname, '../fixtures/apps/mcp-app/logs/tracelog/mcpMiddlewareEnd.log'), 'utf-8');
+
+      assert.ok(middlewareStartTracelog.includes(' POST /mcp/test/stream] mcp middleware start'));
+      assert.ok(middlewareEndTracelog.includes(' POST /mcp/test/stream] mcp middleware end'));
     });
 
     it('multiple stateless streamable should work', async () => {
@@ -924,6 +970,13 @@ describe('plugin/controller/test/mcp/mcp.test.ts', () => {
 
       await streamableTransport.terminateSession();
       await streamableClient.close();
+
+
+      const middlewareStartTracelog = await fs.readFile(path.join(__dirname, '../fixtures/apps/mcp-app/logs/tracelog/mcpMiddlewareStart.log'), 'utf-8');
+      const middlewareEndTracelog = await fs.readFile(path.join(__dirname, '../fixtures/apps/mcp-app/logs/tracelog/mcpMiddlewareEnd.log'), 'utf-8');
+
+      assert.ok(middlewareStartTracelog.includes(' /mcp/test/stateless/stream] mcp middleware start'));
+      assert.ok(middlewareEndTracelog.includes(' /mcp/test/stateless/stream] mcp middleware end'));
     });
   }
 });