feat(scheduler): support multi-scheduler tool aggregation and nested call IDs (#17429)

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

@@ -19,6 +19,7 @@ import {
   type ToolCallsUpdateMessage,
   type AnyDeclarativeTool,
   type AnyToolInvocation,
+  ROOT_SCHEDULER_ID,
 } from '@google/gemini-cli-core';
 import { createMockMessageBus } from '@google/gemini-cli-core/src/test-utils/mock-message-bus.js';
 
@@ -73,6 +74,10 @@ describe('useToolExecutionScheduler', () => {
     } as unknown as Config;
   });
 
+  afterEach(() => {
+    vi.clearAllMocks();
+  });
+
   it('initializes with empty tool calls', () => {
     const { result } = renderHook(() =>
       useToolExecutionScheduler(
@@ -112,6 +117,7 @@ describe('useToolExecutionScheduler', () => {
       void mockMessageBus.publish({
         type: MessageBusType.TOOL_CALLS_UPDATE,
         toolCalls: [mockToolCall],
+        schedulerId: ROOT_SCHEDULER_ID,
       } as ToolCallsUpdateMessage);
     });
 
@@ -156,6 +162,7 @@ describe('useToolExecutionScheduler', () => {
       void mockMessageBus.publish({
         type: MessageBusType.TOOL_CALLS_UPDATE,
         toolCalls: [mockToolCall],
+        schedulerId: ROOT_SCHEDULER_ID,
       } as ToolCallsUpdateMessage);
     });
 
@@ -212,6 +219,7 @@ describe('useToolExecutionScheduler', () => {
       void mockMessageBus.publish({
         type: MessageBusType.TOOL_CALLS_UPDATE,
         toolCalls: [mockToolCall],
+        schedulerId: ROOT_SCHEDULER_ID,
       } as ToolCallsUpdateMessage);
     });
 
@@ -274,6 +282,7 @@ describe('useToolExecutionScheduler', () => {
       void mockMessageBus.publish({
         type: MessageBusType.TOOL_CALLS_UPDATE,
         toolCalls: [mockToolCall],
+        schedulerId: ROOT_SCHEDULER_ID,
       } as ToolCallsUpdateMessage);
     });
 
@@ -290,6 +299,7 @@ describe('useToolExecutionScheduler', () => {
       void mockMessageBus.publish({
         type: MessageBusType.TOOL_CALLS_UPDATE,
         toolCalls: [mockToolCall],
+        schedulerId: ROOT_SCHEDULER_ID,
       } as ToolCallsUpdateMessage);
     });
 
@@ -326,6 +336,7 @@ describe('useToolExecutionScheduler', () => {
             invocation: createMockInvocation(),
           },
         ],
+        schedulerId: ROOT_SCHEDULER_ID,
       } as ToolCallsUpdateMessage);
     });
 
@@ -412,4 +423,103 @@ describe('useToolExecutionScheduler', () => {
     expect(completedResult).toEqual([completedToolCall]);
     expect(onComplete).toHaveBeenCalledWith([completedToolCall]);
   });
+
+  it('setToolCallsForDisplay re-groups tools by schedulerId (Multi-Scheduler support)', () => {
+    const { result } = renderHook(() =>
+      useToolExecutionScheduler(
+        vi.fn().mockResolvedValue(undefined),
+        mockConfig,
+        () => undefined,
+      ),
+    );
+
+    const callRoot = {
+      status: 'success' as const,
+      request: {
+        callId: 'call-root',
+        name: 'test',
+        args: {},
+        isClientInitiated: false,
+        prompt_id: 'p1',
+      },
+      tool: createMockTool(),
+      invocation: createMockInvocation(),
+      response: {
+        callId: 'call-root',
+        responseParts: [],
+        resultDisplay: 'OK',
+        error: undefined,
+        errorType: undefined,
+      },
+      schedulerId: ROOT_SCHEDULER_ID,
+    };
+
+    const callSub = {
+      ...callRoot,
+      request: { ...callRoot.request, callId: 'call-sub' },
+      schedulerId: 'subagent-1',
+    };
+
+    // 1. Populate state with multiple schedulers
+    act(() => {
+      void mockMessageBus.publish({
+        type: MessageBusType.TOOL_CALLS_UPDATE,
+        toolCalls: [callRoot],
+        schedulerId: ROOT_SCHEDULER_ID,
+      } as ToolCallsUpdateMessage);
+
+      void mockMessageBus.publish({
+        type: MessageBusType.TOOL_CALLS_UPDATE,
+        toolCalls: [callSub],
+        schedulerId: 'subagent-1',
+      } as ToolCallsUpdateMessage);
+    });
+
+    let [toolCalls] = result.current;
+    expect(toolCalls).toHaveLength(2);
+    expect(
+      toolCalls.find((t) => t.request.callId === 'call-root')?.schedulerId,
+    ).toBe(ROOT_SCHEDULER_ID);
+    expect(
+      toolCalls.find((t) => t.request.callId === 'call-sub')?.schedulerId,
+    ).toBe('subagent-1');
+
+    // 2. Call setToolCallsForDisplay (e.g., simulate a manual update or clear)
+    act(() => {
+      const [, , , setToolCalls] = result.current;
+      setToolCalls((prev) =>
+        prev.map((t) => ({ ...t, responseSubmittedToGemini: true })),
+      );
+    });
+
+    // 3. Verify that tools are still present and maintain their scheduler IDs
+    // The internal map should have been re-grouped.
+    [toolCalls] = result.current;
+    expect(toolCalls).toHaveLength(2);
+    expect(toolCalls.every((t) => t.responseSubmittedToGemini)).toBe(true);
+
+    const updatedRoot = toolCalls.find((t) => t.request.callId === 'call-root');
+    const updatedSub = toolCalls.find((t) => t.request.callId === 'call-sub');
+
+    expect(updatedRoot?.schedulerId).toBe(ROOT_SCHEDULER_ID);
+    expect(updatedSub?.schedulerId).toBe('subagent-1');
+
+    // 4. Verify that a subsequent update to ONE scheduler doesn't wipe the other
+    act(() => {
+      void mockMessageBus.publish({
+        type: MessageBusType.TOOL_CALLS_UPDATE,
+        toolCalls: [{ ...callRoot, status: 'executing' }],
+        schedulerId: ROOT_SCHEDULER_ID,
+      } as ToolCallsUpdateMessage);
+    });
+
+    [toolCalls] = result.current;
+    expect(toolCalls).toHaveLength(2);
+    expect(
+      toolCalls.find((t) => t.request.callId === 'call-root')?.status,
+    ).toBe('executing');
+    expect(
+      toolCalls.find((t) => t.request.callId === 'call-sub')?.schedulerId,
+    ).toBe('subagent-1');
+  });
 });