feat(mcp): Improve MCP prompt argument parsing (#6779)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

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"',
+      });
+    });
+  });
+});

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<string, unknown> | Error {
@@ -177,28 +186,36 @@ export class McpPromptLoader implements ICommandLoader {
     const promptInputs: Record<string, unknown> = {};
 
     // 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[] = [];
+    if (unfilledArgs.length === 1) {
-    for (let i = 0; i < unfilledArgs.length; i++) {
+      // If we have only one unfilled arg, we don't require quotes we just
-      if (positionalArgs.length > i && positionalArgs[i]) {
+      // join all the given arguments together as if they were quoted.
-        promptInputs[unfilledArgs[i].name] = positionalArgs[i];
+      promptInputs[unfilledArgs[0].name] = positionalArgs.join(' ');
-      } else {
+    } else {
-        missingArgs.push(unfilledArgs[i].name);
+      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;
   }
 }