refactor: Centralize tool scheduling logic and simplify React hook (#670)

1 files changed, 0 insertions, 626 deletions

diff --git a/packages/cli/src/ui/hooks/useToolScheduler.ts b/packages/cli/src/ui/hooks/useToolScheduler.ts
deleted file mode 100644
index 9233ebcf..00000000
--- a/packages/cli/src/ui/hooks/useToolScheduler.ts
+++ /dev/null
@@ -1,626 +0,0 @@
-/**
- * @license
- * Copyright 2025 Google LLC
- * SPDX-License-Identifier: Apache-2.0
- */
-
-import {
-  Config,
-  ToolCallRequestInfo,
-  ToolCallResponseInfo,
-  ToolConfirmationOutcome,
-  Tool,
-  ToolCallConfirmationDetails,
-  ToolResult,
-} from '@gemini-code/core';
-import { Part, PartUnion, PartListUnion } from '@google/genai';
-import { useCallback, useEffect, useState } from 'react';
-import {
-  HistoryItemToolGroup,
-  IndividualToolCallDisplay,
-  ToolCallStatus,
-  HistoryItemWithoutId,
-} from '../types.js';
-
-type ValidatingToolCall = {
-  status: 'validating';
-  request: ToolCallRequestInfo;
-  tool: Tool;
-};
-
-type ScheduledToolCall = {
-  status: 'scheduled';
-  request: ToolCallRequestInfo;
-  tool: Tool;
-};
-
-type ErroredToolCall = {
-  status: 'error';
-  request: ToolCallRequestInfo;
-  response: ToolCallResponseInfo;
-};
-
-type SuccessfulToolCall = {
-  status: 'success';
-  request: ToolCallRequestInfo;
-  tool: Tool;
-  response: ToolCallResponseInfo;
-};
-
-export type ExecutingToolCall = {
-  status: 'executing';
-  request: ToolCallRequestInfo;
-  tool: Tool;
-  liveOutput?: string;
-};
-
-type CancelledToolCall = {
-  status: 'cancelled';
-  request: ToolCallRequestInfo;
-  response: ToolCallResponseInfo;
-  tool: Tool;
-};
-
-type WaitingToolCall = {
-  status: 'awaiting_approval';
-  request: ToolCallRequestInfo;
-  tool: Tool;
-  confirmationDetails: ToolCallConfirmationDetails;
-};
-
-export type Status = ToolCall['status'];
-
-export type ToolCall =
-  | ValidatingToolCall
-  | ScheduledToolCall
-  | ErroredToolCall
-  | SuccessfulToolCall
-  | ExecutingToolCall
-  | CancelledToolCall
-  | WaitingToolCall;
-
-export type ScheduleFn = (
-  request: ToolCallRequestInfo | ToolCallRequestInfo[],
-) => void;
-export type CancelFn = () => void;
-export type CompletedToolCall =
-  | SuccessfulToolCall
-  | CancelledToolCall
-  | ErroredToolCall;
-
-/**
- * Formats a PartListUnion response from a tool into JSON suitable for a Gemini
- * FunctionResponse and additional Parts to include after that response.
- *
- * This is required because FunctionReponse appears to only support JSON
- * and not arbitrary parts. Including parts like inlineData or fileData
- * directly in a FunctionResponse confuses the model resulting in a failure
- * to interpret the multimodal content and context window exceeded errors.
- */
-
-export function formatLlmContentForFunctionResponse(
-  llmContent: PartListUnion,
-): {
-  functionResponseJson: Record<string, string>;
-  additionalParts: PartUnion[];
-} {
-  const additionalParts: PartUnion[] = [];
-  let functionResponseJson: Record<string, string>;
-
-  if (Array.isArray(llmContent) && llmContent.length === 1) {
-    // Ensure that length 1 arrays are treated as a single Part.
-    llmContent = llmContent[0];
-  }
-
-  if (typeof llmContent === 'string') {
-    functionResponseJson = { output: llmContent };
-  } else if (Array.isArray(llmContent)) {
-    functionResponseJson = { status: 'Tool execution succeeded.' };
-    additionalParts.push(...llmContent);
-  } else {
-    if (
-      llmContent.inlineData !== undefined ||
-      llmContent.fileData !== undefined
-    ) {
-      // For Parts like inlineData or fileData, use the returnDisplay as the textual output for the functionResponse.
-      // The actual Part will be added to additionalParts.
-      functionResponseJson = {
-        status: `Binary content of type ${llmContent.inlineData?.mimeType || llmContent.fileData?.mimeType || 'unknown'} was processed.`,
-      };
-      additionalParts.push(llmContent);
-    } else if (llmContent.text !== undefined) {
-      functionResponseJson = { output: llmContent.text };
-    } else {
-      functionResponseJson = { status: 'Tool execution succeeded.' };
-      additionalParts.push(llmContent);
-    }
-  }
-
-  return {
-    functionResponseJson,
-    additionalParts,
-  };
-}
-
-export function useToolScheduler(
-  onComplete: (tools: CompletedToolCall[]) => void,
-  config: Config,
-  setPendingHistoryItem: React.Dispatch<
-    React.SetStateAction<HistoryItemWithoutId | null>
-  >,
-): [ToolCall[], ScheduleFn, CancelFn] {
-  const [toolRegistry] = useState(() => config.getToolRegistry());
-  const [toolCalls, setToolCalls] = useState<ToolCall[]>([]);
-  const [abortController, setAbortController] = useState<AbortController>(
-    () => new AbortController(),
-  );
-
-  const isRunning = toolCalls.some(
-    (t) => t.status === 'executing' || t.status === 'awaiting_approval',
-  );
-  // Note: request array[] typically signal pending tool calls
-  const schedule = useCallback(
-    async (request: ToolCallRequestInfo | ToolCallRequestInfo[]) => {
-      if (isRunning) {
-        throw new Error(
-          'Cannot schedule tool calls while other tool calls are running',
-        );
-      }
-      const requestsToProcess = Array.isArray(request) ? request : [request];
-
-      // Step 1: Create initial calls with 'validating' status (or 'error' if tool not found)
-      // and add them to the state immediately to make the UI busy.
-      const initialNewCalls: ToolCall[] = requestsToProcess.map(
-        (r): ToolCall => {
-          const tool = toolRegistry.getTool(r.name);
-          if (!tool) {
-            return {
-              status: 'error',
-              request: r,
-              response: toolErrorResponse(
-                r,
-                new Error(`tool ${r.name} does not exist`),
-              ),
-            };
-          }
-          // Set to 'validating' immediately. This will make streamingState 'Responding'.
-          return { status: 'validating', request: r, tool };
-        },
-      );
-      setToolCalls((prevCalls) => prevCalls.concat(initialNewCalls));
-
-      // Step 2: Asynchronously check for confirmation and update status for each new call.
-      initialNewCalls.forEach(async (initialCall) => {
-        // If the call was already marked as an error (tool not found), skip further processing.
-        if (initialCall.status !== 'validating') return;
-
-        const { request: r, tool } = initialCall;
-        try {
-          const userApproval = await tool.shouldConfirmExecute(
-            r.args,
-            abortController.signal,
-          );
-          if (userApproval) {
-            // Confirmation is needed. Update status to 'awaiting_approval'.
-            setToolCalls(
-              setStatus(r.callId, 'awaiting_approval', {
-                ...userApproval,
-                onConfirm: async (outcome) => {
-                  // This onConfirm is triggered by user interaction later.
-                  await userApproval.onConfirm(outcome);
-                  setToolCalls(
-                    outcome === ToolConfirmationOutcome.Cancel
-                      ? setStatus(
-                          r.callId,
-                          'cancelled',
-                          'User did not allow tool call',
-                        )
-                      : // If confirmed, it goes to 'scheduled' to be picked up by the execution effect.
-                        setStatus(r.callId, 'scheduled'),
-                  );
-                },
-              }),
-            );
-          } else {
-            // No confirmation needed, move to 'scheduled' for execution.
-            setToolCalls(setStatus(r.callId, 'scheduled'));
-          }
-        } catch (e) {
-          // Handle errors from tool.shouldConfirmExecute() itself.
-          setToolCalls(
-            setStatus(
-              r.callId,
-              'error',
-              toolErrorResponse(
-                r,
-                e instanceof Error ? e : new Error(String(e)),
-              ),
-            ),
-          );
-        }
-      });
-    },
-    [isRunning, setToolCalls, toolRegistry, abortController.signal],
-  );
-
-  const cancel = useCallback(
-    (reason: string = 'unspecified') => {
-      abortController.abort();
-      setAbortController(new AbortController());
-      setToolCalls((tc) =>
-        tc.map((c) =>
-          c.status !== 'error' && c.status !== 'executing'
-            ? {
-                ...c,
-                status: 'cancelled',
-                response: {
-                  callId: c.request.callId,
-                  responseParts: {
-                    functionResponse: {
-                      id: c.request.callId,
-                      name: c.request.name,
-                      response: {
-                        error: `[Operation Cancelled] Reason: ${reason}`,
-                      },
-                    },
-                  },
-                  resultDisplay: undefined,
-                  error: undefined,
-                },
-              }
-            : c,
-        ),
-      );
-    },
-    [abortController],
-  );
-
-  useEffect(() => {
-    // effect for executing scheduled tool calls
-    const allToolsConfirmed = toolCalls.every(
-      (t) => t.status === 'scheduled' || t.status === 'cancelled',
-    );
-    if (allToolsConfirmed) {
-      const signal = abortController.signal;
-      toolCalls
-        .filter((t) => t.status === 'scheduled')
-        .forEach((t) => {
-          const callId = t.request.callId;
-          setToolCalls(setStatus(t.request.callId, 'executing'));
-
-          const updateOutput = t.tool.canUpdateOutput
-            ? (output: string) => {
-                setPendingHistoryItem(
-                  (prevItem: HistoryItemWithoutId | null) => {
-                    if (prevItem?.type === 'tool_group') {
-                      return {
-                        ...prevItem,
-                        tools: prevItem.tools.map(
-                          (toolDisplay: IndividualToolCallDisplay) =>
-                            toolDisplay.callId === callId &&
-                            toolDisplay.status === ToolCallStatus.Executing
-                              ? {
-                                  ...toolDisplay,
-                                  resultDisplay: output,
-                                }
-                              : toolDisplay,
-                        ),
-                      };
-                    }
-                    return prevItem;
-                  },
-                );
-                // Also update the toolCall itself so that mapToDisplay
-                // can pick up the live output if the item is not pending
-                // (e.g. if it's being re-rendered from history)
-                setToolCalls((prevToolCalls) =>
-                  prevToolCalls.map((tc) =>
-                    tc.request.callId === callId && tc.status === 'executing'
-                      ? { ...tc, liveOutput: output }
-                      : tc,
-                  ),
-                );
-              }
-            : undefined;
-
-          t.tool
-            .execute(t.request.args, signal, updateOutput)
-            .then((result: ToolResult) => {
-              if (signal.aborted) {
-                // TODO(jacobr): avoid stringifying the LLM content.
-                setToolCalls(
-                  setStatus(callId, 'cancelled', String(result.llmContent)),
-                );
-                return;
-              }
-              const { functionResponseJson, additionalParts } =
-                formatLlmContentForFunctionResponse(result.llmContent);
-              const functionResponse: Part = {
-                functionResponse: {
-                  name: t.request.name,
-                  id: callId,
-                  response: functionResponseJson,
-                },
-              };
-              const response: ToolCallResponseInfo = {
-                callId,
-                responseParts: [functionResponse, ...additionalParts],
-                resultDisplay: result.returnDisplay,
-                error: undefined,
-              };
-              setToolCalls(setStatus(callId, 'success', response));
-            })
-            .catch((e: Error) =>
-              setToolCalls(
-                setStatus(
-                  callId,
-                  'error',
-                  toolErrorResponse(
-                    t.request,
-                    e instanceof Error ? e : new Error(String(e)),
-                  ),
-                ),
-              ),
-            );
-        });
-    }
-  }, [toolCalls, toolRegistry, abortController.signal, setPendingHistoryItem]);
-
-  useEffect(() => {
-    const allDone = toolCalls.every(
-      (t) =>
-        t.status === 'success' ||
-        t.status === 'error' ||
-        t.status === 'cancelled',
-    );
-    if (toolCalls.length && allDone) {
-      setToolCalls([]);
-      onComplete(toolCalls);
-      setAbortController(() => new AbortController());
-    }
-  }, [toolCalls, onComplete]);
-
-  return [toolCalls, schedule, cancel];
-}
-
-function setStatus(
-  targetCallId: string,
-  status: 'success',
-  response: ToolCallResponseInfo,
-): (t: ToolCall[]) => ToolCall[];
-function setStatus(
-  targetCallId: string,
-  status: 'awaiting_approval',
-  confirm: ToolCallConfirmationDetails,
-): (t: ToolCall[]) => ToolCall[];
-function setStatus(
-  targetCallId: string,
-  status: 'error',
-  response: ToolCallResponseInfo,
-): (t: ToolCall[]) => ToolCall[];
-function setStatus(
-  targetCallId: string,
-  status: 'cancelled',
-  reason: string,
-): (t: ToolCall[]) => ToolCall[];
-function setStatus(
-  targetCallId: string,
-  status: 'executing' | 'scheduled' | 'validating',
-): (t: ToolCall[]) => ToolCall[];
-function setStatus(
-  targetCallId: string,
-  status: Status,
-  auxiliaryData?: unknown,
-): (t: ToolCall[]) => ToolCall[] {
-  return function (tc: ToolCall[]): ToolCall[] {
-    return tc.map((t) => {
-      if (t.request.callId !== targetCallId || t.status === 'error') {
-        return t;
-      }
-      switch (status) {
-        case 'success': {
-          const next: SuccessfulToolCall = {
-            ...t,
-            status: 'success',
-            response: auxiliaryData as ToolCallResponseInfo,
-          };
-          return next;
-        }
-        case 'error': {
-          const next: ErroredToolCall = {
-            ...t,
-            status: 'error',
-            response: auxiliaryData as ToolCallResponseInfo,
-          };
-          return next;
-        }
-        case 'awaiting_approval': {
-          const next: WaitingToolCall = {
-            ...t,
-            status: 'awaiting_approval',
-            confirmationDetails: auxiliaryData as ToolCallConfirmationDetails,
-          };
-          return next;
-        }
-        case 'scheduled': {
-          const next: ScheduledToolCall = {
-            ...t,
-            status: 'scheduled',
-          };
-          return next;
-        }
-        case 'cancelled': {
-          const next: CancelledToolCall = {
-            ...t,
-            status: 'cancelled',
-            response: {
-              callId: t.request.callId,
-              responseParts: {
-                functionResponse: {
-                  id: t.request.callId,
-                  name: t.request.name,
-                  response: {
-                    error: `[Operation Cancelled] Reason: ${auxiliaryData}`,
-                  },
-                },
-              },
-              resultDisplay: undefined,
-              error: undefined,
-            },
-          };
-          return next;
-        }
-        case 'validating': {
-          const next: ValidatingToolCall = {
-            ...(t as ValidatingToolCall), // Added type assertion for safety
-            status: 'validating',
-          };
-          return next;
-        }
-        case 'executing': {
-          const next: ExecutingToolCall = {
-            ...t,
-            status: 'executing',
-          };
-          return next;
-        }
-        default: {
-          // ensures every case is checked for above
-          const exhaustiveCheck: never = status;
-          return exhaustiveCheck;
-        }
-      }
-    });
-  };
-}
-
-const toolErrorResponse = (
-  request: ToolCallRequestInfo,
-  error: Error,
-): ToolCallResponseInfo => ({
-  callId: request.callId,
-  error,
-  responseParts: {
-    functionResponse: {
-      id: request.callId,
-      name: request.name,
-      response: { error: error.message },
-    },
-  },
-  resultDisplay: error.message,
-});
-
-function mapStatus(status: Status): ToolCallStatus {
-  switch (status) {
-    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: {
-      // ensures every case is checked for above
-      const exhaustiveCheck: never = status;
-      return exhaustiveCheck;
-    }
-  }
-}
-
-// convenient function for callers to map ToolCall back to a HistoryItem
-export function mapToDisplay(
-  tool: ToolCall[] | ToolCall,
-): HistoryItemToolGroup {
-  const tools = Array.isArray(tool) ? tool : [tool];
-  const toolsDisplays = tools.map((t): IndividualToolCallDisplay => {
-    switch (t.status) {
-      case 'success':
-        return {
-          callId: t.request.callId,
-          name: t.tool.displayName,
-          description: t.tool.getDescription(t.request.args),
-          resultDisplay: t.response.resultDisplay,
-          status: mapStatus(t.status),
-          confirmationDetails: undefined,
-          renderOutputAsMarkdown: t.tool.isOutputMarkdown,
-        };
-      case 'error':
-        return {
-          callId: t.request.callId,
-          name: t.request.name, // Use request.name as tool might be undefined
-          description: '', // No description available if tool is undefined
-          resultDisplay: t.response.resultDisplay,
-          status: mapStatus(t.status),
-          confirmationDetails: undefined,
-          renderOutputAsMarkdown: false,
-        };
-      case 'cancelled':
-        return {
-          callId: t.request.callId,
-          name: t.tool.displayName,
-          description: t.tool.getDescription(t.request.args),
-          resultDisplay: t.response.resultDisplay,
-          status: mapStatus(t.status),
-          confirmationDetails: undefined,
-          renderOutputAsMarkdown: t.tool.isOutputMarkdown,
-        };
-      case 'awaiting_approval':
-        return {
-          callId: t.request.callId,
-          name: t.tool.displayName,
-          description: t.tool.getDescription(t.request.args),
-          resultDisplay: undefined,
-          status: mapStatus(t.status),
-          confirmationDetails: t.confirmationDetails,
-          renderOutputAsMarkdown: t.tool.isOutputMarkdown,
-        };
-      case 'executing':
-        return {
-          callId: t.request.callId,
-          name: t.tool.displayName,
-          description: t.tool.getDescription(t.request.args),
-          resultDisplay: t.liveOutput ?? undefined,
-          status: mapStatus(t.status),
-          confirmationDetails: undefined,
-          renderOutputAsMarkdown: t.tool.isOutputMarkdown,
-        };
-      case 'validating': // Add this case
-        return {
-          callId: t.request.callId,
-          name: t.tool.displayName,
-          description: t.tool.getDescription(t.request.args),
-          resultDisplay: undefined,
-          status: mapStatus(t.status),
-          confirmationDetails: undefined,
-          renderOutputAsMarkdown: t.tool.isOutputMarkdown,
-        };
-      case 'scheduled':
-        return {
-          callId: t.request.callId,
-          name: t.tool.displayName,
-          description: t.tool.getDescription(t.request.args),
-          resultDisplay: undefined,
-          status: mapStatus(t.status),
-          confirmationDetails: undefined,
-          renderOutputAsMarkdown: t.tool.isOutputMarkdown,
-        };
-      default: {
-        // ensures every case is checked for above
-        const exhaustiveCheck: never = t;
-        return exhaustiveCheck;
-      }
-    }
-  });
-  const historyItem: HistoryItemToolGroup = {
-    type: 'tool_group',
-    tools: toolsDisplays,
-  };
-  return historyItem;
-}