| 1 | import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' |
| 2 | import { |
| 3 | ErrorCode, |
| 4 | GetPromptRequestSchema, |
| 5 | ListPromptsRequestSchema, |
| 6 | McpError, |
| 7 | PromptArgument, |
| 8 | } from '@modelcontextprotocol/sdk/types.js' |
| 9 | import { Inject, Injectable, Scope } from '@nestjs/common' |
| 10 | import { ContextIdFactory, ModuleRef } from '@nestjs/core' |
| 11 | import { getErrorMessage } from '@yikart/common' |
| 12 | import { HttpRequest } from '../../interfaces/http-adapter.interface' |
| 13 | import { McpRegistryService } from '../mcp-registry.service' |
| 14 | import { McpHandlerBase } from './mcp-handler.base' |
| 15 | |
| 16 | @Injectable({ scope: Scope.REQUEST }) |
| 17 | export class McpPromptsHandler extends McpHandlerBase { |
| 18 | constructor( |
| 19 | moduleRef: ModuleRef, |
| 20 | registry: McpRegistryService, |
| 21 | @Inject('MCP_MODULE_ID') private readonly mcpModuleId: string, |
| 22 | ) { |
| 23 | super(moduleRef, registry, McpPromptsHandler.name) |
| 24 | } |
| 25 | |
| 26 | registerHandlers(mcpServer: McpServer, httpRequest: HttpRequest) { |
| 27 | if (this.registry.getPrompts(this.mcpModuleId).length === 0) { |
| 28 | this.logger.debug('No prompts registered, skipping prompt handlers') |
| 29 | return |
| 30 | } |
| 31 | mcpServer.server.setRequestHandler(ListPromptsRequestSchema, () => { |
| 32 | this.logger.debug('ListPromptsRequestSchema is being called') |
| 33 | |
| 34 | const prompts = this.registry |
| 35 | .getPrompts(this.mcpModuleId) |
| 36 | .map(prompt => ({ |
| 37 | name: prompt.metadata.name, |
| 38 | description: prompt.metadata.description, |
| 39 | arguments: prompt.metadata.parameters |
| 40 | ? Object.entries(prompt.metadata.parameters.shape).map( |
| 41 | ([name, field]): PromptArgument => ({ |
| 42 | name, |
| 43 | description: field.description, |
| 44 | required: !field.isOptional(), |
| 45 | }), |
| 46 | ) |
| 47 | : [], |
| 48 | })) |
| 49 | |
| 50 | return { |
| 51 | prompts, |
| 52 | } |
| 53 | }) |
| 54 | |
| 55 | mcpServer.server.setRequestHandler( |
| 56 | GetPromptRequestSchema, |
| 57 | async (request) => { |
| 58 | this.logger.debug('GetPromptRequestSchema is being called') |
| 59 | |
| 60 | try { |
| 61 | const name = request.params.name |
| 62 | const promptInfo = this.registry.findPrompt(this.mcpModuleId, name) |
| 63 | |
| 64 | if (!promptInfo) { |
| 65 | throw new McpError( |
| 66 | ErrorCode.MethodNotFound, |
| 67 | `Unknown prompt: ${name}`, |
| 68 | ) |
| 69 | } |
| 70 | |
| 71 | const contextId = ContextIdFactory.getByRequest(httpRequest) |
| 72 | this.moduleRef.registerRequestByContextId(httpRequest, contextId) |
| 73 | |
| 74 | const promptInstance = await this.moduleRef.resolve( |
| 75 | promptInfo.providerClass, |
| 76 | contextId, |
| 77 | { strict: false }, |
| 78 | ) |
| 79 | |
| 80 | if (!promptInstance) { |
| 81 | throw new McpError( |
| 82 | ErrorCode.MethodNotFound, |
| 83 | `Unknown prompt: ${name}`, |
| 84 | ) |
| 85 | } |
| 86 | |
| 87 | const context = this.createContext(mcpServer, request) |
| 88 | const methodName = promptInfo.methodName |
| 89 | |
| 90 | const result = await promptInstance[methodName]( |
| 91 | request.params.arguments, |
| 92 | context, |
| 93 | httpRequest.raw, |
| 94 | ) |
| 95 | |
| 96 | this.logger.debug(result, 'GetPromptRequestSchema result') |
| 97 | |
| 98 | return result |
| 99 | } |
| 100 | catch (error) { |
| 101 | this.logger.error(error) |
| 102 | const errorMessage = getErrorMessage(error) |
| 103 | return { |
| 104 | contents: [{ mimeType: 'text/plain', text: errorMessage }], |
| 105 | isError: true, |
| 106 | } |
| 107 | } |
| 108 | }, |
| 109 | ) |
| 110 | } |
| 111 | } |
| 112 |