refactor(cli): centralize tool mapping and decouple legacy scheduler (#17044)

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

@@ -0,0 +1,236 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { mapCoreStatusToDisplayStatus, mapToDisplay } from './toolMapping.js';
+import {
+  debugLogger,
+  type AnyDeclarativeTool,
+  type AnyToolInvocation,
+  type ToolCallRequestInfo,
+  type ToolCallResponseInfo,
+  type Status,
+  type ToolCall,
+  type ScheduledToolCall,
+  type SuccessfulToolCall,
+  type ExecutingToolCall,
+  type WaitingToolCall,
+  type CancelledToolCall,
+} from '@google/gemini-cli-core';
+import { ToolCallStatus } from '../types.js';
+
+vi.mock('@google/gemini-cli-core', async (importOriginal) => {
+  const actual =
+    await importOriginal<typeof import('@google/gemini-cli-core')>();
+  return {
+    ...actual,
+    debugLogger: {
+      warn: vi.fn(),
+    },
+  };
+});
+
+describe('toolMapping', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  describe('mapCoreStatusToDisplayStatus', () => {
+    it.each([
+      ['validating', ToolCallStatus.Executing],
+      ['awaiting_approval', ToolCallStatus.Confirming],
+      ['executing', ToolCallStatus.Executing],
+      ['success', ToolCallStatus.Success],
+      ['cancelled', ToolCallStatus.Canceled],
+      ['error', ToolCallStatus.Error],
+      ['scheduled', ToolCallStatus.Pending],
+    ] as const)('maps %s to %s', (coreStatus, expectedDisplayStatus) => {
+      expect(mapCoreStatusToDisplayStatus(coreStatus)).toBe(
+        expectedDisplayStatus,
+      );
+    });
+
+    it('logs warning and defaults to Error for unknown status', () => {
+      const result = mapCoreStatusToDisplayStatus('unknown_status' as Status);
+      expect(result).toBe(ToolCallStatus.Error);
+      expect(debugLogger.warn).toHaveBeenCalledWith(
+        'Unknown core status encountered: unknown_status',
+      );
+    });
+  });
+
+  describe('mapToDisplay', () => {
+    const mockRequest: ToolCallRequestInfo = {
+      callId: 'call-1',
+      name: 'test_tool',
+      args: { arg1: 'val1' },
+      isClientInitiated: false,
+      prompt_id: 'p1',
+    };
+
+    const mockTool = {
+      name: 'test_tool',
+      displayName: 'Test Tool',
+      isOutputMarkdown: true,
+    } as unknown as AnyDeclarativeTool;
+
+    const mockInvocation = {
+      getDescription: () => 'Calling test_tool with args...',
+    } as unknown as AnyToolInvocation;
+
+    const mockResponse: ToolCallResponseInfo = {
+      callId: 'call-1',
+      responseParts: [],
+      resultDisplay: 'Success output',
+      error: undefined,
+      errorType: undefined,
+    };
+
+    it('handles a single tool call input', () => {
+      const toolCall: ScheduledToolCall = {
+        status: 'scheduled',
+        request: mockRequest,
+        tool: mockTool,
+        invocation: mockInvocation,
+      };
+
+      const result = mapToDisplay(toolCall);
+      expect(result.type).toBe('tool_group');
+      expect(result.tools).toHaveLength(1);
+      expect(result.tools[0]?.callId).toBe('call-1');
+    });
+
+    it('handles an array of tool calls', () => {
+      const toolCall1: ScheduledToolCall = {
+        status: 'scheduled',
+        request: mockRequest,
+        tool: mockTool,
+        invocation: mockInvocation,
+      };
+      const toolCall2: ScheduledToolCall = {
+        status: 'scheduled',
+        request: { ...mockRequest, callId: 'call-2' },
+        tool: mockTool,
+        invocation: mockInvocation,
+      };
+
+      const result = mapToDisplay([toolCall1, toolCall2]);
+      expect(result.tools).toHaveLength(2);
+      expect(result.tools[0]?.callId).toBe('call-1');
+      expect(result.tools[1]?.callId).toBe('call-2');
+    });
+
+    it('maps successful tool call properties correctly', () => {
+      const toolCall: SuccessfulToolCall = {
+        status: 'success',
+        request: mockRequest,
+        tool: mockTool,
+        invocation: mockInvocation,
+        response: {
+          ...mockResponse,
+          outputFile: '/tmp/output.txt',
+        },
+      };
+
+      const result = mapToDisplay(toolCall);
+      const displayTool = result.tools[0];
+
+      expect(displayTool).toEqual(
+        expect.objectContaining({
+          callId: 'call-1',
+          name: 'Test Tool',
+          description: 'Calling test_tool with args...',
+          renderOutputAsMarkdown: true,
+          status: ToolCallStatus.Success,
+          resultDisplay: 'Success output',
+          outputFile: '/tmp/output.txt',
+        }),
+      );
+    });
+
+    it('maps executing tool call properties correctly with live output and ptyId', () => {
+      const toolCall: ExecutingToolCall = {
+        status: 'executing',
+        request: mockRequest,
+        tool: mockTool,
+        invocation: mockInvocation,
+        liveOutput: 'Loading...',
+        pid: 12345,
+      };
+
+      const result = mapToDisplay(toolCall);
+      const displayTool = result.tools[0];
+
+      expect(displayTool.status).toBe(ToolCallStatus.Executing);
+      expect(displayTool.resultDisplay).toBe('Loading...');
+      expect(displayTool.ptyId).toBe(12345);
+    });
+
+    it('maps awaiting_approval tool call properties with correlationId', () => {
+      const confirmationDetails = {
+        type: 'exec' as const,
+        title: 'Confirm Exec',
+        command: 'ls',
+        rootCommand: 'ls',
+        rootCommands: ['ls'],
+        onConfirm: vi.fn(),
+      };
+
+      const toolCall: WaitingToolCall = {
+        status: 'awaiting_approval',
+        request: mockRequest,
+        tool: mockTool,
+        invocation: mockInvocation,
+        confirmationDetails,
+        correlationId: 'corr-id-123',
+      };
+
+      const result = mapToDisplay(toolCall);
+      const displayTool = result.tools[0];
+
+      expect(displayTool.status).toBe(ToolCallStatus.Confirming);
+      expect(displayTool.confirmationDetails).toEqual(confirmationDetails);
+    });
+
+    it('maps error tool call missing tool definition', () => {
+      // e.g. "TOOL_NOT_REGISTERED" errors
+      const toolCall: ToolCall = {
+        status: 'error',
+        request: mockRequest, // name: 'test_tool'
+        response: { ...mockResponse, resultDisplay: 'Tool not found' },
+        // notice: no `tool` or `invocation` defined here
+      };
+
+      const result = mapToDisplay(toolCall);
+      const displayTool = result.tools[0];
+
+      expect(displayTool.status).toBe(ToolCallStatus.Error);
+      expect(displayTool.name).toBe('test_tool'); // falls back to request.name
+      expect(displayTool.description).toBe('{"arg1":"val1"}'); // falls back to stringified args
+      expect(displayTool.resultDisplay).toBe('Tool not found');
+      expect(displayTool.renderOutputAsMarkdown).toBe(false);
+    });
+
+    it('maps cancelled tool call properties correctly', () => {
+      const toolCall: CancelledToolCall = {
+        status: 'cancelled',
+        request: mockRequest,
+        tool: mockTool,
+        invocation: mockInvocation,
+        response: {
+          ...mockResponse,
+          resultDisplay: 'User cancelled', // Could be diff output for edits
+        },
+      };
+
+      const result = mapToDisplay(toolCall);
+      const displayTool = result.tools[0];
+
+      expect(displayTool.status).toBe(ToolCallStatus.Canceled);
+      expect(displayTool.resultDisplay).toBe('User cancelled');
+    });
+  });
+});

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

@@ -0,0 +1,133 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+  type ToolCall,
+  type Status as CoreStatus,
+  type ToolCallConfirmationDetails,
+  type ToolResultDisplay,
+  debugLogger,
+} from '@google/gemini-cli-core';
+import {
+  ToolCallStatus,
+  type HistoryItemToolGroup,
+  type IndividualToolCallDisplay,
+} from '../types.js';
+
+export function mapCoreStatusToDisplayStatus(
+  coreStatus: CoreStatus,
+): ToolCallStatus {
+  switch (coreStatus) {
+    case 'validating':
+      return ToolCallStatus.Executing;
+    case 'awaiting_approval':
+      return ToolCallStatus.Confirming;
+    case 'executing':
+      return ToolCallStatus.Executing;
+    case 'success':
+      return ToolCallStatus.Success;
+    case 'cancelled':
+      return ToolCallStatus.Canceled;
+    case 'error':
+      return ToolCallStatus.Error;
+    case 'scheduled':
+      return ToolCallStatus.Pending;
+    default:
+      debugLogger.warn(`Unknown core status encountered: ${coreStatus}`);
+      return ToolCallStatus.Error;
+  }
+}
+
+/**
+ * Transforms `ToolCall` objects into `HistoryItemToolGroup` objects for UI
+ * display. This is a pure projection layer and does not track interaction
+ * state.
+ */
+export function mapToDisplay(
+  toolOrTools: ToolCall[] | ToolCall,
+): HistoryItemToolGroup {
+  const toolCalls = Array.isArray(toolOrTools) ? toolOrTools : [toolOrTools];
+
+  const toolDisplays = toolCalls.map((call): IndividualToolCallDisplay => {
+    let description: string;
+    let renderOutputAsMarkdown = false;
+
+    const displayName = call.tool?.displayName ?? call.request.name;
+
+    if (call.status === 'error') {
+      description = JSON.stringify(call.request.args);
+    } else {
+      description = call.invocation.getDescription();
+      renderOutputAsMarkdown = call.tool.isOutputMarkdown;
+    }
+
+    const baseDisplayProperties = {
+      callId: call.request.callId,
+      name: displayName,
+      description,
+      renderOutputAsMarkdown,
+    };
+
+    let resultDisplay: ToolResultDisplay | undefined = undefined;
+    let confirmationDetails: ToolCallConfirmationDetails | undefined =
+      undefined;
+    let outputFile: string | undefined = undefined;
+    let ptyId: number | undefined = undefined;
+
+    switch (call.status) {
+      case 'success':
+        resultDisplay = call.response.resultDisplay;
+        outputFile = call.response.outputFile;
+        break;
+      case 'error':
+      case 'cancelled':
+        resultDisplay = call.response.resultDisplay;
+        break;
+      case 'awaiting_approval':
+        // Only map if it's the legacy callback-based details.
+        // Serializable details will be handled in a later milestone.
+        if (
+          call.confirmationDetails &&
+          'onConfirm' in call.confirmationDetails &&
+          typeof call.confirmationDetails.onConfirm === 'function'
+        ) {
+          confirmationDetails =
+            call.confirmationDetails as ToolCallConfirmationDetails;
+        }
+        break;
+      case 'executing':
+        resultDisplay = call.liveOutput;
+        ptyId = call.pid;
+        break;
+      case 'scheduled':
+      case 'validating':
+        break;
+      default: {
+        const exhaustiveCheck: never = call;
+        debugLogger.warn(
+          `Unhandled tool call status in mapper: ${
+            (exhaustiveCheck as ToolCall).status
+          }`,
+        );
+        break;
+      }
+    }
+
+    return {
+      ...baseDisplayProperties,
+      status: mapCoreStatusToDisplayStatus(call.status),
+      resultDisplay,
+      confirmationDetails,
+      outputFile,
+      ptyId,
+    };
+  });
+
+  return {
+    type: 'tool_group',
+    tools: toolDisplays,
+  };
+}

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

@@ -63,9 +63,9 @@ import { useStateAndRef } from './useStateAndRef.js';
 import type { UseHistoryManagerReturn } from './useHistoryManager.js';
 import { useLogger } from './useLogger.js';
 import { SHELL_COMMAND_NAME } from '../constants.js';
+import { mapToDisplay as mapTrackedToolCallsToDisplay } from './toolMapping.js';
 import {
   useReactToolScheduler,
-  mapToDisplay as mapTrackedToolCallsToDisplay,
   type TrackedToolCall,
   type TrackedCompletedToolCall,
   type TrackedCancelledToolCall,

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

@@ -7,33 +7,27 @@
 import type {
   Config,
   ToolCallRequestInfo,
-  ExecutingToolCall,
-  ScheduledToolCall,
-  ValidatingToolCall,
-  WaitingToolCall,
-  CompletedToolCall,
-  CancelledToolCall,
   OutputUpdateHandler,
   AllToolCallsCompleteHandler,
   ToolCallsUpdateHandler,
   ToolCall,
-  ToolCallConfirmationDetails,
-  Status as CoreStatus,
   EditorType,
+  CompletedToolCall,
+  ExecutingToolCall,
+  ScheduledToolCall,
+  ValidatingToolCall,
+  WaitingToolCall,
+  CancelledToolCall,
 } from '@google/gemini-cli-core';
-import { CoreToolScheduler, debugLogger } from '@google/gemini-cli-core';
+import { CoreToolScheduler } from '@google/gemini-cli-core';
 import { useCallback, useState, useMemo, useEffect, useRef } from 'react';
-import type {
-  HistoryItemToolGroup,
-  IndividualToolCallDisplay,
-} from '../types.js';
-import { ToolCallStatus } from '../types.js';
 
 export type ScheduleFn = (
   request: ToolCallRequestInfo | ToolCallRequestInfo[],
   signal: AbortSignal,
 ) => Promise<void>;
 export type MarkToolsAsSubmittedFn = (callIds: string[]) => void;
+export type CancelAllFn = (signal: AbortSignal) => void;
 
 export type TrackedScheduledToolCall = ScheduledToolCall & {
   responseSubmittedToGemini?: boolean;
@@ -63,8 +57,12 @@ export type TrackedToolCall =
   | TrackedCompletedToolCall
   | TrackedCancelledToolCall;
 
-export type CancelAllFn = (signal: AbortSignal) => void;
-
+/**
+ * Legacy scheduler implementation based on CoreToolScheduler callbacks.
+ *
+ * This is currently the default implementation used by useGeminiStream.
+ * It will be phased out once the event-driven scheduler migration is complete.
+ */
 export function useReactToolScheduler(
   onComplete: (tools: CompletedToolCall[]) => Promise<void>,
   config: Config,
@@ -82,7 +80,6 @@ export function useReactToolScheduler(
   >([]);
   const [lastToolOutputTime, setLastToolOutputTime] = useState<number>(0);
 
-  // Store callbacks in refs to keep them up-to-date without causing re-renders.
   const onCompleteRef = useRef(onComplete);
   const getPreferredEditorRef = useRef(getPreferredEditor);
 
@@ -131,7 +128,6 @@ export function useReactToolScheduler(
             existingTrackedCall?.responseSubmittedToGemini ?? false;
 
           if (coreTc.status === 'executing') {
-            // Preserve live output if it exists from a previous render.
             const liveOutput = (existingTrackedCall as TrackedExecutingToolCall)
               ?.liveOutput;
             return {
@@ -215,135 +211,3 @@ export function useReactToolScheduler(
     lastToolOutputTime,
   ];
 }
-
-/**
- * Maps a CoreToolScheduler status to the UI's ToolCallStatus enum.
- */
-function mapCoreStatusToDisplayStatus(coreStatus: CoreStatus): ToolCallStatus {
-  switch (coreStatus) {
-    case 'validating':
-      return ToolCallStatus.Executing;
-    case 'awaiting_approval':
-      return ToolCallStatus.Confirming;
-    case 'executing':
-      return ToolCallStatus.Executing;
-    case 'success':
-      return ToolCallStatus.Success;
-    case 'cancelled':
-      return ToolCallStatus.Canceled;
-    case 'error':
-      return ToolCallStatus.Error;
-    case 'scheduled':
-      return ToolCallStatus.Pending;
-    default: {
-      const exhaustiveCheck: never = coreStatus;
-      debugLogger.warn(`Unknown core status encountered: ${exhaustiveCheck}`);
-      return ToolCallStatus.Error;
-    }
-  }
-}
-
-/**
- * Transforms `TrackedToolCall` objects into `HistoryItemToolGroup` objects for UI display.
- */
-export function mapToDisplay(
-  toolOrTools: TrackedToolCall[] | TrackedToolCall,
-): HistoryItemToolGroup {
-  const toolCalls = Array.isArray(toolOrTools) ? toolOrTools : [toolOrTools];
-
-  const toolDisplays = toolCalls.map(
-    (trackedCall): IndividualToolCallDisplay => {
-      let displayName: string;
-      let description: string;
-      let renderOutputAsMarkdown = false;
-
-      if (trackedCall.status === 'error') {
-        displayName =
-          trackedCall.tool === undefined
-            ? trackedCall.request.name
-            : trackedCall.tool.displayName;
-        description = JSON.stringify(trackedCall.request.args);
-      } else {
-        displayName = trackedCall.tool.displayName;
-        description = trackedCall.invocation.getDescription();
-        renderOutputAsMarkdown = trackedCall.tool.isOutputMarkdown;
-      }
-
-      const baseDisplayProperties: Omit<
-        IndividualToolCallDisplay,
-        'status' | 'resultDisplay' | 'confirmationDetails'
-      > = {
-        callId: trackedCall.request.callId,
-        name: displayName,
-        description,
-        renderOutputAsMarkdown,
-      };
-
-      switch (trackedCall.status) {
-        case 'success':
-          return {
-            ...baseDisplayProperties,
-            status: mapCoreStatusToDisplayStatus(trackedCall.status),
-            resultDisplay: trackedCall.response.resultDisplay,
-            confirmationDetails: undefined,
-            outputFile: trackedCall.response.outputFile,
-          };
-        case 'error':
-          return {
-            ...baseDisplayProperties,
-            status: mapCoreStatusToDisplayStatus(trackedCall.status),
-            resultDisplay: trackedCall.response.resultDisplay,
-            confirmationDetails: undefined,
-          };
-        case 'cancelled':
-          return {
-            ...baseDisplayProperties,
-            status: mapCoreStatusToDisplayStatus(trackedCall.status),
-            resultDisplay: trackedCall.response.resultDisplay,
-            confirmationDetails: undefined,
-          };
-        case 'awaiting_approval':
-          return {
-            ...baseDisplayProperties,
-            status: mapCoreStatusToDisplayStatus(trackedCall.status),
-            resultDisplay: undefined,
-            confirmationDetails:
-              trackedCall.confirmationDetails as ToolCallConfirmationDetails,
-          };
-        case 'executing':
-          return {
-            ...baseDisplayProperties,
-            status: mapCoreStatusToDisplayStatus(trackedCall.status),
-            resultDisplay: trackedCall.liveOutput ?? undefined,
-            confirmationDetails: undefined,
-            ptyId: trackedCall.pid,
-          };
-        case 'validating': // Fallthrough
-        case 'scheduled':
-          return {
-            ...baseDisplayProperties,
-            status: mapCoreStatusToDisplayStatus(trackedCall.status),
-            resultDisplay: undefined,
-            confirmationDetails: undefined,
-          };
-        default: {
-          const exhaustiveCheck: never = trackedCall;
-          return {
-            callId: (exhaustiveCheck as TrackedToolCall).request.callId,
-            name: 'Unknown Tool',
-            description: 'Encountered an unknown tool call state.',
-            status: ToolCallStatus.Error,
-            resultDisplay: 'Unknown tool call state',
-            confirmationDetails: undefined,
-            renderOutputAsMarkdown: false,
-          };
-        }
-      }
-    },
-  );
-
-  return {
-    type: 'tool_group',
-    tools: toolDisplays,
-  };
-}

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

@@ -9,10 +9,8 @@ import type { Mock } from 'vitest';
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import { act } from 'react';
 import { renderHook } from '../../test-utils/render.js';
-import {
-  useReactToolScheduler,
-  mapToDisplay,
-} from './useReactToolScheduler.js';
+import { useReactToolScheduler } from './useReactToolScheduler.js';
+import { mapToDisplay } from './toolMapping.js';
 import type { PartUnion, FunctionResponse } from '@google/genai';
 import type {
   Config,