From 409f9c825b832d5b58b13114b9d6282093ce7ccc Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Wed, 14 Jan 2026 23:42:31 -0500 Subject: [PATCH] feat(scheduler): add SchedulerStateManager for reactive tool state (#16651) --- .../core/src/scheduler/state-manager.test.ts | 532 ++++++++++++++++++ packages/core/src/scheduler/state-manager.ts | 482 ++++++++++++++++ 2 files changed, 1014 insertions(+) create mode 100644 packages/core/src/scheduler/state-manager.test.ts create mode 100644 packages/core/src/scheduler/state-manager.ts diff --git a/packages/core/src/scheduler/state-manager.test.ts b/packages/core/src/scheduler/state-manager.test.ts new file mode 100644 index 0000000000..504dbbf007 --- /dev/null +++ b/packages/core/src/scheduler/state-manager.test.ts @@ -0,0 +1,532 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { SchedulerStateManager } from './state-manager.js'; +import type { + ValidatingToolCall, + WaitingToolCall, + SuccessfulToolCall, + ErroredToolCall, + CancelledToolCall, + ExecutingToolCall, + ToolCallRequestInfo, + ToolCallResponseInfo, +} from './types.js'; +import { + ToolConfirmationOutcome, + type AnyDeclarativeTool, + type AnyToolInvocation, +} from '../tools/tools.js'; +import { MessageBusType } from '../confirmation-bus/types.js'; +import type { MessageBus } from '../confirmation-bus/message-bus.js'; + +describe('SchedulerStateManager', () => { + const mockRequest: ToolCallRequestInfo = { + callId: 'call-1', + name: 'test-tool', + args: { foo: 'bar' }, + isClientInitiated: false, + prompt_id: 'prompt-1', + }; + + const mockTool = { + name: 'test-tool', + displayName: 'Test Tool', + } as AnyDeclarativeTool; + + const mockInvocation = { + shouldConfirmExecute: vi.fn(), + } as unknown as AnyToolInvocation; + + const createValidatingCall = (id = 'call-1'): ValidatingToolCall => ({ + status: 'validating', + request: { ...mockRequest, callId: id }, + tool: mockTool, + invocation: mockInvocation, + startTime: Date.now(), + }); + + const createMockResponse = (id: string): ToolCallResponseInfo => ({ + callId: id, + responseParts: [], + resultDisplay: 'Success', + error: undefined, + errorType: undefined, + }); + + let stateManager: SchedulerStateManager; + let mockMessageBus: MessageBus; + let onUpdate: (calls: unknown[]) => void; + + beforeEach(() => { + onUpdate = vi.fn(); + mockMessageBus = { + publish: vi.fn(), + subscribe: vi.fn(), + unsubscribe: vi.fn(), + } as unknown as MessageBus; + + // Capture the update when published + vi.mocked(mockMessageBus.publish).mockImplementation((msg) => { + // Return a Promise to satisfy the void | Promise signature if needed, + // though typically mocks handle it. + if (msg.type === MessageBusType.TOOL_CALLS_UPDATE) { + onUpdate(msg.toolCalls); + } + return Promise.resolve(); + }); + + stateManager = new SchedulerStateManager(mockMessageBus); + }); + + describe('Initialization', () => { + it('should start with empty state', () => { + expect(stateManager.isActive).toBe(false); + expect(stateManager.activeCallCount).toBe(0); + expect(stateManager.queueLength).toBe(0); + expect(stateManager.getSnapshot()).toEqual([]); + }); + }); + + describe('Lookup Operations', () => { + it('should find tool calls in active calls', () => { + const call = createValidatingCall('active-1'); + stateManager.enqueue([call]); + stateManager.dequeue(); + expect(stateManager.getToolCall('active-1')).toEqual(call); + }); + + it('should find tool calls in the queue', () => { + const call = createValidatingCall('queued-1'); + stateManager.enqueue([call]); + expect(stateManager.getToolCall('queued-1')).toEqual(call); + }); + + it('should find tool calls in the completed batch', () => { + const call = createValidatingCall('completed-1'); + stateManager.enqueue([call]); + stateManager.dequeue(); + stateManager.updateStatus( + 'completed-1', + 'success', + createMockResponse('completed-1'), + ); + stateManager.finalizeCall('completed-1'); + expect(stateManager.getToolCall('completed-1')).toBeDefined(); + }); + + it('should return undefined for non-existent callIds', () => { + expect(stateManager.getToolCall('void')).toBeUndefined(); + }); + }); + + describe('Queue Management', () => { + it('should enqueue calls and notify', () => { + const call = createValidatingCall(); + stateManager.enqueue([call]); + + expect(stateManager.queueLength).toBe(1); + expect(onUpdate).toHaveBeenCalledWith([call]); + }); + + it('should dequeue calls and notify', () => { + const call = createValidatingCall(); + stateManager.enqueue([call]); + + const dequeued = stateManager.dequeue(); + + expect(dequeued).toEqual(call); + expect(stateManager.queueLength).toBe(0); + expect(stateManager.activeCallCount).toBe(1); + expect(onUpdate).toHaveBeenCalled(); + }); + + it('should return undefined when dequeueing from empty queue', () => { + const dequeued = stateManager.dequeue(); + expect(dequeued).toBeUndefined(); + }); + }); + + describe('Status Transitions', () => { + it('should transition validating to scheduled', () => { + const call = createValidatingCall(); + stateManager.enqueue([call]); + stateManager.dequeue(); + + stateManager.updateStatus(call.request.callId, 'scheduled'); + + const snapshot = stateManager.getSnapshot(); + expect(snapshot[0].status).toBe('scheduled'); + expect(snapshot[0].request.callId).toBe(call.request.callId); + }); + + it('should transition scheduled to executing', () => { + const call = createValidatingCall(); + stateManager.enqueue([call]); + stateManager.dequeue(); + stateManager.updateStatus(call.request.callId, 'scheduled'); + + stateManager.updateStatus(call.request.callId, 'executing'); + + expect(stateManager.firstActiveCall?.status).toBe('executing'); + }); + + it('should transition to success and move to completed batch', () => { + const call = createValidatingCall(); + stateManager.enqueue([call]); + stateManager.dequeue(); + + const response: ToolCallResponseInfo = { + callId: call.request.callId, + responseParts: [], + resultDisplay: 'Success', + error: undefined, + errorType: undefined, + }; + + stateManager.updateStatus(call.request.callId, 'success', response); + stateManager.finalizeCall(call.request.callId); + + expect(stateManager.isActive).toBe(false); + expect(stateManager.completedBatch).toHaveLength(1); + const completed = stateManager.completedBatch[0] as SuccessfulToolCall; + expect(completed.status).toBe('success'); + expect(completed.response).toEqual(response); + expect(completed.durationMs).toBeDefined(); + }); + + it('should transition to error and move to completed batch', () => { + const call = createValidatingCall(); + stateManager.enqueue([call]); + stateManager.dequeue(); + + const response: ToolCallResponseInfo = { + callId: call.request.callId, + responseParts: [], + resultDisplay: 'Error', + error: new Error('Failed'), + errorType: undefined, + }; + + stateManager.updateStatus(call.request.callId, 'error', response); + stateManager.finalizeCall(call.request.callId); + + expect(stateManager.isActive).toBe(false); + expect(stateManager.completedBatch).toHaveLength(1); + const completed = stateManager.completedBatch[0] as ErroredToolCall; + expect(completed.status).toBe('error'); + expect(completed.response).toEqual(response); + }); + + it('should transition to awaiting_approval with details', () => { + const call = createValidatingCall(); + stateManager.enqueue([call]); + stateManager.dequeue(); + + const details = { + type: 'info' as const, + title: 'Confirm', + prompt: 'Proceed?', + onConfirm: vi.fn(), + }; + + stateManager.updateStatus( + call.request.callId, + 'awaiting_approval', + details, + ); + + const active = stateManager.firstActiveCall as WaitingToolCall; + expect(active.status).toBe('awaiting_approval'); + expect(active.confirmationDetails).toEqual(details); + }); + + it('should transition to awaiting_approval with event-driven format', () => { + const call = createValidatingCall(); + stateManager.enqueue([call]); + stateManager.dequeue(); + + const details = { + type: 'info' as const, + title: 'Confirm', + prompt: 'Proceed?', + }; + const eventDrivenData = { + correlationId: 'corr-123', + confirmationDetails: details, + }; + + stateManager.updateStatus( + call.request.callId, + 'awaiting_approval', + eventDrivenData, + ); + + const active = stateManager.firstActiveCall as WaitingToolCall; + expect(active.status).toBe('awaiting_approval'); + expect(active.correlationId).toBe('corr-123'); + expect(active.confirmationDetails).toEqual(details); + }); + + it('should preserve diff when cancelling an edit tool call', () => { + const call = createValidatingCall(); + stateManager.enqueue([call]); + stateManager.dequeue(); + + const details = { + type: 'edit' as const, + title: 'Edit', + fileName: 'test.txt', + filePath: '/path/to/test.txt', + fileDiff: 'diff', + originalContent: 'old', + newContent: 'new', + onConfirm: vi.fn(), + }; + + stateManager.updateStatus( + call.request.callId, + 'awaiting_approval', + details, + ); + stateManager.updateStatus( + call.request.callId, + 'cancelled', + 'User said no', + ); + stateManager.finalizeCall(call.request.callId); + + const completed = stateManager.completedBatch[0] as CancelledToolCall; + expect(completed.status).toBe('cancelled'); + expect(completed.response.resultDisplay).toEqual({ + fileDiff: 'diff', + fileName: 'test.txt', + filePath: '/path/to/test.txt', + originalContent: 'old', + newContent: 'new', + }); + }); + + it('should ignore status updates for non-existent callIds', () => { + stateManager.updateStatus('unknown', 'scheduled'); + expect(onUpdate).not.toHaveBeenCalled(); + }); + + it('should ignore status updates for terminal calls', () => { + const call = createValidatingCall(); + stateManager.enqueue([call]); + stateManager.dequeue(); + stateManager.updateStatus( + call.request.callId, + 'success', + createMockResponse(call.request.callId), + ); + stateManager.finalizeCall(call.request.callId); + + vi.mocked(onUpdate).mockClear(); + stateManager.updateStatus(call.request.callId, 'scheduled'); + expect(onUpdate).not.toHaveBeenCalled(); + }); + + it('should only finalize terminal calls', () => { + const call = createValidatingCall(); + stateManager.enqueue([call]); + stateManager.dequeue(); + + stateManager.updateStatus(call.request.callId, 'executing'); + stateManager.finalizeCall(call.request.callId); + + expect(stateManager.isActive).toBe(true); + expect(stateManager.completedBatch).toHaveLength(0); + + stateManager.updateStatus( + call.request.callId, + 'success', + createMockResponse(call.request.callId), + ); + stateManager.finalizeCall(call.request.callId); + + expect(stateManager.isActive).toBe(false); + expect(stateManager.completedBatch).toHaveLength(1); + }); + + it('should merge liveOutput and pid during executing updates', () => { + const call = createValidatingCall(); + stateManager.enqueue([call]); + stateManager.dequeue(); + + // Start executing + stateManager.updateStatus(call.request.callId, 'executing'); + let active = stateManager.firstActiveCall as ExecutingToolCall; + expect(active.status).toBe('executing'); + expect(active.liveOutput).toBeUndefined(); + + // Update with live output + stateManager.updateStatus(call.request.callId, 'executing', { + liveOutput: 'chunk 1', + }); + active = stateManager.firstActiveCall as ExecutingToolCall; + expect(active.liveOutput).toBe('chunk 1'); + + // Update with pid (should preserve liveOutput) + stateManager.updateStatus(call.request.callId, 'executing', { + pid: 1234, + }); + active = stateManager.firstActiveCall as ExecutingToolCall; + expect(active.liveOutput).toBe('chunk 1'); + expect(active.pid).toBe(1234); + + // Update live output again (should preserve pid) + stateManager.updateStatus(call.request.callId, 'executing', { + liveOutput: 'chunk 2', + }); + active = stateManager.firstActiveCall as ExecutingToolCall; + expect(active.liveOutput).toBe('chunk 2'); + expect(active.pid).toBe(1234); + }); + }); + + describe('Argument Updates', () => { + it('should update args and invocation', () => { + const call = createValidatingCall(); + stateManager.enqueue([call]); + stateManager.dequeue(); + + const newArgs = { foo: 'updated' }; + const newInvocation = { ...mockInvocation } as AnyToolInvocation; + + stateManager.updateArgs(call.request.callId, newArgs, newInvocation); + + const active = stateManager.firstActiveCall; + if (active && 'invocation' in active) { + expect(active.invocation).toEqual(newInvocation); + } else { + throw new Error('Active call should have invocation'); + } + }); + + it('should ignore arg updates for errored calls', () => { + const call = createValidatingCall(); + stateManager.enqueue([call]); + stateManager.dequeue(); + stateManager.updateStatus( + call.request.callId, + 'error', + createMockResponse(call.request.callId), + ); + stateManager.finalizeCall(call.request.callId); + + stateManager.updateArgs( + call.request.callId, + { foo: 'new' }, + mockInvocation, + ); + + const completed = stateManager.completedBatch[0]; + expect(completed.request.args).toEqual(mockRequest.args); + }); + }); + + describe('Outcome Tracking', () => { + it('should set outcome and notify', () => { + const call = createValidatingCall(); + stateManager.enqueue([call]); + stateManager.dequeue(); + + stateManager.setOutcome( + call.request.callId, + ToolConfirmationOutcome.ProceedAlways, + ); + + const active = stateManager.firstActiveCall; + expect(active?.outcome).toBe(ToolConfirmationOutcome.ProceedAlways); + expect(onUpdate).toHaveBeenCalled(); + }); + }); + + describe('Batch Operations', () => { + it('should cancel all queued calls', () => { + stateManager.enqueue([ + createValidatingCall('1'), + createValidatingCall('2'), + ]); + + stateManager.cancelAllQueued('Batch cancel'); + + expect(stateManager.queueLength).toBe(0); + expect(stateManager.completedBatch).toHaveLength(2); + expect( + stateManager.completedBatch.every((c) => c.status === 'cancelled'), + ).toBe(true); + }); + + it('should clear batch and notify', () => { + const call = createValidatingCall(); + stateManager.enqueue([call]); + stateManager.dequeue(); + stateManager.updateStatus( + call.request.callId, + 'success', + createMockResponse(call.request.callId), + ); + stateManager.finalizeCall(call.request.callId); + + stateManager.clearBatch(); + + expect(stateManager.completedBatch).toHaveLength(0); + expect(onUpdate).toHaveBeenCalledWith([]); + }); + + it('should return a copy of the completed batch (defensive)', () => { + const call = createValidatingCall(); + stateManager.enqueue([call]); + stateManager.dequeue(); + stateManager.updateStatus( + call.request.callId, + 'success', + createMockResponse(call.request.callId), + ); + stateManager.finalizeCall(call.request.callId); + + const batch = stateManager.completedBatch; + expect(batch).toHaveLength(1); + + // Mutate the returned array + batch.pop(); + expect(batch).toHaveLength(0); + + // Verify internal state is unchanged + expect(stateManager.completedBatch).toHaveLength(1); + }); + }); + + describe('Snapshot and Ordering', () => { + it('should return snapshot in order: completed, active, queue', () => { + // 1. Completed + const call1 = createValidatingCall('1'); + stateManager.enqueue([call1]); + stateManager.dequeue(); + stateManager.updateStatus('1', 'success', createMockResponse('1')); + stateManager.finalizeCall('1'); + + // 2. Active + const call2 = createValidatingCall('2'); + stateManager.enqueue([call2]); + stateManager.dequeue(); + + // 3. Queue + const call3 = createValidatingCall('3'); + stateManager.enqueue([call3]); + + const snapshot = stateManager.getSnapshot(); + expect(snapshot).toHaveLength(3); + expect(snapshot[0].request.callId).toBe('1'); + expect(snapshot[1].request.callId).toBe('2'); + expect(snapshot[2].request.callId).toBe('3'); + }); + }); +}); diff --git a/packages/core/src/scheduler/state-manager.ts b/packages/core/src/scheduler/state-manager.ts new file mode 100644 index 0000000000..1817f0771f --- /dev/null +++ b/packages/core/src/scheduler/state-manager.ts @@ -0,0 +1,482 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { + ToolCall, + Status, + WaitingToolCall, + CompletedToolCall, + SuccessfulToolCall, + ErroredToolCall, + CancelledToolCall, + ScheduledToolCall, + ValidatingToolCall, + ExecutingToolCall, + ToolCallResponseInfo, +} from './types.js'; +import type { + ToolConfirmationOutcome, + ToolResultDisplay, + AnyToolInvocation, + ToolCallConfirmationDetails, + AnyDeclarativeTool, +} from '../tools/tools.js'; +import type { MessageBus } from '../confirmation-bus/message-bus.js'; +import { + MessageBusType, + type SerializableConfirmationDetails, +} from '../confirmation-bus/types.js'; + +/** + * Manages the state of tool calls. + * Publishes state changes to the MessageBus via TOOL_CALLS_UPDATE events. + */ +export class SchedulerStateManager { + private readonly activeCalls = new Map(); + private readonly queue: ToolCall[] = []; + private _completedBatch: CompletedToolCall[] = []; + + constructor(private readonly messageBus: MessageBus) {} + + addToolCalls(calls: ToolCall[]): void { + this.enqueue(calls); + } + + getToolCall(callId: string): ToolCall | undefined { + return ( + this.activeCalls.get(callId) || + this.queue.find((c) => c.request.callId === callId) || + this._completedBatch.find((c) => c.request.callId === callId) + ); + } + + enqueue(calls: ToolCall[]): void { + this.queue.push(...calls); + this.emitUpdate(); + } + + dequeue(): ToolCall | undefined { + const next = this.queue.shift(); + if (next) { + this.activeCalls.set(next.request.callId, next); + this.emitUpdate(); + } + return next; + } + + get isActive(): boolean { + return this.activeCalls.size > 0; + } + + get activeCallCount(): number { + return this.activeCalls.size; + } + + get queueLength(): number { + return this.queue.length; + } + + get firstActiveCall(): ToolCall | undefined { + return this.activeCalls.values().next().value; + } + + /** + * Updates the status of a tool call with specific auxiliary data required for certain states. + */ + updateStatus( + callId: string, + status: 'success', + data: ToolCallResponseInfo, + ): void; + updateStatus( + callId: string, + status: 'error', + data: ToolCallResponseInfo, + ): void; + updateStatus( + callId: string, + status: 'awaiting_approval', + data: + | ToolCallConfirmationDetails + | { + correlationId: string; + confirmationDetails: SerializableConfirmationDetails; + }, + ): void; + updateStatus(callId: string, status: 'cancelled', data: string): void; + updateStatus( + callId: string, + status: 'executing', + data?: Partial, + ): void; + updateStatus(callId: string, status: 'scheduled' | 'validating'): void; + updateStatus(callId: string, status: Status, auxiliaryData?: unknown): void { + const call = this.activeCalls.get(callId); + if (!call) return; + + const updatedCall = this.transitionCall(call, status, auxiliaryData); + this.activeCalls.set(callId, updatedCall); + + this.emitUpdate(); + } + + finalizeCall(callId: string): void { + const call = this.activeCalls.get(callId); + if (!call) return; + + if (this.isTerminalCall(call)) { + this._completedBatch.push(call); + this.activeCalls.delete(callId); + } + } + + updateArgs( + callId: string, + newArgs: Record, + newInvocation: AnyToolInvocation, + ): void { + const call = this.activeCalls.get(callId); + if (!call || call.status === 'error') return; + + this.activeCalls.set( + callId, + this.patchCall(call, { + request: { ...call.request, args: newArgs }, + invocation: newInvocation, + }), + ); + this.emitUpdate(); + } + + setOutcome(callId: string, outcome: ToolConfirmationOutcome): void { + const call = this.activeCalls.get(callId); + if (!call) return; + + this.activeCalls.set(callId, this.patchCall(call, { outcome })); + this.emitUpdate(); + } + + cancelAllQueued(reason: string): void { + while (this.queue.length > 0) { + const queuedCall = this.queue.shift()!; + if (queuedCall.status === 'error') { + this._completedBatch.push(queuedCall); + continue; + } + this._completedBatch.push(this.toCancelled(queuedCall, reason)); + } + this.emitUpdate(); + } + + getSnapshot(): ToolCall[] { + return [ + ...this._completedBatch, + ...Array.from(this.activeCalls.values()), + ...this.queue, + ]; + } + + clearBatch(): void { + if (this._completedBatch.length === 0) return; + this._completedBatch = []; + this.emitUpdate(); + } + + get completedBatch(): CompletedToolCall[] { + return [...this._completedBatch]; + } + + private emitUpdate() { + const snapshot = this.getSnapshot(); + + // Fire and forget - The message bus handles the publish and error handling. + void this.messageBus.publish({ + type: MessageBusType.TOOL_CALLS_UPDATE, + toolCalls: snapshot, + }); + } + + private isTerminalCall(call: ToolCall): call is CompletedToolCall { + const { status } = call; + return status === 'success' || status === 'error' || status === 'cancelled'; + } + + private transitionCall( + call: ToolCall, + newStatus: Status, + auxiliaryData?: unknown, + ): ToolCall { + switch (newStatus) { + case 'success': { + if (!this.isToolCallResponseInfo(auxiliaryData)) { + throw new Error( + `Invalid data for 'success' transition (callId: ${call.request.callId})`, + ); + } + return this.toSuccess(call, auxiliaryData); + } + case 'error': { + if (!this.isToolCallResponseInfo(auxiliaryData)) { + throw new Error( + `Invalid data for 'error' transition (callId: ${call.request.callId})`, + ); + } + return this.toError(call, auxiliaryData); + } + case 'awaiting_approval': { + if (!auxiliaryData) { + throw new Error( + `Missing data for 'awaiting_approval' transition (callId: ${call.request.callId})`, + ); + } + return this.toAwaitingApproval(call, auxiliaryData); + } + case 'scheduled': + return this.toScheduled(call); + case 'cancelled': { + if (typeof auxiliaryData !== 'string') { + throw new Error( + `Invalid reason (string) for 'cancelled' transition (callId: ${call.request.callId})`, + ); + } + return this.toCancelled(call, auxiliaryData); + } + case 'validating': + return this.toValidating(call); + case 'executing': { + if ( + auxiliaryData !== undefined && + !this.isExecutingToolCallPatch(auxiliaryData) + ) { + throw new Error( + `Invalid patch for 'executing' transition (callId: ${call.request.callId})`, + ); + } + return this.toExecuting(call, auxiliaryData); + } + default: { + const exhaustiveCheck: never = newStatus; + return exhaustiveCheck; + } + } + } + + private isToolCallResponseInfo(data: unknown): data is ToolCallResponseInfo { + return ( + typeof data === 'object' && + data !== null && + 'callId' in data && + 'responseParts' in data + ); + } + + private isExecutingToolCallPatch( + data: unknown, + ): data is Partial { + // A partial can be an empty object, but it must be a non-null object. + return typeof data === 'object' && data !== null; + } + + // --- Transition Helpers --- + + /** + * Ensures the tool call has an associated tool and invocation before + * transitioning to states that require them. + */ + private validateHasToolAndInvocation( + call: ToolCall, + targetStatus: Status, + ): asserts call is ToolCall & { + tool: AnyDeclarativeTool; + invocation: AnyToolInvocation; + } { + if ( + !('tool' in call && call.tool && 'invocation' in call && call.invocation) + ) { + throw new Error( + `Invalid state transition: cannot transition to ${targetStatus} without tool/invocation (callId: ${call.request.callId})`, + ); + } + } + + private toSuccess( + call: ToolCall, + response: ToolCallResponseInfo, + ): SuccessfulToolCall { + this.validateHasToolAndInvocation(call, 'success'); + const startTime = 'startTime' in call ? call.startTime : undefined; + return { + request: call.request, + tool: call.tool, + invocation: call.invocation, + status: 'success', + response, + durationMs: startTime ? Date.now() - startTime : undefined, + outcome: call.outcome, + }; + } + + private toError( + call: ToolCall, + response: ToolCallResponseInfo, + ): ErroredToolCall { + const startTime = 'startTime' in call ? call.startTime : undefined; + return { + request: call.request, + status: 'error', + tool: 'tool' in call ? call.tool : undefined, + response, + durationMs: startTime ? Date.now() - startTime : undefined, + outcome: call.outcome, + }; + } + + private toAwaitingApproval(call: ToolCall, data: unknown): WaitingToolCall { + this.validateHasToolAndInvocation(call, 'awaiting_approval'); + + let confirmationDetails: + | ToolCallConfirmationDetails + | SerializableConfirmationDetails; + let correlationId: string | undefined; + + if (this.isEventDrivenApprovalData(data)) { + correlationId = data.correlationId; + confirmationDetails = data.confirmationDetails; + } else { + // TODO: Remove legacy callback shape once event-driven migration is complete + confirmationDetails = data as ToolCallConfirmationDetails; + } + + return { + request: call.request, + tool: call.tool, + status: 'awaiting_approval', + correlationId, + confirmationDetails, + startTime: 'startTime' in call ? call.startTime : undefined, + outcome: call.outcome, + invocation: call.invocation, + }; + } + + private isEventDrivenApprovalData(data: unknown): data is { + correlationId: string; + confirmationDetails: SerializableConfirmationDetails; + } { + return ( + typeof data === 'object' && + data !== null && + 'correlationId' in data && + 'confirmationDetails' in data + ); + } + + private toScheduled(call: ToolCall): ScheduledToolCall { + this.validateHasToolAndInvocation(call, 'scheduled'); + return { + request: call.request, + tool: call.tool, + status: 'scheduled', + startTime: 'startTime' in call ? call.startTime : undefined, + outcome: call.outcome, + invocation: call.invocation, + }; + } + + private toCancelled(call: ToolCall, reason: string): CancelledToolCall { + this.validateHasToolAndInvocation(call, 'cancelled'); + const startTime = 'startTime' in call ? call.startTime : undefined; + + // TODO: Refactor this tool-specific logic into the confirmation details payload. + // See: https://github.com/google-gemini/gemini-cli/issues/16716 + let resultDisplay: ToolResultDisplay | undefined = undefined; + if (this.isWaitingToolCall(call)) { + const details = call.confirmationDetails; + if ( + details.type === 'edit' && + 'fileDiff' in details && + 'fileName' in details && + 'filePath' in details && + 'originalContent' in details && + 'newContent' in details + ) { + resultDisplay = { + fileDiff: details.fileDiff, + fileName: details.fileName, + filePath: details.filePath, + originalContent: details.originalContent, + newContent: details.newContent, + }; + } + } + + const errorMessage = `[Operation Cancelled] Reason: ${reason}`; + return { + request: call.request, + tool: call.tool, + invocation: call.invocation, + status: 'cancelled', + response: { + callId: call.request.callId, + responseParts: [ + { + functionResponse: { + id: call.request.callId, + name: call.request.name, + response: { error: errorMessage }, + }, + }, + ], + resultDisplay, + error: undefined, + errorType: undefined, + contentLength: errorMessage.length, + }, + durationMs: startTime ? Date.now() - startTime : undefined, + outcome: call.outcome, + }; + } + + private isWaitingToolCall(call: ToolCall): call is WaitingToolCall { + return call.status === 'awaiting_approval'; + } + + private patchCall(call: T, patch: Partial): T { + return { ...call, ...patch }; + } + + private toValidating(call: ToolCall): ValidatingToolCall { + this.validateHasToolAndInvocation(call, 'validating'); + return { + request: call.request, + tool: call.tool, + status: 'validating', + startTime: 'startTime' in call ? call.startTime : undefined, + outcome: call.outcome, + invocation: call.invocation, + }; + } + + private toExecuting(call: ToolCall, data?: unknown): ExecutingToolCall { + this.validateHasToolAndInvocation(call, 'executing'); + const execData = data as Partial | undefined; + const liveOutput = + execData?.liveOutput ?? + ('liveOutput' in call ? call.liveOutput : undefined); + const pid = execData?.pid ?? ('pid' in call ? call.pid : undefined); + + return { + request: call.request, + tool: call.tool, + status: 'executing', + startTime: 'startTime' in call ? call.startTime : undefined, + outcome: call.outcome, + invocation: call.invocation, + liveOutput, + pid, + }; + } +}