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');
+  });
 });

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

@@ -16,6 +16,7 @@ import {
   Scheduler,
   type EditorType,
   type ToolCallsUpdateMessage,
+  ROOT_SCHEDULER_ID,
 } from '@google/gemini-cli-core';
 import { useCallback, useState, useMemo, useEffect, useRef } from 'react';
 
@@ -54,8 +55,10 @@ export function useToolExecutionScheduler(
   CancelAllFn,
   number,
 ] {
-  // State stores Core objects, not Display objects
-  const [toolCalls, setToolCalls] = useState<TrackedToolCall[]>([]);
+  // State stores tool calls organized by their originating schedulerId
+  const [toolCallsMap, setToolCallsMap] = useState<
+    Record<string, TrackedToolCall[]>
+  >({});
   const [lastToolOutputTime, setLastToolOutputTime] = useState<number>(0);
 
   const messageBus = useMemo(() => config.getMessageBus(), [config]);
@@ -76,6 +79,7 @@ export function useToolExecutionScheduler(
         config,
         messageBus,
         getPreferredEditor: () => getPreferredEditorRef.current(),
+        schedulerId: ROOT_SCHEDULER_ID,
       }),
     [config, messageBus],
   );
@@ -88,15 +92,21 @@ export function useToolExecutionScheduler(
 
   useEffect(() => {
     const handler = (event: ToolCallsUpdateMessage) => {
-      setToolCalls((prev) => {
-        const adapted = internalAdaptToolCalls(event.toolCalls, prev);
+      // Update output timer for UI spinners (Side Effect)
+      if (event.toolCalls.some((tc) => tc.status === 'executing')) {
+        setLastToolOutputTime(Date.now());
+      }
 
-        // Update output timer for UI spinners
-        if (event.toolCalls.some((tc) => tc.status === 'executing')) {
-          setLastToolOutputTime(Date.now());
-        }
+      setToolCallsMap((prev) => {
+        const adapted = internalAdaptToolCalls(
+          event.toolCalls,
+          prev[event.schedulerId] ?? [],
+        );
 
-        return adapted;
+        return {
+          ...prev,
+          [event.schedulerId]: adapted,
+        };
       });
     };
 
@@ -109,12 +119,14 @@ export function useToolExecutionScheduler(
   const schedule: ScheduleFn = useCallback(
     async (request, signal) => {
       // Clear state for new run
-      setToolCalls([]);
+      setToolCallsMap({});
 
       // 1. Await Core Scheduler directly
       const results = await scheduler.schedule(request, signal);
 
       // 2. Trigger legacy reinjection logic (useGeminiStream loop)
+      // Since this hook instance owns the "root" scheduler, we always trigger
+      // onComplete when it finishes its batch.
       await onCompleteRef.current(results);
 
       return results;
@@ -131,13 +143,52 @@ export function useToolExecutionScheduler(
 
   const markToolsAsSubmitted: MarkToolsAsSubmittedFn = useCallback(
     (callIdsToMark: string[]) => {
-      setToolCalls((prevCalls) =>
-        prevCalls.map((tc) =>
-          callIdsToMark.includes(tc.request.callId)
-            ? { ...tc, responseSubmittedToGemini: true }
-            : tc,
-        ),
-      );
+      setToolCallsMap((prevMap) => {
+        const nextMap = { ...prevMap };
+        for (const [sid, calls] of Object.entries(nextMap)) {
+          nextMap[sid] = calls.map((tc) =>
+            callIdsToMark.includes(tc.request.callId)
+              ? { ...tc, responseSubmittedToGemini: true }
+              : tc,
+          );
+        }
+        return nextMap;
+      });
+    },
+    [],
+  );
+
+  // Flatten the map for the UI components that expect a single list of tools.
+  const toolCalls = useMemo(
+    () => Object.values(toolCallsMap).flat(),
+    [toolCallsMap],
+  );
+
+  // Provide a setter that maintains compatibility with legacy [].
+  const setToolCallsForDisplay = useCallback(
+    (action: React.SetStateAction<TrackedToolCall[]>) => {
+      setToolCallsMap((prev) => {
+        const currentFlattened = Object.values(prev).flat();
+        const nextFlattened =
+          typeof action === 'function' ? action(currentFlattened) : action;
+
+        if (nextFlattened.length === 0) {
+          return {};
+        }
+
+        // Re-group by schedulerId to preserve multi-scheduler state
+        const nextMap: Record<string, TrackedToolCall[]> = {};
+        for (const call of nextFlattened) {
+          // All tool calls should have a schedulerId from the core.
+          // Default to ROOT_SCHEDULER_ID as a safeguard.
+          const sid = call.schedulerId ?? ROOT_SCHEDULER_ID;
+          if (!nextMap[sid]) {
+            nextMap[sid] = [];
+          }
+          nextMap[sid].push(call);
+        }
+        return nextMap;
+      });
     },
     [],
   );
@@ -146,7 +197,7 @@ export function useToolExecutionScheduler(
     toolCalls,
     schedule,
     markToolsAsSubmitted,
-    setToolCalls,
+    setToolCallsForDisplay,
     cancelAll,
     lastToolOutputTime,
   ];