gemini-cli/packages/core/src/utils/extensionLoader.test.ts

/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import { describe, expect, it, vi, beforeEach, afterEach } from 'vitest';
import { SimpleExtensionLoader } from './extensionLoader.js';
import type { Config } from '../config/config.js';
import { type McpClientManager } from '../tools/mcp-client-manager.js';

const mockRefreshServerHierarchicalMemory = vi.hoisted(() => vi.fn());

vi.mock('./memoryDiscovery.js', async (importActual) => {
  const actual = await importActual<typeof import('./memoryDiscovery.js')>();
  return {
    ...actual,
    refreshServerHierarchicalMemory: mockRefreshServerHierarchicalMemory,
  };
});

describe('SimpleExtensionLoader', () => {
  let mockConfig: Config;
  let extensionReloadingEnabled: boolean;
  let mockMcpClientManager: McpClientManager;
  const activeExtension = {
    name: 'test-extension',
    isActive: true,
    version: '1.0.0',
    path: '/path/to/extension',
    contextFiles: [],
    id: '123',
  };
  const inactiveExtension = {
    name: 'test-extension',
    isActive: false,
    version: '1.0.0',
    path: '/path/to/extension',
    contextFiles: [],
    id: '123',
  };

  beforeEach(() => {
    mockMcpClientManager = {
      startExtension: vi.fn(),
      stopExtension: vi.fn(),
    } as unknown as McpClientManager;
    extensionReloadingEnabled = false;
    mockConfig = {
      getMcpClientManager: () => mockMcpClientManager,
      getEnableExtensionReloading: () => extensionReloadingEnabled,
    } as unknown as Config;
  });

  afterEach(() => {
    vi.restoreAllMocks();
  });

  it('should start active extensions', async () => {
    const loader = new SimpleExtensionLoader([activeExtension]);
    await loader.start(mockConfig);
    expect(mockMcpClientManager.startExtension).toHaveBeenCalledExactlyOnceWith(
      activeExtension,
    );
  });

  it('should not start inactive extensions', async () => {
    const loader = new SimpleExtensionLoader([inactiveExtension]);
    await loader.start(mockConfig);
    expect(mockMcpClientManager.startExtension).not.toHaveBeenCalled();
  });

  describe('interactive extension loading and unloading', () => {
    it('should not call `start` or `stop` if the loader is not already started', async () => {
      const loader = new SimpleExtensionLoader([]);
      await loader.loadExtension(activeExtension);
      expect(mockMcpClientManager.startExtension).not.toHaveBeenCalled();
      await loader.unloadExtension(activeExtension);
      expect(mockMcpClientManager.stopExtension).not.toHaveBeenCalled();
    });

    it('should start extensions that were explicitly loaded prior to initializing the loader', async () => {
      const loader = new SimpleExtensionLoader([]);
      await loader.loadExtension(activeExtension);
      expect(mockMcpClientManager.startExtension).not.toHaveBeenCalled();
      await loader.start(mockConfig);
      expect(
        mockMcpClientManager.startExtension,
      ).toHaveBeenCalledExactlyOnceWith(activeExtension);
    });

    describe.each([true, false])(
      'when enableExtensionReloading === $i',
      (reloadingEnabled) => {
        beforeEach(() => {
          extensionReloadingEnabled = reloadingEnabled;
        });

        it(`should ${reloadingEnabled ? '' : 'not '}reload extension features`, async () => {
          const loader = new SimpleExtensionLoader([]);
          await loader.start(mockConfig);
          expect(mockMcpClientManager.startExtension).not.toHaveBeenCalled();
          await loader.loadExtension(activeExtension);
          if (reloadingEnabled) {
            expect(
              mockMcpClientManager.startExtension,
            ).toHaveBeenCalledExactlyOnceWith(activeExtension);
            expect(mockRefreshServerHierarchicalMemory).toHaveBeenCalledOnce();
          } else {
            expect(mockMcpClientManager.startExtension).not.toHaveBeenCalled();
            expect(mockRefreshServerHierarchicalMemory).not.toHaveBeenCalled();
          }
          mockRefreshServerHierarchicalMemory.mockClear();

          await loader.unloadExtension(activeExtension);
          if (reloadingEnabled) {
            expect(
              mockMcpClientManager.stopExtension,
            ).toHaveBeenCalledExactlyOnceWith(activeExtension);
            expect(mockRefreshServerHierarchicalMemory).toHaveBeenCalledOnce();
          } else {
            expect(mockMcpClientManager.stopExtension).not.toHaveBeenCalled();
            expect(mockRefreshServerHierarchicalMemory).not.toHaveBeenCalled();
          }
        });

        it.runIf(reloadingEnabled)(
          'Should only reload memory once all extensions are done',
          async () => {
            const anotherExtension = {
              ...activeExtension,
              name: 'another-extension',
            };
            const loader = new SimpleExtensionLoader([]);
            await loader.loadExtension(activeExtension);
            await loader.start(mockConfig);
            expect(mockRefreshServerHierarchicalMemory).not.toHaveBeenCalled();
            await Promise.all([
              loader.unloadExtension(activeExtension),
              loader.loadExtension(anotherExtension),
            ]);
            expect(mockRefreshServerHierarchicalMemory).toHaveBeenCalledOnce();
          },
        );
      },
    );
  });
});