diff --git a/src/index.ts b/src/index.ts index 55fdbb6..5b136f2 100644 --- a/src/index.ts +++ b/src/index.ts @@ -3,3 +3,4 @@ export * from './interfaces'; export * from './mcp.module'; export * from './services/mcp-registry.service'; export * from './services/mcp-executor.service'; +export * from './services/mcp-client.service'; diff --git a/src/interfaces/mcp-options.interface.ts b/src/interfaces/mcp-options.interface.ts index 44a7638..378a7bd 100644 --- a/src/interfaces/mcp-options.interface.ts +++ b/src/interfaces/mcp-options.interface.ts @@ -5,6 +5,7 @@ export enum McpTransportType { SSE = 'sse', STREAMABLE_HTTP = 'streamable-http', STDIO = 'stdio', + IN_MEMORY = 'in-memory', } export interface McpOptions { diff --git a/src/mcp.module.ts b/src/mcp.module.ts index 06dcf9e..22df7f0 100644 --- a/src/mcp.module.ts +++ b/src/mcp.module.ts @@ -3,14 +3,17 @@ import { DiscoveryModule } from '@nestjs/core'; import { McpOptions, McpTransportType } from './interfaces'; import { McpExecutorService } from './services/mcp-executor.service'; import { McpRegistryService } from './services/mcp-registry.service'; +import { McpClientService } from './services/mcp-client.service'; import { SsePingService } from './services/sse-ping.service'; import { createSseController } from './transport/sse.controller.factory'; import { StdioService } from './transport/stdio.service'; import { createStreamableHttpController } from './transport/streamable-http.controller.factory'; +import { InMemoryService } from './transport/in-memory.service'; @Module({ imports: [DiscoveryModule], - providers: [McpRegistryService, McpExecutorService], + providers: [McpRegistryService, McpExecutorService, McpClientService], + exports: [McpClientService], }) export class McpModule { static forRoot(options: McpOptions): DynamicModule { @@ -19,6 +22,7 @@ export class McpModule { McpTransportType.SSE, McpTransportType.STREAMABLE_HTTP, McpTransportType.STDIO, + McpTransportType.IN_MEMORY, ], sseEndpoint: 'sse', messagesEndpoint: 'messages', @@ -38,13 +42,20 @@ export class McpModule { }; const mergedOptions = { ...defaultOptions, ...options } as McpOptions; const providers = this.createProvidersFromOptions(mergedOptions); - const controllers = this.createControllersFromOptions(mergedOptions); + const controllers = this.createControllersFromOptions(mergedOptions); // If IN_MEMORY transport is used, add McpClientService to providers + if ( + Array.isArray(mergedOptions.transport) + ? mergedOptions.transport.includes(McpTransportType.IN_MEMORY) + : mergedOptions.transport === McpTransportType.IN_MEMORY + ) { + providers.push(InMemoryService, McpClientService); + } return { module: McpModule, controllers, providers, - exports: [McpRegistryService], + exports: [McpRegistryService, InMemoryService, McpClientService], }; } @@ -87,9 +98,12 @@ export class McpModule { // STDIO transport is handled by injectable StdioService, no controller } + if (transports.includes(McpTransportType.IN_MEMORY)) { + // IN_MEMORY transport is handled by injectable InMemoryService, no controller + } + return controllers; } - private static createProvidersFromOptions(options: McpOptions): Provider[] { const providers: Provider[] = [ { @@ -112,6 +126,10 @@ export class McpModule { providers.push(StdioService); } + if (transports.includes(McpTransportType.IN_MEMORY)) { + providers.push(InMemoryService, McpClientService); + } + return providers; } } diff --git a/src/services/mcp-client.service.ts b/src/services/mcp-client.service.ts new file mode 100644 index 0000000..39fceb3 --- /dev/null +++ b/src/services/mcp-client.service.ts @@ -0,0 +1,26 @@ +import { Injectable } from '@nestjs/common'; +import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import { InMemoryService } from '../transport/in-memory.service'; + +/** + * Service that provides access to the MCP client. + * This service can be injected in other modules to get access to the MCP client. + */ +@Injectable() +export class McpClientService { + constructor(private readonly inMemoryService: InMemoryService) {} + + /** + * Get the MCP client instance. + * @returns The MCP client instance + * @throws Error if the client is not initialized or transport isn't IN_MEMORY + */ + getClient(): Client { + if (!this.inMemoryService) { + throw new Error( + 'InMemoryService is not available. Make sure the transport includes IN_MEMORY and McpModule is properly imported.', + ); + } + return this.inMemoryService.client; + } +} diff --git a/src/transport/in-memory.service.ts b/src/transport/in-memory.service.ts new file mode 100644 index 0000000..9cd3238 --- /dev/null +++ b/src/transport/in-memory.service.ts @@ -0,0 +1,70 @@ +import { Injectable, Inject, Logger, OnModuleInit } from '@nestjs/common'; +import { ModuleRef, ContextIdFactory } from '@nestjs/core'; +import { McpOptions, McpTransportType } from '../interfaces'; +import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import { InMemoryTransport } from '@modelcontextprotocol/sdk/inMemory.js'; +import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; +import { McpExecutorService } from '../services/mcp-executor.service'; + +@Injectable() +export class InMemoryService implements OnModuleInit { + private readonly logger = new Logger(InMemoryService.name); + + #client: Client | null = null; + + constructor( + @Inject('MCP_OPTIONS') private readonly options: McpOptions, + private readonly moduleRef: ModuleRef, + ) {} + + async onModuleInit() { + if (this.options.transport !== McpTransportType.IN_MEMORY) { + return; + } + this.logger.log('Bootstrapping MCP IN_MEMORY...'); + + const mcpServer = new McpServer( + { name: this.options.name, version: this.options.version }, + { + capabilities: this.options.capabilities || { + tools: {}, + resources: {}, + prompts: {}, + instructions: [], + }, + }, + ); + + const contextId = ContextIdFactory.create(); + const executor = await this.moduleRef.resolve( + McpExecutorService, + contextId, + { strict: false }, + ); + + // eslint-disable-next-line @typescript-eslint/no-unsafe-argument + executor.registerRequestHandlers(mcpServer, {} as any); + + const [serverTransport, clientTransport] = + InMemoryTransport.createLinkedPair(); + + await mcpServer.connect(serverTransport); + + this.#client = new Client({ + name: `${this.options.name} client`, + version: this.options.version, + }); + + await this.#client.connect(clientTransport); + + this.logger.log('MCP IN_MEMORY ready'); + } + + get client(): Client { + if (!this.#client) { + throw new Error('MCP Client is not initialized'); + } + + return this.#client; + } +}