@nestjs-mcp/server 0.1.0-alpha.9 → 0.1.1

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@nestjs-mcp/server",
-  "version": "0.1.0-alpha.9",
+  "version": "0.1.1",
   "description": "Modular library for building scalable MCP servers with NestJS, providing decorators and integration patterns as a wrapper for the official MCP TypeScript SDK.",
   "author": "Adrián Darío Hidalgo Flores",
   "license": "MIT",
@@ -30,25 +30,13 @@
     "model-context-protocol",
     "module",
     "nestjs",
+    "npm",
+    "pnpm",
     "sdk",
     "server",
-    "typescript"
+    "typescript",
+    "yarn"
   ],
-  "scripts": {
-    "build": "nest build",
-    "prepare": "husky install",
-    "format": "prettier --write \"src/**/*.ts\" \"test/**/*.ts\"",
-    "start:example": "npx -y ts-node-dev --respawn examples/$EXAMPLE/main.ts",
-    "start:inspector": "npx -y @modelcontextprotocol/inspector",
-    "lint": "eslint \"{src,test,examples}/**/*.ts\" --fix",
-    "typecheck": "tsc --noEmit",
-    "test": "jest",
-    "test:watch": "jest --watch",
-    "test:cov": "jest --coverage",
-    "test:debug": "node --inspect-brk -r tsconfig-paths/register -r ts-node/register node_modules/.bin/jest --runInBand",
-    "test:e2e": "jest --config ./test/jest-e2e.json",
-    "npm:publish": "node scripts/npm-publish.js"
-  },
   "peerDependencies": {
     "@nestjs/common": "^11.0.1",
     "@nestjs/core": "^11.0.1",
@@ -105,5 +93,19 @@
     ],
     "coverageDirectory": "../coverage",
     "testEnvironment": "node"
+  },
+  "scripts": {
+    "build": "nest build",
+    "format": "prettier --write \"src/**/*.ts\" \"test/**/*.ts\"",
+    "start:example": "npx -y ts-node-dev --respawn examples/$EXAMPLE/main.ts",
+    "start:inspector": "npx -y @modelcontextprotocol/inspector",
+    "lint": "eslint \"{src,test,examples}/**/*.ts\" --fix",
+    "typecheck": "tsc --noEmit",
+    "test": "jest",
+    "test:watch": "jest --watch",
+    "test:cov": "jest --coverage",
+    "test:debug": "node --inspect-brk -r tsconfig-paths/register -r ts-node/register node_modules/.bin/jest --runInBand",
+    "test:e2e": "jest --config ./test/jest-e2e.json",
+    "npm:publish": "node scripts/npm-publish.js"
   }
-}
+}

package/src/controllers/sse/sse.service.ts

@@ -6,18 +6,18 @@ import { Request, Response } from 'express';
 import { McpServerOptions } from '../../interfaces/mcp-server-options.interface';
 import { McpLoggerService } from '../../registry/logger.service';
 import { RegistryService } from '../../registry/registry.service';
+import { SessionManager } from '../../services/session.manager';
 
 @Injectable()
 export class SseService implements OnModuleInit {
   private server: McpServer;
 
-  private transports = {} as Record<string, SSEServerTransport>;
-
   constructor(
     @Inject('MCP_SERVER_OPTIONS')
     private readonly options: McpServerOptions,
     private readonly registry: RegistryService,
     private readonly logger: McpLoggerService,
+    private readonly sessionManager: SessionManager,
   ) {
     this.server = new McpServer(this.options.serverInfo, this.options.options);
   }
@@ -36,7 +36,11 @@ export class SseService implements OnModuleInit {
   async handleSse(req: Request, res: Response) {
     // Create SSE transport for legacy clients
     const transport = new SSEServerTransport('/messages', res);
-    this.transports[transport.sessionId] = transport;
+
+    this.sessionManager.setSession(transport.sessionId, {
+      transport,
+      request: req,
+    });
 
     this.logger.debug(
       `Starting SSE for sessionId: ${transport.sessionId}`,
@@ -44,7 +48,7 @@ export class SseService implements OnModuleInit {
     );
 
     res.on('close', () => {
-      delete this.transports[transport.sessionId];
+      this.sessionManager.deleteSession(transport.sessionId);
     });
 
     await this.server.connect(transport);
@@ -55,7 +59,12 @@ export class SseService implements OnModuleInit {
    */
   async handleMessage(req: Request, res: Response) {
     const sessionId = req.query.sessionId as string;
-    const transport = this.transports[sessionId];
+    const session = this.sessionManager.getSession(sessionId);
+
+    if (!session) {
+      res.status(400).send('Invalid or missing sessionId');
+      return;
+    }
 
     this.logger.debug(
       `Receiving SSE message for sessionId: ${sessionId}`,
@@ -64,6 +73,13 @@ export class SseService implements OnModuleInit {
 
     this.logger.debug(`SSE message: ${JSON.stringify(req.body)}`, 'MCP_SERVER');
 
+    const transport = session.transport;
+
+    if (!(transport instanceof SSEServerTransport)) {
+      res.status(400).send('Invalid transport');
+      return;
+    }
+
     try {
       if (transport) {
         await transport.handlePostMessage(req, res, req.body);

package/src/controllers/streamable/streamable.service.ts

@@ -11,15 +11,13 @@ import {
 } from '../../interfaces/mcp-server-options.interface';
 import { McpLoggerService } from '../../registry/logger.service';
 import { RegistryService } from '../../registry/registry.service';
-
+import { SessionManager } from '../../services/session.manager';
 // TODO: Stateless mode should be handled here or in another service
 
 @Injectable()
 export class StreamableService implements OnModuleInit {
   private server: McpServer;
 
-  private transports = {} as Record<string, StreamableHTTPServerTransport>;
-
   constructor(
     @Inject('MCP_SERVER_OPTIONS')
     private readonly options: McpServerOptions,
@@ -27,6 +25,7 @@ export class StreamableService implements OnModuleInit {
     private readonly transportOptions: McpModuleTransportOptions,
     private readonly registry: RegistryService,
     private readonly logger: McpLoggerService,
+    private readonly sessionManager: SessionManager,
   ) {
     this.server = new McpServer(this.options.serverInfo, this.options.options);
   }
@@ -47,23 +46,34 @@ export class StreamableService implements OnModuleInit {
    * @param req Express Request object (expects sessionId in query)
    * @param res Express Response object
    */
-  async handlePostRequest(
-    req: Request<any, any, any, { sessionId?: string }>,
-    res: Response,
-  ) {
+  async handlePostRequest(req: Request, res: Response) {
     const sessionId = req.headers['mcp-session-id'] as string | undefined;
     let transport: StreamableHTTPServerTransport;
 
     const { options } = this.transportOptions?.streamable || {};
 
-    if (sessionId && this.transports[sessionId]) {
-      transport = this.transports[sessionId];
+    if (sessionId && this.sessionManager.getSession(sessionId)) {
+      const session = this.sessionManager.getSession(sessionId);
+
+      if (!session) {
+        throw new Error('Session not found');
+      }
+
+      if (!(session.transport instanceof StreamableHTTPServerTransport)) {
+        throw new Error('Invalid transport');
+      }
+
+      transport = session.transport;
     } else if (!sessionId && isInitializeRequest(req.body)) {
+      // This is called only when method is initialize
       transport = new StreamableHTTPServerTransport({
         sessionIdGenerator: () =>
           options?.sessionIdGenerator?.() || randomUUID(),
         onsessioninitialized: (sessionId) => {
-          this.transports[sessionId] = transport;
+          this.sessionManager.setSession(sessionId, {
+            transport,
+            request: req,
+          });
         },
         enableJsonResponse: options?.enableJsonResponse,
         eventStore: options?.eventStore,
@@ -71,9 +81,11 @@ export class StreamableService implements OnModuleInit {
 
       transport.onclose = () => {
         if (transport.sessionId) {
-          delete this.transports[transport.sessionId];
+          this.sessionManager.deleteSession(transport.sessionId);
         }
       };
+
+      await this.server.connect(transport);
     } else {
       // Invalid request
       res.status(400).json({
@@ -99,18 +111,25 @@ export class StreamableService implements OnModuleInit {
    * @param req Express Request object (expects sessionId in query)
    * @param res Express Response object
    */
-  async handleGetRequest(
-    req: Request<any, any, any, { sessionId?: string }>,
-    res: Response,
-  ) {
+  async handleGetRequest(req: Request, res: Response) {
     const sessionId = req.headers['mcp-session-id'] as string | undefined;
 
-    if (!sessionId || !this.transports[sessionId]) {
+    if (!sessionId || !this.sessionManager.getSession(sessionId)) {
       res.status(400).send('Invalid or missing session ID');
       return;
     }
 
-    const transport = this.transports[sessionId];
+    const session = this.sessionManager.getSession(sessionId);
+
+    if (!session) {
+      throw new Error('Session not found');
+    }
+
+    const { transport } = session;
+
+    if (!(transport instanceof StreamableHTTPServerTransport)) {
+      throw new Error('Invalid transport');
+    }
 
     await transport.handleRequest(req, res);
   }
@@ -125,10 +144,7 @@ export class StreamableService implements OnModuleInit {
    * @param req Express Request object
    * @param res Express Response object
    */
-  async handleDeleteRequest(
-    req: Request<any, any, any, { sessionId?: string }>,
-    res: Response,
-  ) {
+  async handleDeleteRequest(req: Request, res: Response) {
     const sessionId = req.headers['mcp-session-id'] as string | undefined;
 
     if (!sessionId) {
@@ -136,7 +152,13 @@ export class StreamableService implements OnModuleInit {
       return;
     }
 
-    const transport = this.transports[sessionId];
+    const session = this.sessionManager.getSession(sessionId);
+
+    if (!session) {
+      throw new Error('Session not found');
+    }
+
+    const { transport } = session;
 
     if (transport) {
       this.logger.debug(
@@ -154,7 +176,7 @@ export class StreamableService implements OnModuleInit {
         return;
       }
 
-      delete this.transports[sessionId];
+      this.sessionManager.deleteSession(sessionId);
 
       res.status(200).json({ success: true, sessionId });
     } else {

package/src/interfaces/capabilities.interface.ts

@@ -2,7 +2,19 @@ import {
   CompleteResourceTemplateCallback,
   ListResourcesCallback,
 } from '@modelcontextprotocol/sdk/server/mcp.js';
-import { ZodOptional, ZodRawShape, ZodType, ZodTypeDef } from 'zod';
+import { RequestHandlerExtra as SdkRequestHandlerExtra } from '@modelcontextprotocol/sdk/shared/protocol.js';
+import {
+  ServerNotification,
+  ServerRequest,
+} from '@modelcontextprotocol/sdk/types.js';
+import {
+  z,
+  ZodOptional,
+  ZodRawShape,
+  ZodType,
+  ZodTypeAny,
+  ZodTypeDef,
+} from 'zod';
 
 export interface ResourceBaseOptions {
   name: string;
@@ -93,3 +105,102 @@ export interface TemplateCallbacks {
     [variable: string]: CompleteResourceTemplateCallback;
   };
 }
+
+export type RequestHandlerExtra = SdkRequestHandlerExtra<
+  ServerRequest,
+  ServerNotification
+>;
+
+export class ResourceUriHandlerParams {
+  public readonly uri: URL;
+  public readonly extra: RequestHandlerExtra;
+
+  private constructor(uri: URL, extra: RequestHandlerExtra) {
+    this.uri = uri;
+    this.extra = extra;
+  }
+
+  static from(uri: URL, extra: RequestHandlerExtra): ResourceUriHandlerParams {
+    return new ResourceUriHandlerParams(uri, extra);
+  }
+}
+
+export class ResourceTemplateHandlerParams {
+  public readonly uri: URL;
+  public readonly variables?: Record<string, string>;
+  public readonly extra: RequestHandlerExtra;
+
+  private constructor(
+    uri: URL,
+    extra: RequestHandlerExtra,
+    variables?: Record<string, string>,
+  ) {
+    this.uri = uri;
+    this.extra = extra;
+    this.variables = variables;
+  }
+
+  static from(
+    uri: URL,
+    extra: RequestHandlerExtra,
+    variables?: Record<string, string>,
+  ): ResourceTemplateHandlerParams {
+    return new ResourceTemplateHandlerParams(uri, extra, variables);
+  }
+}
+
+export class PromptHandlerParams<
+  Args extends PromptArgsRawShape | undefined = undefined,
+> {
+  public readonly args?: Args extends PromptArgsRawShape
+    ? z.objectOutputType<Args, ZodTypeAny>
+    : undefined;
+  public readonly extra: RequestHandlerExtra;
+
+  private constructor(
+    extra: RequestHandlerExtra,
+    args?: Args extends PromptArgsRawShape
+      ? z.objectOutputType<Args, ZodTypeAny>
+      : undefined,
+  ) {
+    this.extra = extra;
+    this.args = args;
+  }
+
+  static from<Args extends PromptArgsRawShape | undefined>(
+    extra: RequestHandlerExtra,
+    args?: Args extends PromptArgsRawShape
+      ? z.objectOutputType<Args, ZodTypeAny>
+      : undefined,
+  ): PromptHandlerParams<Args> {
+    return new PromptHandlerParams<Args>(extra, args);
+  }
+}
+
+export class ToolHandlerParams<
+  Args extends ZodRawShape | undefined = undefined,
+> {
+  public readonly args?: Args extends ZodRawShape
+    ? z.objectOutputType<Args, ZodTypeAny>
+    : undefined;
+  public readonly extra: RequestHandlerExtra;
+
+  private constructor(
+    extra: RequestHandlerExtra,
+    args?: Args extends ZodRawShape
+      ? z.objectOutputType<Args, ZodTypeAny>
+      : undefined,
+  ) {
+    this.extra = extra;
+    this.args = args;
+  }
+
+  static from<Args extends ZodRawShape | undefined>(
+    extra: RequestHandlerExtra,
+    args?: Args extends ZodRawShape
+      ? z.objectOutputType<Args, ZodTypeAny>
+      : undefined,
+  ): ToolHandlerParams<Args> {
+    return new ToolHandlerParams<Args>(extra, args);
+  }
+}

package/src/interfaces/context.interface.ts

@@ -1,6 +1,11 @@
 import { ExecutionContext } from '@nestjs/common';
 
-import { McpMessage } from './message.types';
+import {
+  PromptHandlerParams,
+  ResourceTemplateHandlerParams,
+  ResourceUriHandlerParams,
+  ToolHandlerParams,
+} from './capabilities.interface';
 /**
  * Custom execution context for MCP guards.
  * Extends NestJS ExecutionContext and adds args for MCP method arguments.
@@ -8,11 +13,21 @@ import { McpMessage } from './message.types';
  * @property args - The arguments passed to the MCP method
  * @property message - The current message from the request
  */
-// TODO: Type Args correctly
+// TODO: Remove extends ExecutionContext we don't need it
 export interface McpExecutionContext extends ExecutionContext {
-  /** The arguments passed to the MCP method */
-  args: unknown[];
+  // TODO: Remove this once the getArgs method implementation is complete.
+  args:
+    | ResourceUriHandlerParams
+    | ResourceTemplateHandlerParams
+    | PromptHandlerParams
+    | ToolHandlerParams;
 
-  /** The current message from the request */
-  message: McpMessage | undefined;
+  // TODO: Uncomment this once the getArgs type is fixed
+  // getArgs: () =>
+  //   | ResourceUriHandlerParams
+  //   | ResourceTemplateHandlerParams
+  //   | PromptHandlerParams
+  //   | ToolHandlerParams;
+
+  getSessionId: () => string;
 }

package/src/mcp.module.ts

@@ -1,6 +1,6 @@
 import { Implementation } from '@modelcontextprotocol/sdk/types.js';
 import { DynamicModule, Module, Provider, Type } from '@nestjs/common';
-import { APP_INTERCEPTOR, DiscoveryModule } from '@nestjs/core';
+import { DiscoveryModule } from '@nestjs/core';
 import { AsyncLocalStorage } from 'async_hooks';
 
 import { SseController, SseService } from './controllers/sse';
@@ -8,7 +8,6 @@ import {
   StreamableController,
   StreamableService,
 } from './controllers/streamable';
-import { RequestContextInterceptor } from './interceptors/message.interceptor';
 import {
   McpFeatureOptions,
   McpLoggingOptions,
@@ -19,8 +18,7 @@ import {
 import { DiscoveryService } from './registry/discovery.service';
 import { McpLoggerService } from './registry/logger.service';
 import { RegistryService } from './registry/registry.service';
-import { MessageService } from './services/message.service';
-
+import { SessionManager } from './services/session.manager';
 @Module({
   imports: [DiscoveryModule],
   providers: [
@@ -31,13 +29,9 @@ import { MessageService } from './services/message.service';
       useValue: new AsyncLocalStorage(),
     },
     McpLoggerService,
-    MessageService,
-    {
-      provide: APP_INTERCEPTOR,
-      useClass: RequestContextInterceptor,
-    },
+    SessionManager,
   ],
-  exports: [MessageService],
+  exports: [SessionManager],
 })
 export class McpModule {
   /**

package/src/registry/registry.service.ts

@@ -4,6 +4,7 @@ import {
 } from '@modelcontextprotocol/sdk/server/mcp.js';
 import type { CanActivate, Type } from '@nestjs/common';
 import { Injectable } from '@nestjs/common';
+import { Reflector } from '@nestjs/core';
 
 import { MCP_RESOLVER } from '../decorators';
 import {
@@ -13,20 +14,27 @@ import {
 } from '../decorators/capabilities.constants';
 import { MCP_GUARDS } from '../decorators/capabilities.decorators';
 import {
+  PromptHandlerParams,
   PromptOptions,
+  RequestHandlerExtra,
   ResourceOptions,
+  ResourceTemplateHandlerParams,
+  ResourceUriHandlerParams,
+  ToolHandlerParams,
   ToolOptions,
 } from '../interfaces/capabilities.interface';
 import { McpExecutionContext } from '../interfaces/context.interface';
-import { MessageService } from '../services/message.service';
+import { SessionManager } from '../services/session.manager';
 import { DiscoveryService } from './discovery.service';
 import { McpLoggerService } from './logger.service';
+
 @Injectable()
 export class RegistryService {
   constructor(
     private readonly discoveryService: DiscoveryService,
     private readonly logger: McpLoggerService,
-    private readonly messageService: MessageService,
+    private readonly reflector: Reflector,
+    private readonly sessionManager: SessionManager,
   ) {}
 
   async registerAll(server: McpServer): Promise<void> {
@@ -42,6 +50,57 @@ export class RegistryService {
     ]);
   }
 
+  private getDecoratorType(method: Type<any> | undefined): string | null {
+    if (!method) return null;
+
+    if (this.reflector.get(MCP_TOOL, method)) return 'TOOL';
+    if (this.reflector.get(MCP_PROMPT, method)) return 'PROMPT';
+    if (this.reflector.get(MCP_RESOURCE, method)) return 'RESOURCE';
+
+    return null;
+  }
+
+  private getHandlerArgs(
+    method: Type<any> | undefined,
+    args: unknown[],
+  ):
+    | ResourceUriHandlerParams
+    | ResourceTemplateHandlerParams
+    | PromptHandlerParams
+    | ToolHandlerParams {
+    if (!method) throw new Error('Method not found');
+
+    switch (this.getDecoratorType(method)) {
+      case 'RESOURCE':
+        return args[0] instanceof URL
+          ? ResourceUriHandlerParams.from(
+              args[0],
+              args[1] as RequestHandlerExtra,
+            )
+          : ResourceTemplateHandlerParams.from(
+              args[0] as any,
+              args[2] as RequestHandlerExtra,
+              args[1] as Record<string, string>,
+            );
+      case 'PROMPT':
+        return args.length === 1
+          ? PromptHandlerParams.from(args[0] as RequestHandlerExtra)
+          : PromptHandlerParams.from(
+              args[1] as RequestHandlerExtra,
+              args[0] as any,
+            );
+      case 'TOOL':
+        return args.length === 1
+          ? ToolHandlerParams.from(args[0] as RequestHandlerExtra)
+          : ToolHandlerParams.from(
+              args[1] as RequestHandlerExtra,
+              args[0] as any,
+            );
+      default:
+        throw new Error(`Unknown decorator type for method ${method.name}`);
+    }
+  }
+
   /**
    * Executes all guards attached to the resolver class and method.
    * Throws an error if any guard denies access.
@@ -58,6 +117,7 @@ export class RegistryService {
   ): Promise<void> {
     // Retrieve class-level guards
     const classConstructor = instance.constructor;
+
     const classGuards: (CanActivate | { new (): CanActivate })[] =
       (Reflect.getMetadata(MCP_GUARDS, classConstructor) as (
         | CanActivate
@@ -69,7 +129,8 @@ export class RegistryService {
       string,
       unknown
     >;
-    const methodKey = prototype[methodName] as object | undefined;
+
+    const methodKey = prototype[methodName] as Type<any> | undefined;
 
     const methodGuards: (CanActivate | { new (): CanActivate })[] =
       (methodKey &&
@@ -83,16 +144,39 @@ export class RegistryService {
     const allGuards = [...classGuards, ...methodGuards];
 
     if (!allGuards.length) return Promise.resolve();
-    // Build a minimal context (customize as needed)
-    const context: McpExecutionContext = {
-      args,
-      message: this.messageService.get(),
 
-      // @ts-expect-error: Default types are 'http' | 'ws' | 'rpc' but in our case
-      // we are using 'mcp'
+    const handlerArgs = this.getHandlerArgs(methodKey, args);
+
+    const { sessionId } = handlerArgs.extra;
+
+    if (!sessionId) {
+      throw new Error('Session not found');
+    }
+
+    const session = this.sessionManager.getSession(sessionId);
+
+    if (!session) {
+      throw new Error('Session not found');
+    }
+
+    const context: McpExecutionContext = {
+      args: handlerArgs,
+      // @ts-expect-error: Default types are 'http' | 'ws' | 'rpc' but in our case, we are using 'mcp'
       getType: () => 'mcp',
       getClass: () => instance.constructor as Type<any>,
-      getArgs: <T extends Array<any>>() => args as T,
+      getArgs: <T = any>() => args as T,
+      getArgByIndex: <T = any>(index: number) => args[index] as T,
+      getSessionId: () => sessionId,
+      getHandler: () => methodKey as unknown as Type<any>,
+      switchToHttp: () => ({
+        getRequest: <R = Request>() => session.request as R,
+        getResponse: () => {
+          throw new Error('Response not available in MCP context');
+        },
+        getNext: () => {
+          throw new Error('Next not available in MCP context');
+        },
+      }),
     };
 
     return (async () => {

package/src/services/session.manager.ts

@@ -0,0 +1,26 @@
+import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse';
+import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
+import { Injectable } from '@nestjs/common';
+import { Request } from 'express';
+
+type Session = {
+  transport: StreamableHTTPServerTransport | SSEServerTransport;
+  request: Request;
+};
+
+@Injectable()
+export class SessionManager {
+  private sessions = new Map<string, Session>();
+
+  public getSession(sessionId: string): Session | undefined {
+    return this.sessions.get(sessionId);
+  }
+
+  public setSession(sessionId: string, session: Session): void {
+    this.sessions.set(sessionId, session);
+  }
+
+  public deleteSession(sessionId: string): void {
+    this.sessions.delete(sessionId);
+  }
+}

package/dist/interceptors/message.interceptor.d.ts

@@ -1,10 +0,0 @@
-import { CallHandler, ExecutionContext, NestInterceptor } from '@nestjs/common';
-import { Observable } from 'rxjs';
-import { McpLoggerService } from '../registry/logger.service';
-import { MessageService } from '../services/message.service';
-export declare class RequestContextInterceptor implements NestInterceptor {
-    private readonly logger;
-    private readonly messageService;
-    constructor(logger: McpLoggerService, messageService: MessageService);
-    intercept(context: ExecutionContext, next: CallHandler): Observable<any>;
-}