diff --git a/packages/cli/src/services/McpPromptLoader.test.ts b/packages/cli/src/services/McpPromptLoader.test.ts new file mode 100644 index 00000000..d407d739 --- /dev/null +++ b/packages/cli/src/services/McpPromptLoader.test.ts @@ -0,0 +1,128 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { McpPromptLoader } from './McpPromptLoader.js'; +import { Config } from '@google/gemini-cli-core'; +import { PromptArgument } from '@modelcontextprotocol/sdk/types.js'; +import { describe, it, expect } from 'vitest'; + +describe('McpPromptLoader', () => { + const mockConfig = {} as Config; + + describe('parseArgs', () => { + it('should handle multi-word positional arguments', () => { + const loader = new McpPromptLoader(mockConfig); + const promptArgs: PromptArgument[] = [ + { name: 'arg1', required: true }, + { name: 'arg2', required: true }, + ]; + const userArgs = 'hello world'; + const result = loader.parseArgs(userArgs, promptArgs); + expect(result).toEqual({ arg1: 'hello', arg2: 'world' }); + }); + + it('should handle quoted multi-word positional arguments', () => { + const loader = new McpPromptLoader(mockConfig); + const promptArgs: PromptArgument[] = [ + { name: 'arg1', required: true }, + { name: 'arg2', required: true }, + ]; + const userArgs = '"hello world" foo'; + const result = loader.parseArgs(userArgs, promptArgs); + expect(result).toEqual({ arg1: 'hello world', arg2: 'foo' }); + }); + + it('should handle a single positional argument with multiple words', () => { + const loader = new McpPromptLoader(mockConfig); + const promptArgs: PromptArgument[] = [{ name: 'arg1', required: true }]; + const userArgs = 'hello world'; + const result = loader.parseArgs(userArgs, promptArgs); + expect(result).toEqual({ arg1: 'hello world' }); + }); + + it('should handle escaped quotes in positional arguments', () => { + const loader = new McpPromptLoader(mockConfig); + const promptArgs: PromptArgument[] = [{ name: 'arg1', required: true }]; + const userArgs = '"hello \\"world\\""'; + const result = loader.parseArgs(userArgs, promptArgs); + expect(result).toEqual({ arg1: 'hello "world"' }); + }); + + it('should handle escaped backslashes in positional arguments', () => { + const loader = new McpPromptLoader(mockConfig); + const promptArgs: PromptArgument[] = [{ name: 'arg1', required: true }]; + const userArgs = '"hello\\\\world"'; + const result = loader.parseArgs(userArgs, promptArgs); + expect(result).toEqual({ arg1: 'hello\\world' }); + }); + + it('should handle named args followed by positional args', () => { + const loader = new McpPromptLoader(mockConfig); + const promptArgs: PromptArgument[] = [ + { name: 'named', required: true }, + { name: 'pos', required: true }, + ]; + const userArgs = '--named="value" positional'; + const result = loader.parseArgs(userArgs, promptArgs); + expect(result).toEqual({ named: 'value', pos: 'positional' }); + }); + + it('should handle positional args followed by named args', () => { + const loader = new McpPromptLoader(mockConfig); + const promptArgs: PromptArgument[] = [ + { name: 'pos', required: true }, + { name: 'named', required: true }, + ]; + const userArgs = 'positional --named="value"'; + const result = loader.parseArgs(userArgs, promptArgs); + expect(result).toEqual({ pos: 'positional', named: 'value' }); + }); + + it('should handle positional args interspersed with named args', () => { + const loader = new McpPromptLoader(mockConfig); + const promptArgs: PromptArgument[] = [ + { name: 'pos1', required: true }, + { name: 'named', required: true }, + { name: 'pos2', required: true }, + ]; + const userArgs = 'p1 --named="value" p2'; + const result = loader.parseArgs(userArgs, promptArgs); + expect(result).toEqual({ pos1: 'p1', named: 'value', pos2: 'p2' }); + }); + + it('should treat an escaped quote at the start as a literal', () => { + const loader = new McpPromptLoader(mockConfig); + const promptArgs: PromptArgument[] = [ + { name: 'arg1', required: true }, + { name: 'arg2', required: true }, + ]; + const userArgs = '\\"hello world'; + const result = loader.parseArgs(userArgs, promptArgs); + expect(result).toEqual({ arg1: '"hello', arg2: 'world' }); + }); + + it('should handle a complex mix of args', () => { + const loader = new McpPromptLoader(mockConfig); + const promptArgs: PromptArgument[] = [ + { name: 'pos1', required: true }, + { name: 'named1', required: true }, + { name: 'pos2', required: true }, + { name: 'named2', required: true }, + { name: 'pos3', required: true }, + ]; + const userArgs = + 'p1 --named1="value 1" "p2 has spaces" --named2=value2 "p3 \\"with quotes\\""'; + const result = loader.parseArgs(userArgs, promptArgs); + expect(result).toEqual({ + pos1: 'p1', + named1: 'value 1', + pos2: 'p2 has spaces', + named2: 'value2', + pos3: 'p3 "with quotes"', + }); + }); + }); +}); diff --git a/packages/cli/src/services/McpPromptLoader.ts b/packages/cli/src/services/McpPromptLoader.ts index c2862911..1ad37b34 100644 --- a/packages/cli/src/services/McpPromptLoader.ts +++ b/packages/cli/src/services/McpPromptLoader.ts @@ -169,7 +169,16 @@ export class McpPromptLoader implements ICommandLoader { return Promise.resolve(promptCommands); } - private parseArgs( + /** + * Parses the `userArgs` string representing the prompt arguments (all the text + * after the command) into a record matching the shape of the `promptArgs`. + * + * @param userArgs + * @param promptArgs + * @returns A record of the parsed arguments + * @visibleForTesting + */ + parseArgs( userArgs: string, promptArgs: PromptArgument[] | undefined, ): Record | Error { @@ -177,28 +186,36 @@ export class McpPromptLoader implements ICommandLoader { const promptInputs: Record = {}; // arg parsing: --key="value" or --key=value - const namedArgRegex = /--([^=]+)=(?:"((?:\\.|[^"\\])*)"|([^ ]*))/g; + const namedArgRegex = /--([^=]+)=(?:"((?:\\.|[^"\\])*)"|([^ ]+))/g; let match; - const remainingArgs: string[] = []; let lastIndex = 0; + const positionalParts: string[] = []; while ((match = namedArgRegex.exec(userArgs)) !== null) { const key = match[1]; - const value = match[2] ?? match[3]; // Quoted or unquoted value + // Extract the quoted or unquoted argument and remove escape chars. + const value = (match[2] ?? match[3]).replace(/\\(.)/g, '$1'); argValues[key] = value; // Capture text between matches as potential positional args if (match.index > lastIndex) { - remainingArgs.push(userArgs.substring(lastIndex, match.index).trim()); + positionalParts.push(userArgs.substring(lastIndex, match.index)); } lastIndex = namedArgRegex.lastIndex; } // Capture any remaining text after the last named arg if (lastIndex < userArgs.length) { - remainingArgs.push(userArgs.substring(lastIndex).trim()); + positionalParts.push(userArgs.substring(lastIndex)); } - const positionalArgs = remainingArgs.join(' ').split(/ +/); + const positionalArgsString = positionalParts.join('').trim(); + // extracts either quoted strings or non-quoted sequences of non-space characters. + const positionalArgRegex = /(?:"((?:\\.|[^"\\])*)"|([^ ]+))/g; + const positionalArgs: string[] = []; + while ((match = positionalArgRegex.exec(positionalArgsString)) !== null) { + // Extract the quoted or unquoted argument and remove escape chars. + positionalArgs.push((match[1] ?? match[2]).replace(/\\(.)/g, '$1')); + } if (!promptArgs) { return promptInputs; @@ -213,19 +230,27 @@ export class McpPromptLoader implements ICommandLoader { (arg) => arg.required && !promptInputs[arg.name], ); - const missingArgs: string[] = []; - for (let i = 0; i < unfilledArgs.length; i++) { - if (positionalArgs.length > i && positionalArgs[i]) { - promptInputs[unfilledArgs[i].name] = positionalArgs[i]; - } else { - missingArgs.push(unfilledArgs[i].name); + if (unfilledArgs.length === 1) { + // If we have only one unfilled arg, we don't require quotes we just + // join all the given arguments together as if they were quoted. + promptInputs[unfilledArgs[0].name] = positionalArgs.join(' '); + } else { + const missingArgs: string[] = []; + for (let i = 0; i < unfilledArgs.length; i++) { + if (positionalArgs.length > i) { + promptInputs[unfilledArgs[i].name] = positionalArgs[i]; + } else { + missingArgs.push(unfilledArgs[i].name); + } + } + if (missingArgs.length > 0) { + const missingArgNames = missingArgs + .map((name) => `--${name}`) + .join(', '); + return new Error(`Missing required argument(s): ${missingArgNames}`); } } - if (missingArgs.length > 0) { - const missingArgNames = missingArgs.map((name) => `--${name}`).join(', '); - return new Error(`Missing required argument(s): ${missingArgNames}`); - } return promptInputs; } }