/** * @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 mockConfig: Mocked; let mockToolRegistry: Mocked; let mockMessageBus: Mocked; beforeEach(() => { mockMessageBus = {} as Mocked; mockToolRegistry = { getTool: vi.fn(), } as unknown as Mocked; mockConfig = { getMessageBus: vi.fn().mockReturnValue(mockMessageBus), getToolRegistry: vi.fn().mockReturnValue(mockToolRegistry), } as unknown as Mocked; }); it('should create a scheduler with agent-specific config', async () => { 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, }), ); // Verify that the scheduler's config has the overridden tool registry const schedulerConfig = vi.mocked(Scheduler).mock.calls[0][0].config; expect(schedulerConfig.getToolRegistry()).toBe(mockToolRegistry); }); });