Initial support for reloading extensions in the CLI - mcp servers only (#12239)

integration-tests/extensions-reload.test.ts

@@ -0,0 +1,116 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { expect, it, describe } from 'vitest';
+import { TestRig } from './test-helper.js';
+import { TestMcpServer } from './test-mcp-server.js';
+import { writeFileSync } from 'node:fs';
+import { join } from 'node:path';
+import { safeJsonStringify } from '@google/gemini-cli-core/src/utils/safeJsonStringify.js';
+import { env } from 'node:process';
+import { platform } from 'node:os';
+
+const itIf = (condition: boolean) => (condition ? it : it.skip);
+
+describe('extension reloading', () => {
+  const sandboxEnv = env['GEMINI_SANDBOX'];
+
+  // Fails in sandbox mode, can't check for local extension updates.
+  itIf((!sandboxEnv || sandboxEnv === 'false') && platform() !== 'win32')(
+    'installs a local extension, updates it, checks it was reloaded properly',
+    async () => {
+      const serverA = new TestMcpServer();
+      const portA = await serverA.start({
+        hello: () => ({ content: [{ type: 'text', text: 'world' }] }),
+      });
+      const extension = {
+        name: 'test-extension',
+        version: '0.0.1',
+        mcpServers: {
+          'test-server': {
+            httpUrl: `http://localhost:${portA}/mcp`,
+          },
+        },
+      };
+
+      const rig = new TestRig();
+      rig.setup('extension reload test', {
+        settings: {
+          experimental: { extensionReloading: true },
+        },
+      });
+      const testServerPath = join(rig.testDir!, 'gemini-extension.json');
+      writeFileSync(testServerPath, safeJsonStringify(extension, 2));
+      // defensive cleanup from previous tests.
+      try {
+        await rig.runCommand(['extensions', 'uninstall', 'test-extension']);
+      } catch {
+        /* empty */
+      }
+
+      const result = await rig.runCommand(
+        ['extensions', 'install', `${rig.testDir!}`],
+        { stdin: 'y\n' },
+      );
+      expect(result).toContain('test-extension');
+
+      // Now create the update, but its not installed yet
+      const serverB = new TestMcpServer();
+      const portB = await serverB.start({
+        goodbye: () => ({ content: [{ type: 'text', text: 'world' }] }),
+      });
+      extension.version = '0.0.2';
+      extension.mcpServers['test-server'].httpUrl =
+        `http://localhost:${portB}/mcp`;
+      writeFileSync(testServerPath, safeJsonStringify(extension, 2));
+
+      // Start the CLI.
+      const run = await rig.runInteractive('--debug');
+      await run.expectText('You have 1 extension with an update available');
+      // See the outdated extension
+      await run.sendText('/extensions list');
+      await run.type('\r');
+      await run.expectText(
+        'test-extension (v0.0.1) - active (update available)',
+      );
+      await run.sendText('/mcp list');
+      await run.type('\r');
+      await run.expectText(
+        'test-server (from test-extension) - Ready (1 tool)',
+      );
+      await run.expectText('- hello');
+
+      // Update the extension, expect the list to update, and mcp servers as well.
+      await run.sendText('/extensions update test-extension');
+      await run.type('\r');
+      await run.expectText(
+        ` * test-server (remote): http://localhost:${portB}/mcp`,
+      );
+      await run.type('\r'); // consent
+      await run.expectText(
+        'Extension "test-extension" successfully updated: 0.0.1 → 0.0.2',
+      );
+      await new Promise((resolve) => setTimeout(resolve, 1000));
+      await run.sendText('/extensions list');
+      await run.type('\r');
+      await run.expectText('test-extension (v0.0.2) - active (updated)');
+      await run.sendText('/mcp list');
+      await run.type('\r');
+      await run.expectText(
+        'test-server (from test-extension) - Ready (1 tool)',
+      );
+      await run.expectText('- goodbye');
+      await run.sendText('/quit');
+      await run.sendKeys('\r');
+
+      // Clean things up.
+      await serverA.stop();
+      await serverB.stop();
+      await rig.runCommand(['extensions', 'uninstall', 'test-extension']);
+      await rig.cleanup();
+    },
+  );
+});

integration-tests/test-helper.ts

@@ -220,6 +220,13 @@ export class InteractiveRun {
     }
   }
 
+  // Types an entire string at once, necessary for some things like commands
+  // but may run into paste detection issues for larger strings.
+  async sendText(text: string) {
+    this.ptyProcess.write(text);
+    await new Promise((resolve) => setTimeout(resolve, 5));
+  }
+
   // Simulates typing a string one character at a time to avoid paste detection.
   async sendKeys(text: string) {
     const delay = 5;
@@ -311,6 +318,8 @@ export class TestRig {
       model: DEFAULT_GEMINI_MODEL,
       sandbox:
         env['GEMINI_SANDBOX'] !== 'false' ? env['GEMINI_SANDBOX'] : false,
+      // Don't show the IDE connection dialog when running from VsCode
+      ide: { enabled: false, hasSeenNudge: true },
       ...options.settings, // Allow tests to override/add settings
     };
     writeFileSync(

integration-tests/test-mcp-server.ts

@@ -4,17 +4,21 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
+import {
+  McpServer,
+  type ToolCallback,
+} from '@modelcontextprotocol/sdk/server/mcp.js';
 import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
 import express from 'express';
 import { type Server as HTTPServer } from 'node:http';
-
-import { randomUUID } from 'node:crypto';
+import { type ZodRawShape } from 'zod';
 
 export class TestMcpServer {
   private server: HTTPServer | undefined;
 
-  async start(): Promise<number> {
+  async start(
+    tools?: Record<string, ToolCallback<ZodRawShape>>,
+  ): Promise<number> {
     const app = express();
     app.use(express.json());
     const mcpServer = new McpServer(
@@ -22,18 +26,30 @@ export class TestMcpServer {
         name: 'test-mcp-server',
         version: '1.0.0',
       },
-      { capabilities: {} },
+      { capabilities: { tools: {} } },
     );
-
-    const transport = new StreamableHTTPServerTransport({
-      sessionIdGenerator: () => randomUUID(),
-    });
-    mcpServer.connect(transport);
+    if (tools) {
+      for (const [name, cb] of Object.entries(tools)) {
+        mcpServer.registerTool(name, {}, cb);
+      }
+    }
 
     app.post('/mcp', async (req, res) => {
+      const transport = new StreamableHTTPServerTransport({
+        sessionIdGenerator: undefined,
+        enableJsonResponse: true,
+      });
+      res.on('close', () => {
+        transport.close();
+      });
+      await mcpServer.connect(transport);
       await transport.handleRequest(req, res, req.body);
     });
 
+    app.get('/mcp', async (req, res) => {
+      res.status(405).send('Not supported');
+    });
+
     return new Promise((resolve, reject) => {
       this.server = app.listen(0, () => {
         const address = this.server!.address();