fix: update sessionIdGenerator

Author: akitaSummer

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

@@ -1,3 +1,4 @@
+import type { Context } from 'egg';
 import { randomUUID } from 'node:crypto';
 import { EventStore } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
 import { InMemoryEventStore } from '@modelcontextprotocol/sdk/examples/shared/inMemoryEventStore.js';
@@ -7,7 +8,7 @@ export interface MCPConfigOptions {
   sseMessagePath: string;
   streamPath: string;
   statelessStreamPath: string;
-  sessionIdGenerator?: () => string;
+  sessionIdGenerator?: ((ctx?: Context) => string) | ((...args: any[]) => string);
   eventStore?: EventStore;
   sseHeartTime?: number;
 }
@@ -17,7 +18,7 @@ export class MCPConfig {
   private _sseMessagePath: string;
   private _streamPath: string;
   private _statelessStreamPath: string;
-  private _sessionIdGenerator: () => string;
+  private _sessionIdGenerator: ((ctx?: Context) => string) | ((...args: any[]) => string);
   private _eventStore: EventStore;
   private _sseHeartTime: number;
 

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

@@ -220,72 +220,58 @@ 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'] ?? '');
 
-        let body;
+      let body;
 
-        try {
-          const rawBody = await getRawBody(ctx.req, {
-            limit: '4mb',
-            encoding: ct.parameters.charset ?? 'utf-8',
-          });
-
-          body = JSON.parse(rawBody);
-        } catch (e) {
-          ctx.status = 400;
-          ctx.body = {
-            jsonrpc: '2.0',
-            error: {
-              code: -32000,
-              message: `Bad Request: body should is json, ${e.toString()}`,
-            },
-            id: null,
-          };
-          return;
-        }
+      try {
+        const rawBody = await getRawBody(ctx.req, {
+          limit: '4mb',
+          encoding: ct.parameters.charset ?? 'utf-8',
+        });
 
-        if (isInitializeRequest(body)) {
-          ctx.respond = false;
-          const eventStore = this.mcpConfig.eventStore;
-          const self = this;
-          const transport = new StreamableHTTPServerTransport({
-            sessionIdGenerator: () => this.mcpConfig.sessionIdGenerator(),
-            eventStore,
-            onsessioninitialized: async () => {
-              if (MCPControllerRegister.hooks.length > 0) {
-                for (const hook of MCPControllerRegister.hooks) {
-                  await hook.onStreamSessionInitialized?.(self.app.currentContext, transport, self);
-                }
+        body = JSON.parse(rawBody);
+      } catch (e) {
+        ctx.status = 400;
+        ctx.body = {
+          jsonrpc: '2.0',
+          error: {
+            code: -32000,
+            message: `Bad Request: body should is json, ${e.toString()}`,
+          },
+          id: null,
+        };
+        return;
+      }
+      if (isInitializeRequest(body)) {
+        ctx.respond = false;
+        const eventStore = this.mcpConfig.eventStore;
+        const self = this;
+        const transport = new StreamableHTTPServerTransport({
+          sessionIdGenerator: () => this.mcpConfig.sessionIdGenerator(ctx),
+          eventStore,
+          onsessioninitialized: async () => {
+            if (MCPControllerRegister.hooks.length > 0) {
+              for (const hook of MCPControllerRegister.hooks) {
+                await hook.onStreamSessionInitialized?.(self.app.currentContext, transport, self);
               }
-            },
-          });
+            }
+          },
+        });
 
-          ctx.set({
-            'content-type': 'text/event-stream',
-            'transfer-encoding': 'chunked',
-          });
+        ctx.set({
+          'content-type': 'text/event-stream',
+          'transfer-encoding': 'chunked',
+        });
 
-          await self.mcpServer.connect(transport);
+        await self.mcpServer.connect(transport);
 
-          await ctx.app.ctxStorage.run(ctx, async () => {
-            await mw(ctx, async () => {
-              await transport.handleRequest(ctx.req, ctx.res, body);
-              await awaitEvent(ctx.res, 'close');
-            });
+        await ctx.app.ctxStorage.run(ctx, async () => {
+          await mw(ctx, async () => {
+            await transport.handleRequest(ctx.req, ctx.res, body);
+            await awaitEvent(ctx.res, 'close');
           });
-        } else {
-          ctx.status = 400;
-          ctx.body = {
-            jsonrpc: '2.0',
-            error: {
-              code: -32000,
-              message: 'Bad Request: No valid session ID provided',
-            },
-            id: null,
-          };
-          return;
-        }
+        });
       } else if (sessionId) {
         const transport = self.streamTransports[sessionId];
         if (transport) {
@@ -301,7 +287,7 @@ export class MCPControllerRegister implements ControllerRegister {
           });
           await ctx.app.ctxStorage.run(ctx, async () => {
             await mw(ctx, async () => {
-              await transport.handleRequest(ctx.req, ctx.res);
+              await transport.handleRequest(ctx.req, ctx.res, body);
               await awaitEvent(ctx.res, 'close');
             });
           });
@@ -315,6 +301,17 @@ export class MCPControllerRegister implements ControllerRegister {
             }
           }
         }
+      } else {
+        ctx.status = 400;
+        ctx.body = {
+          jsonrpc: '2.0',
+          error: {
+            code: -32000,
+            message: 'Bad Request: No valid session ID provided',
+          },
+          id: null,
+        };
+        return;
       }
       return;
     };