gemini-cli/packages/core/src/agents/agent-scheduler.test.ts

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

import { describe, it, expect, vi, beforeEach, type Mocked } from 'vitest';
import { scheduleAgentTools } from './agent-scheduler.js';
import { Scheduler } from '../scheduler/scheduler.js';
import type { Config } from '../config/config.js';
import type { ToolRegistry } from '../tools/tool-registry.js';
import type { ToolCallRequestInfo } from '../scheduler/types.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';

vi.mock('../scheduler/scheduler.js', () => ({
  Scheduler: vi.fn().mockImplementation(() => ({
    schedule: vi.fn().mockResolvedValue([{ status: 'success' }]),
  })),
}));

describe('agent-scheduler', () => {
  let mockToolRegistry: Mocked<ToolRegistry>;
  let mockMessageBus: Mocked<MessageBus>;

  beforeEach(() => {
    vi.mocked(Scheduler).mockClear();
    mockMessageBus = {} as Mocked<MessageBus>;
    mockToolRegistry = {
      getTool: vi.fn(),
      getMessageBus: vi.fn().mockReturnValue(mockMessageBus),
    } as unknown as Mocked<ToolRegistry>;
  });

  it('should create a scheduler with agent-specific config', async () => {
    const mockConfig = {
      getMessageBus: vi.fn().mockReturnValue(mockMessageBus),
      toolRegistry: mockToolRegistry,
    } as unknown as Mocked<Config>;

    const requests: ToolCallRequestInfo[] = [
      {
        callId: 'call-1',
        name: 'test-tool',
        args: {},
        isClientInitiated: false,
        prompt_id: 'prompt-1',
      },
    ];

    const options = {
      schedulerId: 'subagent-1',
      parentCallId: 'parent-1',
      toolRegistry: mockToolRegistry as unknown as ToolRegistry,
      signal: new AbortController().signal,
    };

    const results = await scheduleAgentTools(
      mockConfig as unknown as Config,
      requests,
      options,
    );

    expect(results).toEqual([{ status: 'success' }]);
    expect(Scheduler).toHaveBeenCalledWith(
      expect.objectContaining({
        schedulerId: 'subagent-1',
        parentCallId: 'parent-1',
        messageBus: mockMessageBus,
      }),
    );

    const schedulerConfig = vi.mocked(Scheduler).mock.calls[0][0].config;
    expect(schedulerConfig.toolRegistry).toBe(mockToolRegistry);
  });

  it('should override toolRegistry getter from prototype chain', async () => {
    const mainRegistry = { _id: 'main' } as unknown as Mocked<ToolRegistry>;
    const agentRegistry = {
      _id: 'agent',
      getMessageBus: vi.fn().mockReturnValue(mockMessageBus),
    } as unknown as Mocked<ToolRegistry>;

    const config = {
      getMessageBus: vi.fn().mockReturnValue(mockMessageBus),
    } as unknown as Mocked<Config>;
    Object.defineProperty(config, 'toolRegistry', {
      get: () => mainRegistry,
      configurable: true,
    });

    await scheduleAgentTools(
      config as unknown as Config,
      [
        {
          callId: 'c1',
          name: 'new_page',
          args: {},
          isClientInitiated: false,
          prompt_id: 'p1',
        },
      ],
      {
        schedulerId: 'browser-1',
        toolRegistry: agentRegistry as unknown as ToolRegistry,
        signal: new AbortController().signal,
      },
    );

    const schedulerConfig = vi.mocked(Scheduler).mock.calls[0][0].config;
    expect(schedulerConfig.toolRegistry).toBe(agentRegistry);
    expect(schedulerConfig.toolRegistry).not.toBe(mainRegistry);
    expect(schedulerConfig.getToolRegistry()).toBe(agentRegistry);
  });
});