Initial support for reloading extensions in the CLI - mcp servers only

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,36 +4,52 @@
  * 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(
       {
         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();

packages/a2a-server/src/utils/testing_utils.ts

@@ -53,6 +53,7 @@ export function createMockConfig(
     getEnableMessageBusIntegration: vi.fn().mockReturnValue(false),
     getMessageBus: vi.fn(),
     getPolicyEngine: vi.fn(),
+    getEnableExtensionReloading: vi.fn().mockReturnValue(false),
     ...overrides,
   } as unknown as Config;
 

packages/cli/src/commands/extensions/update.ts

@@ -30,11 +30,12 @@ const updateOutput = (info: ExtensionUpdateInfo) =>
 
 export async function handleUpdate(args: UpdateArgs) {
   const workspaceDir = process.cwd();
+  const settings = loadSettings(workspaceDir).merged;
   const extensionManager = new ExtensionManager({
     workspaceDir,
     requestConsent: requestConsentNonInteractive,
     requestSetting: promptForSetting,
-    settings: loadSettings(workspaceDir).merged,
+    settings,
   });
 
   const extensions = await extensionManager.loadExtensions();
@@ -67,6 +68,7 @@ export async function handleUpdate(args: UpdateArgs) {
         extensionManager,
         updateState,
         () => {},
+        settings.experimental?.extensionReloading,
       ))!;
       if (
         updatedExtensionInfo.originalVersion !==

packages/cli/src/config/config.ts

@@ -680,6 +680,7 @@ export async function loadCliConfig(
     listExtensions: argv.listExtensions || false,
     enabledExtensions: argv.extensions,
     extensionLoader: extensionManager,
+    enableExtensionReloading: settings.experimental?.extensionReloading,
     blockedMcpServers,
     noBrowser: !!process.env['NO_BROWSER'],
     summarizeToolOutput: settings.model?.summarizeToolOutput,

packages/cli/src/config/extensions/update.ts

@@ -28,6 +28,7 @@ export async function updateExtension(
   extensionManager: ExtensionManager,
   currentState: ExtensionUpdateState,
   dispatchExtensionStateUpdate: (action: ExtensionUpdateAction) => void,
+  enableExtensionReloading?: boolean,
 ): Promise<ExtensionUpdateInfo | undefined> {
   if (currentState === ExtensionUpdateState.UPDATING) {
     return undefined;
@@ -81,7 +82,9 @@ export async function updateExtension(
       type: 'SET_STATE',
       payload: {
         name: extension.name,
-        state: ExtensionUpdateState.UPDATED_NEEDS_RESTART,
+        state: enableExtensionReloading
+          ? ExtensionUpdateState.UPDATED
+          : ExtensionUpdateState.UPDATED_NEEDS_RESTART,
       },
     });
     return {
@@ -109,6 +112,7 @@ export async function updateAllUpdatableExtensions(
   extensionsState: Map<string, ExtensionUpdateStatus>,
   extensionManager: ExtensionManager,
   dispatch: (action: ExtensionUpdateAction) => void,
+  enableExtensionReloading?: boolean,
 ): Promise<ExtensionUpdateInfo[]> {
   return (
     await Promise.all(
@@ -124,6 +128,7 @@ export async function updateAllUpdatableExtensions(
             extensionManager,
             extensionsState.get(extension.name)!.status,
             dispatch,
+            enableExtensionReloading,
           ),
         ),
     )
@@ -141,34 +146,37 @@ export async function checkForAllExtensionUpdates(
   dispatch: (action: ExtensionUpdateAction) => void,
 ): Promise<void> {
   dispatch({ type: 'BATCH_CHECK_START' });
-  const promises: Array<Promise<void>> = [];
-  for (const extension of extensions) {
-    if (!extension.installMetadata) {
+  try {
+    const promises: Array<Promise<void>> = [];
+    for (const extension of extensions) {
+      if (!extension.installMetadata) {
+        dispatch({
+          type: 'SET_STATE',
+          payload: {
+            name: extension.name,
+            state: ExtensionUpdateState.NOT_UPDATABLE,
+          },
+        });
+        continue;
+      }
       dispatch({
         type: 'SET_STATE',
         payload: {
           name: extension.name,
-          state: ExtensionUpdateState.NOT_UPDATABLE,
+          state: ExtensionUpdateState.CHECKING_FOR_UPDATES,
         },
       });
-      continue;
+      promises.push(
+        checkForExtensionUpdate(extension, extensionManager).then((state) =>
+          dispatch({
+            type: 'SET_STATE',
+            payload: { name: extension.name, state },
+          }),
+        ),
+      );
     }
-    dispatch({
-      type: 'SET_STATE',
-      payload: {
-        name: extension.name,
-        state: ExtensionUpdateState.CHECKING_FOR_UPDATES,
-      },
-    });
-    promises.push(
-      checkForExtensionUpdate(extension, extensionManager).then((state) =>
-        dispatch({
-          type: 'SET_STATE',
-          payload: { name: extension.name, state },
-        }),
-      ),
-    );
+    await Promise.all(promises);
+  } finally {
+    dispatch({ type: 'BATCH_CHECK_END' });
   }
-  await Promise.all(promises);
-  dispatch({ type: 'BATCH_CHECK_END' });
 }

packages/cli/src/config/settingsSchema.ts

@@ -1075,6 +1075,16 @@ const SETTINGS_SCHEMA = {
         description: 'Enable extension management features.',
         showInDialog: false,
       },
+      extensionReloading: {
+        type: 'boolean',
+        label: 'Extension Reloading',
+        category: 'Experimental',
+        requiresRestart: true,
+        default: false,
+        description:
+          'Enables extension loading/unloading within the CLI session.',
+        showInDialog: false,
+      },
       useModelRouter: {
         type: 'boolean',
         label: 'Use Model Router',

packages/cli/src/ui/AppContainer.tsx

@@ -182,7 +182,11 @@ export const AppContainer = (props: AppContainerProps) => {
     extensionsUpdateState,
     extensionsUpdateStateInternal,
     dispatchExtensionStateUpdate,
-  } = useExtensionUpdates(extensionManager, historyManager.addItem);
+  } = useExtensionUpdates(
+    extensionManager,
+    historyManager.addItem,
+    config.getEnableExtensionReloading(),
+  );
 
   const [isPermissionsDialogOpen, setPermissionsDialogOpen] = useState(false);
   const openPermissionsDialog = useCallback(

packages/cli/src/ui/components/views/ExtensionsList.test.tsx

@@ -97,6 +97,10 @@ describe('<ExtensionsList />', () => {
       state: ExtensionUpdateState.UPDATED_NEEDS_RESTART,
       expectedText: '(updated, needs restart)',
     },
+    {
+      state: ExtensionUpdateState.UPDATED,
+      expectedText: '(updated)',
+    },
     {
       state: ExtensionUpdateState.ERROR,
       expectedText: '(error)',

packages/cli/src/ui/components/views/ExtensionsList.tsx

@@ -48,6 +48,7 @@ export const ExtensionsList: React.FC<ExtensionsList> = ({ extensions }) => {
               break;
             case ExtensionUpdateState.UP_TO_DATE:
             case ExtensionUpdateState.NOT_UPDATABLE:
+            case ExtensionUpdateState.UPDATED:
               stateColor = 'green';
               break;
             case undefined:

packages/cli/src/ui/hooks/atCommandProcessor.test.ts

@@ -84,6 +84,7 @@ describe('handleAtCommand', () => {
         getReadManyFilesExcludes: () => [],
       }),
       getUsageStatisticsEnabled: () => false,
+      getEnableExtensionReloading: () => false,
     } as unknown as Config;
 
     const registry = new ToolRegistry(mockConfig);