basant307/AI_Governance_Project

1/**2 * @license3 * Copyright 2025 Google LLC4 * SPDX-License-Identifier: Apache-2.05 */6 7import type { Config } from '@qwen-code/qwen-code-core';8import {9  getErrorMessage,10  getMCPServerPrompts,11} from '@qwen-code/qwen-code-core';12import type {13  CommandContext,14  SlashCommand,15  SlashCommandActionReturn,16} from '../ui/commands/types.js';17import { CommandKind } from '../ui/commands/types.js';18import type { ICommandLoader } from './types.js';19import type { PromptArgument } from '@modelcontextprotocol/sdk/types.js';20 21/**22 * Discovers and loads executable slash commands from prompts exposed by23 * Model-Context-Protocol (MCP) servers.24 */25export class McpPromptLoader implements ICommandLoader {26  constructor(private readonly config: Config | null) {}27 28  /**29   * Loads all available prompts from all configured MCP servers and adapts30   * them into executable SlashCommand objects.31   *32   * @param _signal An AbortSignal (unused for this synchronous loader).33   * @returns A promise that resolves to an array of loaded SlashCommands.34   */35  loadCommands(_signal: AbortSignal): Promise<SlashCommand[]> {36    const promptCommands: SlashCommand[] = [];37    if (!this.config) {38      return Promise.resolve([]);39    }40    const mcpServers = this.config.getMcpServers() || {};41    for (const serverName in mcpServers) {42      const prompts = getMCPServerPrompts(this.config, serverName) || [];43      for (const prompt of prompts) {44        const commandName = `${prompt.name}`;45        const description =46          prompt.description || `Invoke prompt ${prompt.name}`;47        const newPromptCommand: SlashCommand = {48          name: commandName,49          description,50          modelDescription: description,51          kind: CommandKind.MCP_PROMPT,52          source: 'mcp-prompt' as const,53          sourceLabel: `MCP: ${serverName}`,54          subCommands: [55            {56              name: 'help',57              description: 'Show help for this prompt',58              modelDescription: 'Show help for this prompt',59              kind: CommandKind.MCP_PROMPT,60              source: 'mcp-prompt' as const,61              action: async (): Promise<SlashCommandActionReturn> => {62                if (!prompt.arguments || prompt.arguments.length === 0) {63                  return {64                    type: 'message',65                    messageType: 'info',66                    content: `Prompt "${prompt.name}" has no arguments.`,67                  };68                }69 70                let helpMessage = `Arguments for "${prompt.name}":\n\n`;71                if (prompt.arguments && prompt.arguments.length > 0) {72                  helpMessage += `You can provide arguments by name (e.g., --argName="value") or by position.\n\n`;73                  helpMessage += `e.g., ${prompt.name} ${prompt.arguments?.map((_) => `"foo"`)} is equivalent to ${prompt.name} ${prompt.arguments?.map((arg) => `--${arg.name}="foo"`)}\n\n`;74                }75                for (const arg of prompt.arguments) {76                  helpMessage += `  --${arg.name}\n`;77                  if (arg.description) {78                    helpMessage += `    ${arg.description}\n`;79                  }80                  helpMessage += `    (required: ${81                    arg.required ? 'yes' : 'no'82                  })\n\n`;83                }84                return {85                  type: 'message',86                  messageType: 'info',87                  content: helpMessage,88                };89              },90            },91          ],92          action: async (93            context: CommandContext,94            args: string,95          ): Promise<SlashCommandActionReturn> => {96            if (!this.config) {97              return {98                type: 'message',99                messageType: 'error',100                content: 'Config not loaded.',101              };102            }103 104            const promptInputs = this.parseArgs(args, prompt.arguments);105            if (promptInputs instanceof Error) {106              return {107                type: 'message',108                messageType: 'error',109                content: promptInputs.message,110              };111            }112 113            try {114              const mcpServers = this.config.getMcpServers() || {};115              const mcpServerConfig = mcpServers[serverName];116              if (!mcpServerConfig) {117                return {118                  type: 'message',119                  messageType: 'error',120                  content: `MCP server config not found for '${serverName}'.`,121                };122              }123              const result = await prompt.invoke(promptInputs);124 125              if (result['error']) {126                return {127                  type: 'message',128                  messageType: 'error',129                  content: `Error invoking prompt: ${result['error']}`,130                };131              }132 133              const firstMessage = result.messages?.[0];134              const content = firstMessage?.content;135 136              if (content?.type !== 'text') {137                return {138                  type: 'message',139                  messageType: 'error',140                  content:141                    'Received an empty or invalid prompt response from the server.',142                };143              }144 145              return {146                type: 'submit_prompt',147                content: JSON.stringify(content.text),148              };149            } catch (error) {150              return {151                type: 'message',152                messageType: 'error',153                content: `Error: ${getErrorMessage(error)}`,154              };155            }156          },157          completion: async (158            commandContext: CommandContext,159            partialArg: string,160          ) => {161            const invocation = commandContext.invocation;162            if (!prompt || !prompt.arguments || !invocation) {163              return [];164            }165            const indexOfFirstSpace = invocation.raw.indexOf(' ') + 1;166            let promptInputs =167              indexOfFirstSpace === 0168                ? {}169                : this.parseArgs(170                    invocation.raw.substring(indexOfFirstSpace),171                    prompt.arguments,172                  );173            if (promptInputs instanceof Error) {174              promptInputs = {};175            }176 177            const providedArgNames = Object.keys(promptInputs);178            const unusedArguments =179              prompt.arguments180                .filter((arg) => {181                  // If this arguments is not in the prompt inputs182                  // add it to unusedArguments183                  if (!providedArgNames.includes(arg.name)) {184                    return true;185                  }186 187                  // The parseArgs method assigns the value188                  // at the end of the prompt as a final value189                  // The argument should still be suggested190                  // Example /add --numberOne="34" --num191                  // numberTwo would be assigned a value of --num192                  // numberTwo should still be considered unused193                  const argValue = promptInputs[arg.name];194                  return argValue === partialArg;195                })196                .map((argument) => `--${argument.name}="`) || [];197 198            const exactlyMatchingArgumentAtTheEnd = prompt.arguments199              .map((argument) => `--${argument.name}="`)200              .filter((flagArgument) => {201                const regex = new RegExp(`${flagArgument}[^"]*$`);202                return regex.test(invocation.raw);203              });204 205            if (exactlyMatchingArgumentAtTheEnd.length === 1) {206              if (exactlyMatchingArgumentAtTheEnd[0] === partialArg) {207                return [`${partialArg}"`];208              }209              if (partialArg.endsWith('"')) {210                return [partialArg];211              }212              return [`${partialArg}"`];213            }214 215            const matchingArguments = unusedArguments.filter((flagArgument) =>216              flagArgument.startsWith(partialArg),217            );218 219            return matchingArguments;220          },221        };222        promptCommands.push(newPromptCommand);223      }224    }225    return Promise.resolve(promptCommands);226  }227 228  /**229   * Parses the `userArgs` string representing the prompt arguments (all the text230   * after the command) into a record matching the shape of the `promptArgs`.231   *232   * @param userArgs233   * @param promptArgs234   * @returns A record of the parsed arguments235   * @visibleForTesting236   */237  parseArgs(238    userArgs: string,239    promptArgs: PromptArgument[] | undefined,240  ): Record<string, unknown> | Error {241    const argValues: { [key: string]: string } = {};242    const promptInputs: Record<string, unknown> = {};243 244    // arg parsing: --key="value" or --key=value245    const namedArgRegex = /--([^=]+)=(?:"((?:\\.|[^"\\])*)"|([^ ]+))/g;246    let match;247    let lastIndex = 0;248    const positionalParts: string[] = [];249 250    while ((match = namedArgRegex.exec(userArgs)) !== null) {251      const key = match[1];252      // Extract the quoted or unquoted argument and remove escape chars.253      const value = (match[2] ?? match[3]).replace(/\\(.)/g, '$1');254      argValues[key] = value;255      // Capture text between matches as potential positional args256      if (match.index > lastIndex) {257        positionalParts.push(userArgs.substring(lastIndex, match.index));258      }259      lastIndex = namedArgRegex.lastIndex;260    }261 262    // Capture any remaining text after the last named arg263    if (lastIndex < userArgs.length) {264      positionalParts.push(userArgs.substring(lastIndex));265    }266 267    const positionalArgsString = positionalParts.join('').trim();268    // extracts either quoted strings or non-quoted sequences of non-space characters.269    const positionalArgRegex = /(?:"((?:\\.|[^"\\])*)"|([^ ]+))/g;270    const positionalArgs: string[] = [];271    while ((match = positionalArgRegex.exec(positionalArgsString)) !== null) {272      // Extract the quoted or unquoted argument and remove escape chars.273      positionalArgs.push((match[1] ?? match[2]).replace(/\\(.)/g, '$1'));274    }275 276    if (!promptArgs) {277      return promptInputs;278    }279    for (const arg of promptArgs) {280      if (Object.hasOwn(argValues, arg.name)) {281        promptInputs[arg.name] = argValues[arg.name];282      }283    }284 285    const unfilledArgs = promptArgs.filter(286      (arg) => arg.required && !Object.hasOwn(promptInputs, arg.name),287    );288 289    if (unfilledArgs.length === 1) {290      // If we have only one unfilled arg, we don't require quotes we just291      // join all the given arguments together as if they were quoted.292      promptInputs[unfilledArgs[0].name] = positionalArgs.join(' ');293    } else {294      const missingArgs: string[] = [];295      for (let i = 0; i < unfilledArgs.length; i++) {296        if (positionalArgs.length > i) {297          promptInputs[unfilledArgs[i].name] = positionalArgs[i];298        } else {299          missingArgs.push(unfilledArgs[i].name);300        }301      }302      if (missingArgs.length > 0) {303        const missingArgNames = missingArgs304          .map((name) => `--${name}`)305          .join(', ');306        return new Error(`Missing required argument(s): ${missingArgNames}`);307      }308    }309 310    return promptInputs;311  }312}313