/**
 * @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 { CoreToolCallStatus, ROOT_SCHEDULER_ID, } from './types.js';
import { ToolConfirmationOutcome, } from '../tools/tools.js';
import { MessageBusType } from '../confirmation-bus/types.js';
import { ApprovalMode } from '../policy/types.js';
describe('SchedulerStateManager', () => {
    const mockRequest = {
        callId: 'call-1',
        name: 'test-tool',
        args: { foo: 'bar' },
        isClientInitiated: false,
        prompt_id: 'prompt-1',
    };
    const mockTool = {
        name: 'test-tool',
        displayName: 'Test Tool',
    };
    const mockInvocation = {
        shouldConfirmExecute: vi.fn(),
    };
    const createValidatingCall = (id = 'call-1', mode = ApprovalMode.DEFAULT) => ({
        status: CoreToolCallStatus.Validating,
        request: { ...mockRequest, callId: id },
        tool: mockTool,
        invocation: mockInvocation,
        startTime: Date.now(),
        approvalMode: mode,
    });
    const createMockResponse = (id) => ({
        callId: id,
        responseParts: [],
        resultDisplay: 'Success',
        error: undefined,
        errorType: undefined,
    });
    let stateManager;
    let mockMessageBus;
    let onUpdate;
    beforeEach(() => {
        onUpdate = vi.fn();
        mockMessageBus = {
            publish: vi.fn(),
            subscribe: vi.fn(),
            unsubscribe: vi.fn(),
        };
        // Capture the update when published
        vi.mocked(mockMessageBus.publish).mockImplementation((msg) => {
            // Return a Promise to satisfy the void | Promise<void> 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('Observer Callback', () => {
        it('should trigger onTerminalCall when finalizing a call', () => {
            const onTerminalCall = vi.fn();
            const manager = new SchedulerStateManager(mockMessageBus, ROOT_SCHEDULER_ID, onTerminalCall);
            const call = createValidatingCall();
            manager.enqueue([call]);
            manager.dequeue();
            manager.updateStatus(call.request.callId, CoreToolCallStatus.Success, createMockResponse(call.request.callId));
            manager.finalizeCall(call.request.callId);
            expect(onTerminalCall).toHaveBeenCalledTimes(1);
            expect(onTerminalCall).toHaveBeenCalledWith(expect.objectContaining({
                status: CoreToolCallStatus.Success,
                request: expect.objectContaining({ callId: call.request.callId }),
            }));
        });
        it('should trigger onTerminalCall for every call in cancelAllQueued', () => {
            const onTerminalCall = vi.fn();
            const manager = new SchedulerStateManager(mockMessageBus, ROOT_SCHEDULER_ID, onTerminalCall);
            manager.enqueue([createValidatingCall('1'), createValidatingCall('2')]);
            manager.cancelAllQueued('Test cancel');
            expect(onTerminalCall).toHaveBeenCalledTimes(2);
            expect(onTerminalCall).toHaveBeenCalledWith(expect.objectContaining({
                status: CoreToolCallStatus.Cancelled,
                request: expect.objectContaining({ callId: '1' }),
            }));
            expect(onTerminalCall).toHaveBeenCalledWith(expect.objectContaining({
                status: CoreToolCallStatus.Cancelled,
                request: expect.objectContaining({ callId: '2' }),
            }));
        });
    });
    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', CoreToolCallStatus.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('call-1', ApprovalMode.PLAN);
            stateManager.enqueue([call]);
            stateManager.dequeue();
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.Scheduled);
            const snapshot = stateManager.getSnapshot();
            expect(snapshot[0].status).toBe(CoreToolCallStatus.Scheduled);
            expect(snapshot[0].request.callId).toBe(call.request.callId);
            expect(snapshot[0].approvalMode).toBe(ApprovalMode.PLAN);
        });
        it('should transition scheduled to executing', () => {
            const call = createValidatingCall();
            stateManager.enqueue([call]);
            stateManager.dequeue();
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.Scheduled);
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.Executing);
            expect(stateManager.firstActiveCall?.status).toBe(CoreToolCallStatus.Executing);
        });
        it('should transition to success and move to completed batch', () => {
            const call = createValidatingCall('call-1', ApprovalMode.PLAN);
            stateManager.enqueue([call]);
            stateManager.dequeue();
            const response = {
                callId: call.request.callId,
                responseParts: [],
                resultDisplay: 'Success',
                error: undefined,
                errorType: undefined,
            };
            vi.mocked(onUpdate).mockClear();
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.Success, response);
            expect(onUpdate).toHaveBeenCalledTimes(1);
            vi.mocked(onUpdate).mockClear();
            stateManager.finalizeCall(call.request.callId);
            expect(onUpdate).toHaveBeenCalledTimes(1);
            expect(stateManager.isActive).toBe(false);
            expect(stateManager.completedBatch).toHaveLength(1);
            const completed = stateManager.completedBatch[0];
            expect(completed.status).toBe(CoreToolCallStatus.Success);
            expect(completed.response).toEqual(response);
            expect(completed.durationMs).toBeDefined();
            expect(completed.approvalMode).toBe(ApprovalMode.PLAN);
        });
        it('should transition to error and move to completed batch', () => {
            const call = createValidatingCall();
            stateManager.enqueue([call]);
            stateManager.dequeue();
            const response = {
                callId: call.request.callId,
                responseParts: [],
                resultDisplay: 'Error',
                error: new Error('Failed'),
                errorType: undefined,
            };
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.Error, response);
            stateManager.finalizeCall(call.request.callId);
            expect(stateManager.isActive).toBe(false);
            expect(stateManager.completedBatch).toHaveLength(1);
            const completed = stateManager.completedBatch[0];
            expect(completed.status).toBe(CoreToolCallStatus.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',
                title: 'Confirm',
                prompt: 'Proceed?',
                onConfirm: vi.fn(),
            };
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.AwaitingApproval, details);
            const active = stateManager.firstActiveCall;
            expect(active.status).toBe(CoreToolCallStatus.AwaitingApproval);
            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',
                title: 'Confirm',
                prompt: 'Proceed?',
            };
            const eventDrivenData = {
                correlationId: 'corr-123',
                confirmationDetails: details,
            };
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.AwaitingApproval, eventDrivenData);
            const active = stateManager.firstActiveCall;
            expect(active.status).toBe(CoreToolCallStatus.AwaitingApproval);
            expect(active.correlationId).toBe('corr-123');
            expect(active.confirmationDetails).toEqual(details);
        });
        it('should preserve diff and derive stats when cancelling an edit tool call', () => {
            const call = createValidatingCall();
            stateManager.enqueue([call]);
            stateManager.dequeue();
            const details = {
                type: 'edit',
                title: 'Edit',
                fileName: 'test.txt',
                filePath: '/path/to/test.txt',
                fileDiff: '@@ -1,1 +1,1 @@\n-old line\n+new line',
                originalContent: 'old line',
                newContent: 'new line',
                onConfirm: vi.fn(),
            };
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.AwaitingApproval, details);
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.Cancelled, 'User said no');
            stateManager.finalizeCall(call.request.callId);
            const completed = stateManager.completedBatch[0];
            expect(completed.status).toBe(CoreToolCallStatus.Cancelled);
            const result = completed.response.resultDisplay;
            expect(result.fileDiff).toBe(details.fileDiff);
            expect(result.diffStat).toEqual(expect.objectContaining({
                model_added_lines: 1,
                model_removed_lines: 1,
            }));
        });
        it('should ignore status updates for non-existent callIds', () => {
            stateManager.updateStatus('unknown', CoreToolCallStatus.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, CoreToolCallStatus.Success, createMockResponse(call.request.callId));
            stateManager.finalizeCall(call.request.callId);
            vi.mocked(onUpdate).mockClear();
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.Scheduled);
            expect(onUpdate).not.toHaveBeenCalled();
        });
        it('should only finalize terminal calls', () => {
            const call = createValidatingCall();
            stateManager.enqueue([call]);
            stateManager.dequeue();
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.Executing);
            stateManager.finalizeCall(call.request.callId);
            expect(stateManager.isActive).toBe(true);
            expect(stateManager.completedBatch).toHaveLength(0);
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.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, CoreToolCallStatus.Executing);
            let active = stateManager.firstActiveCall;
            expect(active.status).toBe(CoreToolCallStatus.Executing);
            expect(active.liveOutput).toBeUndefined();
            // Update with live output
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.Executing, {
                liveOutput: 'chunk 1',
            });
            active = stateManager.firstActiveCall;
            expect(active.liveOutput).toBe('chunk 1');
            // Update with pid (should preserve liveOutput)
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.Executing, {
                pid: 1234,
            });
            active = stateManager.firstActiveCall;
            expect(active.liveOutput).toBe('chunk 1');
            expect(active.pid).toBe(1234);
            // Update live output again (should preserve pid)
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.Executing, {
                liveOutput: 'chunk 2',
            });
            active = stateManager.firstActiveCall;
            expect(active.liveOutput).toBe('chunk 2');
            expect(active.pid).toBe(1234);
        });
        it('should update progressMessage and progressPercent during executing updates', () => {
            const call = createValidatingCall();
            stateManager.enqueue([call]);
            stateManager.dequeue();
            // Update with progress
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.Executing, {
                progressMessage: 'Starting...',
                progressPercent: 10,
            });
            let active = stateManager.firstActiveCall;
            expect(active.progressMessage).toBe('Starting...');
            expect(active.progressPercent).toBe(10);
            // Update progress further
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.Executing, {
                progressMessage: 'Halfway!',
                progressPercent: 50,
            });
            active = stateManager.firstActiveCall;
            expect(active.progressMessage).toBe('Halfway!');
            expect(active.progressPercent).toBe(50);
        });
    });
    describe('Argument Updates', () => {
        it('should update args and invocation', () => {
            const call = createValidatingCall();
            stateManager.enqueue([call]);
            stateManager.dequeue();
            const newArgs = { foo: 'updated' };
            const newInvocation = { ...mockInvocation };
            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, CoreToolCallStatus.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'),
            ]);
            vi.mocked(onUpdate).mockClear();
            stateManager.cancelAllQueued('Batch cancel');
            expect(stateManager.queueLength).toBe(0);
            expect(stateManager.completedBatch).toHaveLength(2);
            expect(stateManager.completedBatch.every((c) => c.status === CoreToolCallStatus.Cancelled)).toBe(true);
            expect(onUpdate).toHaveBeenCalledTimes(1);
        });
        it('should not notify if cancelAllQueued is called on an empty queue', () => {
            vi.mocked(onUpdate).mockClear();
            stateManager.cancelAllQueued('Batch cancel');
            expect(onUpdate).not.toHaveBeenCalled();
        });
        it('should clear batch and notify', () => {
            const call = createValidatingCall();
            stateManager.enqueue([call]);
            stateManager.dequeue();
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.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, CoreToolCallStatus.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', CoreToolCallStatus.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');
        });
    });
    describe('progress field preservation', () => {
        it('should preserve progress and progressTotal in toExecuting', () => {
            const call = createValidatingCall('progress-1');
            stateManager.enqueue([call]);
            stateManager.dequeue();
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.Executing, {
                progress: 5,
                progressTotal: 10,
                progressMessage: 'Working',
                progressPercent: 50,
            });
            const active = stateManager.firstActiveCall;
            expect(active.status).toBe(CoreToolCallStatus.Executing);
            expect(active.progress).toBe(5);
            expect(active.progressTotal).toBe(10);
            expect(active.progressMessage).toBe('Working');
            expect(active.progressPercent).toBe(50);
        });
        it('should preserve progress fields after a liveOutput update', () => {
            const call = createValidatingCall('progress-2');
            stateManager.enqueue([call]);
            stateManager.dequeue();
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.Executing, {
                progress: 5,
                progressTotal: 10,
                progressMessage: 'Working',
                progressPercent: 50,
            });
            stateManager.updateStatus(call.request.callId, CoreToolCallStatus.Executing, {
                liveOutput: 'some output',
            });
            const active = stateManager.firstActiveCall;
            expect(active.status).toBe(CoreToolCallStatus.Executing);
            expect(active.liveOutput).toBe('some output');
            expect(active.progress).toBe(5);
            expect(active.progressTotal).toBe(10);
            expect(active.progressMessage).toBe('Working');
            expect(active.progressPercent).toBe(50);
        });
    });
});
//# sourceMappingURL=state-manager.test.js.map