diff options
Diffstat (limited to 'packages/cli/src/ui/hooks/useToolScheduler.ts')
| -rw-r--r-- | packages/cli/src/ui/hooks/useToolScheduler.ts | 626 |
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; -} |