add mcpServer and mcpRequest to context (#45)

* feat: add mcpServer and mcpRequest to context

* adds testing

---------

Co-authored-by: rinormaloku <[email protected]>

Author: lp6moon

src/interfaces/mcp-tool.interface.ts

@@ -1,4 +1,11 @@
-import { Progress } from '@modelcontextprotocol/sdk/types.js';
+import {
+  CallToolRequestSchema,
+  GetPromptRequestSchema,
+  Progress,
+  ReadResourceRequestSchema,
+} from '@modelcontextprotocol/sdk/types.js';
+import { z } from 'zod';
+import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
 
 export type Literal = boolean | null | number | string | undefined;
 
@@ -7,6 +14,13 @@ export type SerializableValue =
   | SerializableValue[]
   | { [key: string]: SerializableValue };
 
+export type McpRequestSchema =
+  | typeof CallToolRequestSchema
+  | typeof ReadResourceRequestSchema
+  | typeof GetPromptRequestSchema;
+
+export type McpRequest = z.infer<McpRequestSchema>;
+
 /**
  * Enhanced execution context that includes user information
  */
@@ -18,4 +32,6 @@ export type Context = {
     info: (message: string, data?: SerializableValue) => void;
     warn: (message: string, data?: SerializableValue) => void;
   };
+  mcpServer: McpServer;
+  mcpRequest: McpRequest;
 };

src/services/handlers/mcp-handler.base.ts

@@ -1,17 +1,8 @@
 import { Logger } from '@nestjs/common';
 import { ModuleRef } from '@nestjs/core';
 import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
-import { z } from 'zod';
-import {
-  CallToolRequestSchema,
-  GetPromptRequestSchema,
-  Progress,
-  ReadResourceRequestSchema,
-} from '@modelcontextprotocol/sdk/types.js';
-import {
-  Context,
-  SerializableValue,
-} from '../../interfaces/mcp-tool.interface';
+import { Progress } from '@modelcontextprotocol/sdk/types.js';
+import { Context, McpRequest, SerializableValue } from '../../interfaces';
 import { McpRegistryService } from '../mcp-registry.service';
 
 export abstract class McpHandlerBase {
@@ -27,15 +18,11 @@ export abstract class McpHandlerBase {
 
   protected createContext(
     mcpServer: McpServer,
-    mcpRequest: z.infer<
-      | typeof CallToolRequestSchema
-      | typeof ReadResourceRequestSchema
-      | typeof GetPromptRequestSchema
-    >,
+    mcpRequest: McpRequest,
   ): Context {
     // handless stateless traffic where notifications and progress are not supported
     if ((mcpServer.server.transport as any).sessionId === undefined) {
-      return this.createStatelessContext();
+      return this.createStatelessContext(mcpServer, mcpRequest);
     }
 
     const progressToken = mcpRequest.params?._meta?.progressToken;
@@ -77,10 +64,15 @@ export abstract class McpHandlerBase {
           });
         },
       },
+      mcpServer,
+      mcpRequest,
     };
   }
 
-  protected createStatelessContext(): Context {
+  protected createStatelessContext(
+    mcpServer: McpServer,
+    mcpRequest: McpRequest,
+  ): Context {
     const warn = (fn: string) => {
       this.logger.warn(`Stateless context: '${fn}' is not supported.`);
     };
@@ -107,6 +99,8 @@ export abstract class McpHandlerBase {
           warn('server report logging not supported in stateless');
         },
       },
+      mcpServer,
+      mcpRequest,
     };
   }
 }

tests/mcp-tool.e2e.spec.ts

@@ -43,6 +43,14 @@ export class GreetingTool {
     }),
   })
   async sayHello({ name }, context: Context) {
+    // Validate that mcpServer and mcpRequest properties exist
+    if (!context.mcpServer) {
+      throw new Error('mcpServer is not defined in the context');
+    }
+    if (!context.mcpRequest) {
+      throw new Error('mcpRequest is not defined in the context');
+    }
+
     const user = await this.userRepository.findByName(name);
     for (let i = 0; i < 5; i++) {
       await new Promise((resolve) => setTimeout(resolve, 50));
@@ -125,12 +133,15 @@ export class ToolRequestScoped {
 
 describe('E2E: MCP ToolServer', () => {
   let app: INestApplication;
-  let testPort: number;
+  let statelessApp: INestApplication;
+  let statefulServerPort: number;
+  let statelessServerPort: number;
 
   // Set timeout for all tests in this describe block to 15000ms
   jest.setTimeout(15000);
 
   beforeAll(async () => {
+    // Create stateful server (original)
     const moduleFixture: TestingModule = await Test.createTestingModule({
       imports: [
         McpModule.forRoot({
@@ -159,21 +170,63 @@ describe('E2E: MCP ToolServer', () => {
     if (!server.address()) {
       throw new Error('Server address not found after listen');
     }
-    testPort = (server.address() as import('net').AddressInfo).port;
+    statefulServerPort = (server.address() as import('net').AddressInfo).port;
+
+    // Create stateless server
+    const statelessModuleFixture: TestingModule =
+      await Test.createTestingModule({
+        imports: [
+          McpModule.forRoot({
+            name: 'test-stateless-mcp-server',
+            version: '0.0.1',
+            guards: [],
+            transport: McpTransportType.STREAMABLE_HTTP,
+            streamableHttp: {
+              enableJsonResponse: true,
+              sessionIdGenerator: undefined,
+              statelessMode: true,
+            },
+          }),
+        ],
+        providers: [
+          GreetingTool,
+          GreetingToolRequestScoped,
+          MockUserRepository,
+          ToolRequestScoped,
+        ],
+      }).compile();
+
+    statelessApp = statelessModuleFixture.createNestApplication();
+    await statelessApp.listen(0);
+
+    const statelessServer = statelessApp.getHttpServer();
+    if (!statelessServer.address()) {
+      throw new Error('Stateless server address not found after listen');
+    }
+    statelessServerPort = (
+      statelessServer.address() as import('net').AddressInfo
+    ).port;
   });
 
   afterAll(async () => {
     await app.close();
+    await statelessApp.close();
   });
 
   const runClientTests = (
     clientType: 'http+sse' | 'streamable http' | 'stdio',
     clientCreator: (port: number, options?: any) => Promise<Client>,
     requestScopedHeaderValue: string,
+    stateless = false,
   ) => {
     describe(`using ${clientType} client (${clientCreator.name})`, () => {
+      let port: number;
+
+      beforeAll(async () => {
+        port = stateless ? statelessServerPort : statefulServerPort;
+      });
       it('should list tools', async () => {
-        const client = await clientCreator(testPort);
+        const client = await clientCreator(port);
         try {
           const tools = await client.listTools();
           expect(tools.tools.length).toBeGreaterThan(0);
@@ -194,7 +247,7 @@ describe('E2E: MCP ToolServer', () => {
       it.each([{ tool: 'hello-world' }, { tool: 'hello-world-scoped' }])(
         'should call the tool $tool and receive results',
         async ({ tool }) => {
-          const client = await clientCreator(testPort);
+          const client = await clientCreator(port);
           try {
             let progressCount = 1;
             const result: any = await client.callTool(
@@ -209,7 +262,7 @@ describe('E2E: MCP ToolServer', () => {
               },
             );
 
-            if (clientType != 'stdio') {
+            if (clientType != 'stdio' && !stateless) {
               // stdio has no support for progress
               expect(progressCount).toBe(5);
             }
@@ -224,7 +277,7 @@ describe('E2E: MCP ToolServer', () => {
       );
 
       it('should call the tool get-request-scoped and receive header', async () => {
-        const client = await clientCreator(testPort, {
+        const client = await clientCreator(port, {
           requestInit: {
             headers: { 'any-header': requestScopedHeaderValue },
           },
@@ -243,7 +296,7 @@ describe('E2E: MCP ToolServer', () => {
       });
 
       it('should reject invalid arguments for hello-world', async () => {
-        const client = await clientCreator(testPort);
+        const client = await clientCreator(port);
 
         try {
           await client.callTool({
@@ -259,7 +312,7 @@ describe('E2E: MCP ToolServer', () => {
       });
 
       it('should reject missing arguments for hello-world', async () => {
-        const client = await clientCreator(testPort);
+        const client = await clientCreator(port);
 
         try {
           await client.callTool({
@@ -275,7 +328,7 @@ describe('E2E: MCP ToolServer', () => {
       });
 
       it('should call the tool and receive an error', async () => {
-        const client = await clientCreator(testPort);
+        const client = await clientCreator(port);
         try {
           const result: any = await client.callTool({
             name: 'hello-world-error',
@@ -297,9 +350,17 @@ describe('E2E: MCP ToolServer', () => {
   // Run tests using the HTTP+SSE MCP client
   runClientTests('http+sse', createSseClient, 'any-value');
 
-  // Run tests using the Streamable HTTP MCP client
+  // Run tests using the [Stateful] Streamable HTTP MCP client
   runClientTests('streamable http', createStreamableClient, 'streamable-value');
 
+  // Run tests using the [Stateless] Streamable HTTP MCP client
+  runClientTests(
+    'streamable http',
+    createStreamableClient,
+    'stateless-streamable-value',
+    true,
+  );
+
   runClientTests(
     'stdio',
     () =>