[MCP] Added Support for MCP Prompts (#28408)

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

Author: notyashhh

tools/Mcp/package-lock.json

@@ -19,7 +19,7 @@
   "license": "ISC",
   "description": "",
   "dependencies": {
-    "@modelcontextprotocol/sdk": "^1.9.0",
+    "@modelcontextprotocol/sdk": "^1.17.3",
     "js-yaml": "^4.1.0",
     "zod": "^3.24.2"
   },

tools/Mcp/package.json

@@ -1,8 +1,9 @@
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
 import { z } from "zod";
-import { responseSchema, toolParameterSchema, toolSchema } from "./types.js";
+import { responseSchema, toolParameterSchema, toolSchema, promptSchema } from "./types.js";
 import { ToolsService } from "./services/toolsService.js";
+import { PromptsService } from "./services/promptsService.js";
 import { readFileSync } from "fs";
 import path from "path";
 import { fileURLToPath } from "url";
@@ -59,7 +60,6 @@ export class CodegenServer {
         params: ElicitRequest["params"],
         options?: RequestOptions
     ): Promise<ElicitResult> {
-        //TODO: add log
         return this._mcp.server.elicitInput(params, options);
     }
 
@@ -91,28 +91,33 @@ export class CodegenServer {
     }
 
     initPrompts() {
-        this._mcp.prompt(
-            "create-greeting", 
-            "Generate a customized greeting message", 
-            { name: z.string().describe("Name of the person to greet"), style: z.string().describe("The style of greeting, such a formal, excited, or casual. If not specified casual will be used")}, 
-            ({ name, style = "casual" }: { name: string, style?: string }) => {
-            return {
-                messages: [
-                    {
-                        role: "user",
-                        content: {
-                            type: "text",
-                            text: `Please generate a greeting in ${style} style to ${name}.`,
-                        },
-                    },
-                ],
-            };
-        });
+        const promptsService = PromptsService.getInstance().setServer(this);
+        const promptsSchemas = (specs.prompts || []) as promptSchema[];
+        for (const schema of promptsSchemas) {
+            const parameter = promptsService.createPromptParametersFromSchema(schema.parameters);
+            const callback = promptsService.getPrompts(schema.callbackName, this._responses.get(schema.name));
+            this._mcp.prompt(
+                schema.name,
+                schema.description,
+                parameter,
+                (args: any) => callback(args)
+            );
+        }
     }
 
     initResponses() {
         (responses as responseSchema[])?.forEach((response: responseSchema) => {
-            this._responses.set(response.name, response.text);
+            let text = response.text;
+            if (text.startsWith("@file:")) {
+                const relPath = text.replace("@file:", "");
+                const absPath = path.join(srcPath, "specs", relPath);
+                try {
+                    text = readFileSync(absPath, "utf-8");
+                } catch (e) {
+                    console.error(`Failed to load prompt file ${absPath}:`, e);
+                }
+            }
+            this._responses.set(response.name, text);
         });
     }
 }

tools/Mcp/src/CodegenServer.ts

@@ -0,0 +1,119 @@
+import { z, ZodRawShape } from "zod";
+import { promptSchema, promptParameterSchema } from "../types.js";
+import { CodegenServer } from "../CodegenServer.js";
+
+
+export class PromptsService {
+    private static _instance: PromptsService;
+    private _server: CodegenServer | null = null;
+    private constructor() {}
+
+    static getInstance(): PromptsService {
+        if (!PromptsService._instance) {
+            PromptsService._instance = new PromptsService();
+        }
+        return PromptsService._instance;
+    }
+
+    setServer(server: CodegenServer): PromptsService {
+        this._server = server;
+        return this;
+    }
+
+    getPrompts<Args extends ZodRawShape>(name: string, responseTemplate: string | undefined) {
+        let func;
+        switch (name) {
+            case "createGreetingPrompt":
+                func = this.createGreetingPrompt<Args>;
+                break;
+            case "createPartnerModuleWorkflow":
+                func = this.createPartnerModuleWorkflow<Args>;
+                break;
+            default:
+                throw new Error(`Prompt ${name} not found`);
+        }
+        return this.constructCallback<Args>(func, responseTemplate);
+    }
+
+    constructCallback<Args extends ZodRawShape>(fn: (arr: Args) => Promise<string[]>, responseTemplate: string | undefined) {
+        return async (args: Args) => {
+            const argsArray = await fn(args);
+            const response = this.getResponseString(argsArray, responseTemplate) ?? "";
+            return {
+                messages: [
+                    {
+                        role: "user" as const,
+                        content: {
+                            type: "text" as const,
+                            text: response
+                        }
+                    }
+                ]
+            };
+        };
+    }
+
+    getResponseString(args: string[], responseTemplate: string | undefined): string | undefined {
+        if (!args || args.length === 0) {
+            return responseTemplate;
+        }
+        let response = responseTemplate;
+        for (let i = 0; i < args.length; i++) {
+            response = response?.replaceAll(`{${i}}`, args[i]);
+        }
+        return response;
+    }
+
+    createPromptParametersFromSchema(schemas: promptParameterSchema[]) {
+        const parameter: { [k: string]: any } = {};
+        for (const schema of schemas) {
+            const base = schema.optional ? z.any().optional() : z.any();
+            switch (schema.type) {
+                case "string":
+                    parameter[schema.name] = (schema.optional ? z.string().optional() : z.string()).describe(schema.description);
+                    break;
+                case "number":
+                    parameter[schema.name] = (schema.optional ? z.number().optional() : z.number()).describe(schema.description);
+                    break;
+                case "boolean":
+                    parameter[schema.name] = (schema.optional ? z.boolean().optional() : z.boolean()).describe(schema.description);
+                    break;
+                case "array":
+                    parameter[schema.name] = (schema.optional ? z.array(z.string()).optional() : z.array(z.string())).describe(schema.description);
+                    break;
+                default:
+                    throw new Error(`Unsupported parameter type: ${schema.type}`);
+            }
+        }
+        return parameter;
+    }
+
+    // prompt implementations
+    createGreetingPrompt = async <Args extends ZodRawShape>(args: Args): Promise<string[]> => {
+        const values = Object.values(args);
+        const name = values[0] as unknown as string; // required
+        const style = (values[1] as unknown as string) || "casual"; // optional fallback
+        return [name, style];
+    };
+
+
+    createPartnerModuleWorkflow = async <Args extends ZodRawShape>(args: Args): Promise<string[]> => {
+        const { } = args as any;
+        return [];
+    };
+}
+
+
+// Some Testing Specs:
+
+        //  {
+        //     "name": "partner-module-workflow",
+        //     "description": "Full autonomous workflow instructions to generate a partner Azure PowerShell module via Autorest.",
+        //     "parameters": [
+        //         {"name": "serviceName", "description": "Service name placeholder. This also often corresponds with the Name of the Powershell Module.", "type": "string", "optional": true},
+        //         {"name": "commitId", "description": "Commit id of the swagger from azure-rest-api-specs", "type": "string", "optional": true},
+        //         {"name": "serviceSpecs", "description": "Service specs path under specification. Path of a swagger upto the resource-manager.", "type": "string", "optional": true},
+        //         {"name": "swaggerFileSpecs", "description": "Swagger JSON relative path. Entire path of the swagger down to the openapi file.", "type": "string", "optional": true}
+        //     ],
+        //     "callbackName": "createPartnerModuleWorkflow"
+        // }

tools/Mcp/src/services/promptsService.ts

@@ -1,4 +1,4 @@
-import { CallToolResult, ElicitRequest, ElicitResult } from '@modelcontextprotocol/sdk/types.js';
+import { CallToolResult } from '@modelcontextprotocol/sdk/types.js';
 import { z, ZodRawShape, ZodType, ZodTypeAny } from "zod";
 import * as utils from "./utils.js";
 import path from 'path';