@machina.ai/cell-cli-core 1.20.2-rc1 → 1.25.0-rc1

package/dist/src/core/coreToolScheduler.js

@@ -3,80 +3,19 @@
  * Copyright 2025 Google LLC
  * SPDX-License-Identifier: Apache-2.0
  */
-import { ToolConfirmationOutcome, ApprovalMode, logToolCall, ToolErrorType, ToolCallEvent, logToolOutputTruncated, ToolOutputTruncatedEvent, runInDevTraceSpan, } from '../index.js';
-import { READ_FILE_TOOL_NAME, SHELL_TOOL_NAME } from '../tools/tool-names.js';
-import { getResponseTextFromParts } from '../utils/generateContentResponseUtilities.js';
-import { isModifiableDeclarativeTool, modifyWithEditor, } from '../tools/modifiable-tool.js';
-import * as Diff from 'diff';
-import * as fs from 'node:fs/promises';
-import * as path from 'node:path';
-import { isShellInvocationAllowlisted, SHELL_TOOL_NAMES, } from '../utils/shell-utils.js';
-import { doesToolInvocationMatch } from '../utils/tool-utils.js';
-import levenshtein from 'fast-levenshtein';
-import { ShellToolInvocation } from '../tools/shell.js';
+import { ToolConfirmationOutcome, } from '../tools/tools.js';
+import { PolicyDecision } from '../policy/types.js';
+import { logToolCall } from '../telemetry/loggers.js';
+import { ToolErrorType } from '../tools/tool-error.js';
+import { ToolCallEvent } from '../telemetry/types.js';
+import { runInDevTraceSpan } from '../telemetry/trace.js';
+import { ToolModificationHandler } from '../scheduler/tool-modifier.js';
+import { getToolSuggestion } from '../utils/tool-utils.js';
 import { MessageBusType } from '../confirmation-bus/types.js';
-import { fireToolNotificationHook, executeToolWithHooks, } from './coreToolHookTriggers.js';
-/**
- * Formats tool output for a Gemini FunctionResponse.
- */
-function createFunctionResponsePart(callId, toolName, output) {
-    return {
-        functionResponse: {
-            id: callId,
-            name: toolName,
-            response: { output },
-        },
-    };
-}
-export function convertToFunctionResponse(toolName, callId, llmContent) {
-    const contentToProcess = Array.isArray(llmContent) && llmContent.length === 1
-        ? llmContent[0]
-        : llmContent;
-    if (typeof contentToProcess === 'string') {
-        return [createFunctionResponsePart(callId, toolName, contentToProcess)];
-    }
-    if (Array.isArray(contentToProcess)) {
-        const functionResponse = createFunctionResponsePart(callId, toolName, 'Tool execution succeeded.');
-        return [functionResponse, ...toParts(contentToProcess)];
-    }
-    // After this point, contentToProcess is a single Part object.
-    if (contentToProcess.functionResponse) {
-        if (contentToProcess.functionResponse.response?.['content']) {
-            const stringifiedOutput = getResponseTextFromParts(contentToProcess.functionResponse.response['content']) || '';
-            return [createFunctionResponsePart(callId, toolName, stringifiedOutput)];
-        }
-        // It's a functionResponse that we should pass through as is.
-        return [contentToProcess];
-    }
-    if (contentToProcess.inlineData || contentToProcess.fileData) {
-        const mimeType = contentToProcess.inlineData?.mimeType ||
-            contentToProcess.fileData?.mimeType ||
-            'unknown';
-        const functionResponse = createFunctionResponsePart(callId, toolName, `Binary content of type ${mimeType} was processed.`);
-        return [functionResponse, contentToProcess];
-    }
-    if (contentToProcess.text !== undefined) {
-        return [
-            createFunctionResponsePart(callId, toolName, contentToProcess.text),
-        ];
-    }
-    // Default case for other kinds of parts.
-    return [
-        createFunctionResponsePart(callId, toolName, 'Tool execution succeeded.'),
-    ];
-}
-function toParts(input) {
-    const parts = [];
-    for (const part of Array.isArray(input) ? input : [input]) {
-        if (typeof part === 'string') {
-            parts.push({ text: part });
-        }
-        else if (part) {
-            parts.push(part);
-        }
-    }
-    return parts;
-}
+import { fireToolNotificationHook } from './coreToolHookTriggers.js';
+import {} from '../scheduler/types.js';
+import { ToolExecutor } from '../scheduler/tool-executor.js';
+import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
 const createErrorResponse = (request, error, errorType) => ({
     callId: request.callId,
     error,
@@ -93,58 +32,6 @@ const createErrorResponse = (request, error, errorType) => ({
     errorType,
     contentLength: error.message.length,
 });
-export async function truncateAndSaveToFile(content, callId, projectTempDir, threshold, truncateLines) {
-    if (content.length <= threshold) {
-        return { content };
-    }
-    let lines = content.split('\n');
-    let fileContent = content;
-    // If the content is long but has few lines, wrap it to enable line-based truncation.
-    if (lines.length <= truncateLines) {
-        const wrapWidth = 120; // A reasonable width for wrapping.
-        const wrappedLines = [];
-        for (const line of lines) {
-            if (line.length > wrapWidth) {
-                for (let i = 0; i < line.length; i += wrapWidth) {
-                    wrappedLines.push(line.substring(i, i + wrapWidth));
-                }
-            }
-            else {
-                wrappedLines.push(line);
-            }
-        }
-        lines = wrappedLines;
-        fileContent = lines.join('\n');
-    }
-    const head = Math.floor(truncateLines / 5);
-    const beginning = lines.slice(0, head);
-    const end = lines.slice(-(truncateLines - head));
-    const truncatedContent = beginning.join('\n') + '\n... [CONTENT TRUNCATED] ...\n' + end.join('\n');
-    // Sanitize callId to prevent path traversal.
-    const safeFileName = `${path.basename(callId)}.output`;
-    const outputFile = path.join(projectTempDir, safeFileName);
-    try {
-        await fs.writeFile(outputFile, fileContent);
-        return {
-            content: `Tool output was too large and has been truncated.
-The full output has been saved to: ${outputFile}
-To read the complete output, use the ${READ_FILE_TOOL_NAME} tool with the absolute file path above. For large files, you can use the offset and limit parameters to read specific sections:
-- ${READ_FILE_TOOL_NAME} tool with offset=0, limit=100 to see the first 100 lines
-- ${READ_FILE_TOOL_NAME} tool with offset=N to skip N lines from the beginning
-- ${READ_FILE_TOOL_NAME} tool with limit=M to read only M lines at a time
-The truncated output below shows the beginning and end of the content. The marker '... [CONTENT TRUNCATED] ...' indicates where content was removed.
-This allows you to efficiently examine different parts of the output without loading the entire file.
-Truncated part of the output:
-${truncatedContent}`,
-            outputFile,
-        };
-    }
-    catch (_error) {
-        return {
-            content: truncatedContent + `\n[Note: Could not save full output to file]`,
-        };
-    }
-}
 export class CoreToolScheduler {
     // Static WeakMap to track which MessageBus instances already have a handler subscribed
     // This prevents duplicate subscriptions when multiple CoreToolScheduler instances are created
@@ -161,35 +48,38 @@ export class CoreToolScheduler {
     requestQueue = [];
     toolCallQueue = [];
     completedToolCallsForBatch = [];
+    toolExecutor;
+    toolModifier;
     constructor(options) {
         this.config = options.config;
         this.outputUpdateHandler = options.outputUpdateHandler;
         this.onAllToolCallsComplete = options.onAllToolCallsComplete;
         this.onToolCallsUpdate = options.onToolCallsUpdate;
         this.getPreferredEditor = options.getPreferredEditor;
+        this.toolExecutor = new ToolExecutor(this.config);
+        this.toolModifier = new ToolModificationHandler();
         // Subscribe to message bus for ASK_USER policy decisions
         // Use a static WeakMap to ensure we only subscribe ONCE per MessageBus instance
         // This prevents memory leaks when multiple CoreToolScheduler instances are created
         // (e.g., on every React render, or for each non-interactive tool call)
-        if (this.config.getEnableMessageBusIntegration()) {
-            const messageBus = this.config.getMessageBus();
-            // Check if we've already subscribed a handler to this message bus
-            if (!CoreToolScheduler.subscribedMessageBuses.has(messageBus)) {
-                // Create a shared handler that will be used for this message bus
-                const sharedHandler = (request) => {
-                    // When ASK_USER policy decision is made, respond with requiresUserConfirmation=true
-                    // to tell tools to use their legacy confirmation flow
-                    messageBus.publish({
-                        type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
-                        correlationId: request.correlationId,
-                        confirmed: false,
-                        requiresUserConfirmation: true,
-                    });
-                };
-                messageBus.subscribe(MessageBusType.TOOL_CONFIRMATION_REQUEST, sharedHandler);
-                // Store the handler in the WeakMap so we don't subscribe again
-                CoreToolScheduler.subscribedMessageBuses.set(messageBus, sharedHandler);
-            }
+        const messageBus = this.config.getMessageBus();
+        // Check if we've already subscribed a handler to this message bus
+        if (!CoreToolScheduler.subscribedMessageBuses.has(messageBus)) {
+            // Create a shared handler that will be used for this message bus
+            const sharedHandler = (request) => {
+                // When ASK_USER policy decision is made, respond with requiresUserConfirmation=true
+                // to tell tools to use their legacy confirmation flow
+                // eslint-disable-next-line @typescript-eslint/no-floating-promises
+                messageBus.publish({
+                    type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
+                    correlationId: request.correlationId,
+                    confirmed: false,
+                    requiresUserConfirmation: true,
+                });
+            };
+            messageBus.subscribe(MessageBusType.TOOL_CONFIRMATION_REQUEST, sharedHandler);
+            // Store the handler in the WeakMap so we don't subscribe again
+            CoreToolScheduler.subscribedMessageBuses.set(messageBus, sharedHandler);
         }
     }
     setStatusInternal(targetCallId, newStatus, signal, auxiliaryData) {
@@ -266,6 +156,7 @@ export class CoreToolScheduler {
                                 fileName: waitingCall.confirmationDetails.fileName,
                                 originalContent: waitingCall.confirmationDetails.originalContent,
                                 newContent: waitingCall.confirmationDetails.newContent,
+                                filePath: waitingCall.confirmationDetails.filePath,
                             };
                         }
                     }
@@ -362,34 +253,6 @@ export class CoreToolScheduler {
             return new Error(String(e));
         }
     }
-    /**
-     * Generates a suggestion string for a tool name that was not found in the registry.
-     * It finds the closest matches based on Levenshtein distance.
-     * @param unknownToolName The tool name that was not found.
-     * @param topN The number of suggestions to return. Defaults to 3.
-     * @returns A suggestion string like " Did you mean 'tool'?" or " Did you mean one of: 'tool1', 'tool2'?", or an empty string if no suggestions are found.
-     */
-    getToolSuggestion(unknownToolName, topN = 3) {
-        const allToolNames = this.config.getToolRegistry().getAllToolNames();
-        const matches = allToolNames.map((toolName) => ({
-            name: toolName,
-            distance: levenshtein.get(unknownToolName, toolName),
-        }));
-        matches.sort((a, b) => a.distance - b.distance);
-        const topNResults = matches.slice(0, topN);
-        if (topNResults.length === 0) {
-            return '';
-        }
-        const suggestedNames = topNResults
-            .map((match) => `"${match.name}"`)
-            .join(', ');
-        if (topNResults.length > 1) {
-            return ` Did you mean one of: ${suggestedNames}?`;
-        }
-        else {
-            return ` Did you mean ${suggestedNames}?`;
-        }
-    }
     schedule(request, signal) {
         return runInDevTraceSpan({ name: 'schedule' }, async ({ metadata: spanMetadata }) => {
             spanMetadata.input = request;
@@ -456,7 +319,7 @@ export class CoreToolScheduler {
                     .getToolRegistry()
                     .getTool(reqInfo.name);
                 if (!toolInstance) {
-                    const suggestion = this.getToolSuggestion(reqInfo.name);
+                    const suggestion = getToolSuggestion(reqInfo.name, this.config.getToolRegistry().getAllToolNames());
                     const errorMessage = `Tool "${reqInfo.name}" not found in registry. Tools must use the exact names that are registered.${suggestion}`;
                     return {
                         status: 'error',
@@ -523,17 +386,40 @@ export class CoreToolScheduler {
                     await this.checkAndNotifyCompletion(signal);
                     return;
                 }
-                const confirmationDetails = await invocation.shouldConfirmExecute(signal);
-                if (!confirmationDetails) {
+                // Policy Check using PolicyEngine
+                // We must reconstruct the FunctionCall format expected by PolicyEngine
+                const toolCallForPolicy = {
+                    name: toolCall.request.name,
+                    args: toolCall.request.args,
+                };
+                const serverName = toolCall.tool instanceof DiscoveredMCPTool
+                    ? toolCall.tool.serverName
+                    : undefined;
+                const { decision } = await this.config
+                    .getPolicyEngine()
+                    .check(toolCallForPolicy, serverName);
+                if (decision === PolicyDecision.DENY) {
+                    const errorMessage = `Tool execution denied by policy.`;
+                    this.setStatusInternal(reqInfo.callId, 'error', signal, createErrorResponse(reqInfo, new Error(errorMessage), ToolErrorType.POLICY_VIOLATION));
+                    await this.checkAndNotifyCompletion(signal);
+                    return;
+                }
+                if (decision === PolicyDecision.ALLOW) {
                     this.setToolCallOutcome(reqInfo.callId, ToolConfirmationOutcome.ProceedAlways);
                     this.setStatusInternal(reqInfo.callId, 'scheduled', signal);
                 }
                 else {
-                    if (this.isAutoApproved(toolCall)) {
+                    // PolicyDecision.ASK_USER
+                    // We need confirmation details to show to the user
+                    const confirmationDetails = await invocation.shouldConfirmExecute(signal);
+                    if (!confirmationDetails) {
                         this.setToolCallOutcome(reqInfo.callId, ToolConfirmationOutcome.ProceedAlways);
                         this.setStatusInternal(reqInfo.callId, 'scheduled', signal);
                     }
                     else {
+                        if (!this.config.isInteractive()) {
+                            throw new Error(`Tool execution for "${toolCall.tool.displayName || toolCall.tool.name}" requires user confirmation, which is not supported in non-interactive mode.`);
+                        }
                         // Fire Notification hook before showing confirmation to user
                         const messageBus = this.config.getMessageBus();
                         const hooksEnabled = this.config.getEnableHooks();
@@ -543,11 +429,14 @@ export class CoreToolScheduler {
                         // Allow IDE to resolve confirmation
                         if (confirmationDetails.type === 'edit' &&
                             confirmationDetails.ideConfirmation) {
+                            // eslint-disable-next-line @typescript-eslint/no-floating-promises
                             confirmationDetails.ideConfirmation.then((resolution) => {
                                 if (resolution.status === 'accepted') {
+                                    // eslint-disable-next-line @typescript-eslint/no-floating-promises
                                     this.handleConfirmationResponse(reqInfo.callId, confirmationDetails.onConfirm, ToolConfirmationOutcome.ProceedOnce, signal);
                                 }
                                 else {
+                                    // eslint-disable-next-line @typescript-eslint/no-floating-promises
                                     this.handleConfirmationResponse(reqInfo.callId, confirmationDetails.onConfirm, ToolConfirmationOutcome.Cancel, signal);
                                 }
                             });
@@ -587,61 +476,50 @@ export class CoreToolScheduler {
         }
         else if (outcome === ToolConfirmationOutcome.ModifyWithEditor) {
             const waitingToolCall = toolCall;
-            if (isModifiableDeclarativeTool(waitingToolCall.tool)) {
-                const modifyContext = waitingToolCall.tool.getModifyContext(signal);
-                const editorType = this.getPreferredEditor();
-                if (!editorType) {
-                    return;
-                }
+            const editorType = this.getPreferredEditor();
+            if (!editorType) {
+                return;
+            }
+            this.setStatusInternal(callId, 'awaiting_approval', signal, {
+                ...waitingToolCall.confirmationDetails,
+                isModifying: true,
+            });
+            const result = await this.toolModifier.handleModifyWithEditor(waitingToolCall, editorType, signal);
+            // Restore status (isModifying: false) and update diff if result exists
+            if (result) {
+                this.setArgsInternal(callId, result.updatedParams);
                 this.setStatusInternal(callId, 'awaiting_approval', signal, {
                     ...waitingToolCall.confirmationDetails,
-                    isModifying: true,
+                    fileDiff: result.updatedDiff,
+                    isModifying: false,
                 });
-                const contentOverrides = waitingToolCall.confirmationDetails.type === 'edit'
-                    ? {
-                        currentContent: waitingToolCall.confirmationDetails.originalContent,
-                        proposedContent: waitingToolCall.confirmationDetails.newContent,
-                    }
-                    : undefined;
-                const { updatedParams, updatedDiff } = await modifyWithEditor(waitingToolCall.request.args, modifyContext, editorType, signal, contentOverrides);
-                this.setArgsInternal(callId, updatedParams);
+            }
+            else {
                 this.setStatusInternal(callId, 'awaiting_approval', signal, {
                     ...waitingToolCall.confirmationDetails,
-                    fileDiff: updatedDiff,
                     isModifying: false,
                 });
             }
         }
         else {
-            // If the client provided new content, apply it before scheduling.
+            // If the client provided new content, apply it and wait for
+            // re-confirmation.
             if (payload?.newContent && toolCall) {
-                await this._applyInlineModify(toolCall, payload, signal);
+                const result = await this.toolModifier.applyInlineModify(toolCall, payload, signal);
+                if (result) {
+                    this.setArgsInternal(callId, result.updatedParams);
+                    this.setStatusInternal(callId, 'awaiting_approval', signal, {
+                        ...toolCall.confirmationDetails,
+                        fileDiff: result.updatedDiff,
+                    });
+                    // After an inline modification, wait for another user confirmation.
+                    return;
+                }
             }
             this.setStatusInternal(callId, 'scheduled', signal);
         }
         await this.attemptExecutionOfScheduledCalls(signal);
     }
-    /**
-     * Applies user-provided content changes to a tool call that is awaiting confirmation.
-     * This method updates the tool's arguments and refreshes the confirmation prompt with a new diff
-     * before the tool is scheduled for execution.
-     * @private
-     */
-    async _applyInlineModify(toolCall, payload, signal) {
-        if (toolCall.confirmationDetails.type !== 'edit' ||
-            !isModifiableDeclarativeTool(toolCall.tool)) {
-            return;
-        }
-        const modifyContext = toolCall.tool.getModifyContext(signal);
-        const currentContent = await modifyContext.getCurrentContent(toolCall.request.args);
-        const updatedParams = modifyContext.createUpdatedParams(currentContent, payload.newContent, toolCall.request.args);
-        const updatedDiff = Diff.createPatch(modifyContext.getFilePath(toolCall.request.args), currentContent, payload.newContent, 'Current', 'Proposed');
-        this.setArgsInternal(toolCall.request.callId, updatedParams);
-        this.setStatusInternal(toolCall.request.callId, 'awaiting_approval', signal, {
-            ...toolCall.confirmationDetails,
-            fileDiff: updatedDiff,
-        });
-    }
     async attemptExecutionOfScheduledCalls(signal) {
         const allCallsFinalOrScheduled = this.toolCalls.every((call) => call.status === 'scheduled' ||
             call.status === 'cancelled' ||
@@ -652,111 +530,36 @@ export class CoreToolScheduler {
             for (const toolCall of callsToExecute) {
                 if (toolCall.status !== 'scheduled')
                     continue;
-                const scheduledCall = toolCall;
-                const { callId, name: toolName } = scheduledCall.request;
-                const invocation = scheduledCall.invocation;
-                this.setStatusInternal(callId, 'executing', signal);
-                const liveOutputCallback = scheduledCall.tool.canUpdateOutput && this.outputUpdateHandler
-                    ? (outputChunk) => {
+                this.setStatusInternal(toolCall.request.callId, 'executing', signal);
+                const executingCall = this.toolCalls.find((c) => c.request.callId === toolCall.request.callId);
+                if (!executingCall) {
+                    // Should not happen, but safe guard
+                    continue;
+                }
+                const completedCall = await this.toolExecutor.execute({
+                    call: executingCall,
+                    signal,
+                    outputUpdateHandler: (callId, output) => {
                         if (this.outputUpdateHandler) {
-                            this.outputUpdateHandler(callId, outputChunk);
+                            this.outputUpdateHandler(callId, output);
                         }
                         this.toolCalls = this.toolCalls.map((tc) => tc.request.callId === callId && tc.status === 'executing'
-                            ? { ...tc, liveOutput: outputChunk }
+                            ? { ...tc, liveOutput: output }
                             : tc);
                         this.notifyToolCallsUpdate();
-                    }
-                    : undefined;
-                const shellExecutionConfig = this.config.getShellExecutionConfig();
-                const hooksEnabled = this.config.getEnableHooks();
-                const messageBus = this.config.getMessageBus();
-                await runInDevTraceSpan({
-                    name: toolCall.tool.name,
-                    attributes: { type: 'tool-call' },
-                }, async ({ metadata: spanMetadata }) => {
-                    spanMetadata.input = {
-                        request: toolCall.request,
-                    };
-                    // TODO: Refactor to remove special casing for ShellToolInvocation.
-                    // Introduce a generic callbacks object for the execute method to handle
-                    // things like `onPid` and `onLiveOutput`. This will make the scheduler
-                    // agnostic to the invocation type.
-                    let promise;
-                    if (invocation instanceof ShellToolInvocation) {
-                        const setPidCallback = (pid) => {
-                            this.toolCalls = this.toolCalls.map((tc) => tc.request.callId === callId && tc.status === 'executing'
-                                ? { ...tc, pid }
-                                : tc);
-                            this.notifyToolCallsUpdate();
-                        };
-                        promise = executeToolWithHooks(invocation, toolName, signal, messageBus, hooksEnabled, liveOutputCallback, shellExecutionConfig, setPidCallback);
-                    }
-                    else {
-                        promise = executeToolWithHooks(invocation, toolName, signal, messageBus, hooksEnabled, liveOutputCallback, shellExecutionConfig);
-                    }
-                    try {
-                        const toolResult = await promise;
-                        spanMetadata.output = toolResult;
-                        if (signal.aborted) {
-                            this.setStatusInternal(callId, 'cancelled', signal, 'User cancelled tool execution.');
-                        }
-                        else if (toolResult.error === undefined) {
-                            let content = toolResult.llmContent;
-                            let outputFile = undefined;
-                            const contentLength = typeof content === 'string' ? content.length : undefined;
-                            if (typeof content === 'string' &&
-                                toolName === SHELL_TOOL_NAME &&
-                                this.config.getEnableToolOutputTruncation() &&
-                                this.config.getTruncateToolOutputThreshold() > 0 &&
-                                this.config.getTruncateToolOutputLines() > 0) {
-                                const originalContentLength = content.length;
-                                const threshold = this.config.getTruncateToolOutputThreshold();
-                                const lines = this.config.getTruncateToolOutputLines();
-                                const truncatedResult = await truncateAndSaveToFile(content, callId, this.config.storage.getProjectTempDir(), threshold, lines);
-                                content = truncatedResult.content;
-                                outputFile = truncatedResult.outputFile;
-                                if (outputFile) {
-                                    logToolOutputTruncated(this.config, new ToolOutputTruncatedEvent(scheduledCall.request.prompt_id, {
-                                        toolName,
-                                        originalContentLength,
-                                        truncatedContentLength: content.length,
-                                        threshold,
-                                        lines,
-                                    }));
-                                }
-                            }
-                            const response = convertToFunctionResponse(toolName, callId, content);
-                            const successResponse = {
-                                callId,
-                                responseParts: response,
-                                resultDisplay: toolResult.returnDisplay,
-                                error: undefined,
-                                errorType: undefined,
-                                outputFile,
-                                contentLength,
-                            };
-                            this.setStatusInternal(callId, 'success', signal, successResponse);
-                        }
-                        else {
-                            // It is a failure
-                            const error = new Error(toolResult.error.message);
-                            const errorResponse = createErrorResponse(scheduledCall.request, error, toolResult.error.type);
-                            this.setStatusInternal(callId, 'error', signal, errorResponse);
-                        }
-                    }
-                    catch (executionError) {
-                        spanMetadata.error = executionError;
-                        if (signal.aborted) {
-                            this.setStatusInternal(callId, 'cancelled', signal, 'User cancelled tool execution.');
-                        }
-                        else {
-                            this.setStatusInternal(callId, 'error', signal, createErrorResponse(scheduledCall.request, executionError instanceof Error
-                                ? executionError
-                                : new Error(String(executionError)), ToolErrorType.UNHANDLED_EXCEPTION));
-                        }
-                    }
-                    await this.checkAndNotifyCompletion(signal);
+                    },
+                    onUpdateToolCall: (updatedCall) => {
+                        this.toolCalls = this.toolCalls.map((tc) => tc.request.callId === updatedCall.request.callId
+                            ? updatedCall
+                            : tc);
+                        this.notifyToolCallsUpdate();
+                    },
                 });
+                this.toolCalls = this.toolCalls.map((tc) => tc.request.callId === completedCall.request.callId
+                    ? completedCall
+                    : tc);
+                this.notifyToolCallsUpdate();
+                await this.checkAndNotifyCompletion(signal);
             }
         }
     }
@@ -791,20 +594,35 @@ export class CoreToolScheduler {
             if (signal.aborted) {
                 this._cancelAllQueuedCalls();
             }
+            // If we are already finalizing, another concurrent call to
+            // checkAndNotifyCompletion will just return. The ongoing finalized loop
+            // will pick up any new tools added to completedToolCallsForBatch.
+            if (this.isFinalizingToolCalls) {
+                return;
+            }
             // If there's nothing to report and we weren't cancelled, we can stop.
             // But if we were cancelled, we must proceed to potentially start the next queued request.
             if (this.completedToolCallsForBatch.length === 0 && !signal.aborted) {
                 return;
             }
-            if (this.onAllToolCallsComplete) {
-                this.isFinalizingToolCalls = true;
-                // Use the batch array, not the (now empty) active array.
-                await this.onAllToolCallsComplete(this.completedToolCallsForBatch);
-                this.completedToolCallsForBatch = []; // Clear after reporting.
+            this.isFinalizingToolCalls = true;
+            try {
+                // We use a while loop here to ensure that if new tools are added to the
+                // batch (e.g., via cancellation) while we are awaiting
+                // onAllToolCallsComplete, they are also reported before we finish.
+                while (this.completedToolCallsForBatch.length > 0) {
+                    const batchToReport = [...this.completedToolCallsForBatch];
+                    this.completedToolCallsForBatch = [];
+                    if (this.onAllToolCallsComplete) {
+                        await this.onAllToolCallsComplete(batchToReport);
+                    }
+                }
+            }
+            finally {
                 this.isFinalizingToolCalls = false;
+                this.isCancelling = false;
+                this.notifyToolCallsUpdate();
             }
-            this.isCancelling = false;
-            this.notifyToolCallsUpdate();
             // After completion of the entire batch, process the next item in the main request queue.
             if (this.requestQueue.length > 0) {
                 const next = this.requestQueue.shift();
@@ -877,17 +695,5 @@ export class CoreToolScheduler {
             };
         });
     }
-    isAutoApproved(toolCall) {
-        if (this.config.getApprovalMode() === ApprovalMode.YOLO) {
-            return true;
-        }
-        const allowedTools = this.config.getAllowedTools() || [];
-        const { tool, invocation } = toolCall;
-        const toolName = typeof tool === 'string' ? tool : tool.name;
-        if (SHELL_TOOL_NAMES.includes(toolName)) {
-            return isShellInvocationAllowlisted(invocation, allowedTools);
-        }
-        return doesToolInvocationMatch(tool, invocation, allowedTools);
-    }
 }
 //# sourceMappingURL=coreToolScheduler.js.map