feat(test-utils): add TestMcpServerBuilder and support in TestRig (#23491)

packages/test-utils/src/index.ts

@@ -7,3 +7,4 @@
 export * from './file-system-test-helpers.js';
 export * from './test-rig.js';
 export * from './mock-utils.js';
+export * from './test-mcp-server.js';

packages/test-utils/src/test-mcp-server-template.mjs

@@ -0,0 +1,69 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { Server } from '@modelcontextprotocol/sdk/server/index.js';
+import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
+import {
+  ListToolsRequestSchema,
+  CallToolRequestSchema,
+} from '@modelcontextprotocol/sdk/types.js';
+import fs from 'fs';
+
+const configPath = process.argv[2];
+if (!configPath) {
+  console.error('Usage: node template.mjs <config-path>');
+  process.exit(1);
+}
+
+const config = JSON.parse(fs.readFileSync(configPath, 'utf-8'));
+
+const server = new Server(
+  {
+    name: config.name,
+    version: config.version || '1.0.0',
+  },
+  {
+    capabilities: {
+      tools: {},
+    },
+  },
+);
+
+// Add tools handler
+server.setRequestHandler(ListToolsRequestSchema, async () => {
+  return {
+    tools: (config.tools || []).map((tool) => ({
+      name: tool.name,
+      description: tool.description,
+      inputSchema: tool.inputSchema || { type: 'object', properties: {} },
+    })),
+  };
+});
+
+// Add call handler
+server.setRequestHandler(CallToolRequestSchema, async (request) => {
+  const toolName = request.params.name;
+  const tool = (config.tools || []).find((t) => t.name === toolName);
+
+  if (!tool) {
+    return {
+      content: [
+        {
+          type: 'text',
+          text: `Error: Tool ${toolName} not found`,
+        },
+      ],
+      isError: true,
+    };
+  }
+
+  return tool.response;
+});
+
+const transport = new StdioServerTransport();
+await server.connect(transport);
+// server.connect resolves when transport connects, but listening continues
+console.error(`Test MCP Server '${config.name}' connected and listening.`);

packages/test-utils/src/test-mcp-server.ts

@@ -0,0 +1,75 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+/**
+ * Response structure for a test tool call.
+ */
+export interface TestToolResponse {
+  content: { type: 'text'; text: string }[];
+  isError?: boolean;
+}
+
+/**
+ * Definition of a test tool.
+ */
+export interface TestTool {
+  name: string;
+  description: string;
+  /** JSON Schema for input arguments */
+  inputSchema?: Record<string, unknown>;
+  response: TestToolResponse;
+}
+
+/**
+ * Configuration structure for the generic test MCP server template.
+ */
+export interface TestMcpConfig {
+  name: string;
+  version?: string;
+  tools: TestTool[];
+}
+
+/**
+ * Builder to easily configure a Test MCP Server in tests.
+ */
+export class TestMcpServerBuilder {
+  private config: TestMcpConfig;
+
+  constructor(name: string) {
+    this.config = { name, tools: [] };
+  }
+
+  /**
+   * Adds a tool to the test server configuration.
+   * @param name Tool name
+   * @param description Tool description
+   * @param response The response to return. Can be a string for simple text responses.
+   * @param inputSchema Optional JSON Schema for validation/documentation
+   */
+  addTool(
+    name: string,
+    description: string,
+    response: TestToolResponse | string,
+    inputSchema?: Record<string, unknown>,
+  ): this {
+    const responseObj =
+      typeof response === 'string'
+        ? { content: [{ type: 'text' as const, text: response }] }
+        : response;
+
+    this.config.tools.push({
+      name,
+      description,
+      inputSchema,
+      response: responseObj,
+    });
+    return this;
+  }
+
+  build(): TestMcpConfig {
+    return this.config;
+  }
+}

packages/test-utils/src/test-rig.ts

@@ -16,6 +16,7 @@ export { GEMINI_DIR };
 import * as pty from '@lydell/node-pty';
 import stripAnsi from 'strip-ansi';
 import * as os from 'node:os';
+import type { TestMcpConfig } from './test-mcp-server.js';
 
 const __dirname = dirname(fileURLToPath(import.meta.url));
 const BUNDLE_PATH = join(__dirname, '..', '..', '..', 'bundle/gemini.js');
@@ -551,7 +552,95 @@ export class TestRig {
     }
     const scriptPath = join(this.testDir, fileName);
     writeFileSync(scriptPath, content);
-    return normalizePath(scriptPath);
+    return normalizePath(scriptPath)!;
+  }
+
+  /**
+   * Adds a test MCP server to the test workspace.
+   * @param name The name of the server
+   * @param config Configuration object or name of predefined config (e.g. 'github')
+   */
+  addTestMcpServer(name: string, config: TestMcpConfig | string) {
+    if (!this.testDir) {
+      throw new Error(
+        'TestRig.setup must be called before adding test servers',
+      );
+    }
+
+    let testConfig: TestMcpConfig;
+    if (typeof config === 'string') {
+      const assetsDir = join(__dirname, '..', 'assets', 'test-servers');
+      const configPath = join(assetsDir, `${config}.json`);
+      if (!fs.existsSync(configPath)) {
+        throw new Error(
+          `Predefined test server config not found: ${configPath}`,
+        );
+      }
+      testConfig = JSON.parse(fs.readFileSync(configPath, 'utf-8'));
+      testConfig.name = name; // Override name
+    } else {
+      testConfig = config;
+    }
+
+    const configFileName = `test-mcp-${name}.json`;
+    const scriptFileName = `test-mcp-${name}.mjs`;
+
+    const configFilePath = join(this.testDir, configFileName);
+    const scriptFilePath = join(this.testDir, scriptFileName);
+
+    // Write config
+    fs.writeFileSync(configFilePath, JSON.stringify(testConfig, null, 2));
+
+    // Copy template script
+    const templatePath = join(__dirname, 'test-mcp-server-template.mjs');
+    if (!fs.existsSync(templatePath)) {
+      throw new Error(`Test template not found at ${templatePath}`);
+    }
+
+    fs.copyFileSync(templatePath, scriptFilePath);
+
+    // Calculate path to monorepo node_modules
+    const monorepoNodeModules = join(
+      __dirname,
+      '..',
+      '..',
+      '..',
+      'node_modules',
+    );
+
+    // Create symlink to node_modules in testDir for ESM resolution
+    const testNodeModules = join(this.testDir, 'node_modules');
+    if (!fs.existsSync(testNodeModules)) {
+      fs.symlinkSync(monorepoNodeModules, testNodeModules, 'dir');
+    }
+
+    // Update settings in workspace and home
+    const updateSettings = (dir: string) => {
+      const settingsPath = join(dir, GEMINI_DIR, 'settings.json');
+      let settings: any = {};
+      if (fs.existsSync(settingsPath)) {
+        settings = JSON.parse(fs.readFileSync(settingsPath, 'utf-8'));
+      } else {
+        fs.mkdirSync(join(dir, GEMINI_DIR), { recursive: true });
+      }
+
+      if (!settings.mcpServers) {
+        settings.mcpServers = {};
+      }
+
+      settings.mcpServers[name] = {
+        command: 'node',
+        args: [scriptFilePath, configFilePath],
+        // Removed env.NODE_PATH as it is ignored in ESM
+      };
+
+      fs.writeFileSync(settingsPath, JSON.stringify(settings, null, 2));
+    };
+
+    updateSettings(this.testDir);
+    if (this.homeDir) {
+      updateSettings(this.homeDir);
+    }
   }
 
   private _getCleanEnv(