qwen-code/packages/cli/src/acp-integration/session/SubAgentTracker.test.ts

/**
 * @license
 * Copyright 2025 Qwen
 * SPDX-License-Identifier: Apache-2.0
 */

import { describe, it, expect, vi, beforeEach } from 'vitest';
import { SubAgentTracker } from './SubAgentTracker.js';
import type { SessionContext } from './types.js';
import type {
  Config,
  ToolRegistry,
  AgentEventEmitter,
  AgentToolCallEvent,
  AgentToolResultEvent,
  AgentApprovalRequestEvent,
  AgentStreamTextEvent,
  ToolEditConfirmationDetails,
  ToolInfoConfirmationDetails,
} from '@qwen-code/qwen-code-core';
import {
  AgentEventType,
  ToolConfirmationOutcome,
  ToolNames,
} from '@qwen-code/qwen-code-core';
import type { AgentSideConnection } from '@agentclientprotocol/sdk';
import { EventEmitter } from 'node:events';

// Helper to create a mock AgentToolCallEvent with required fields
function createToolCallEvent(
  overrides: Partial<AgentToolCallEvent> & { name: string; callId: string },
): AgentToolCallEvent {
  return {
    subagentId: 'test-subagent',
    round: 1,
    timestamp: Date.now(),
    description: `Calling ${overrides.name}`,
    args: {},
    ...overrides,
  };
}

// Helper to create a mock AgentToolResultEvent with required fields
function createToolResultEvent(
  overrides: Partial<AgentToolResultEvent> & {
    name: string;
    callId: string;
    success: boolean;
  },
): AgentToolResultEvent {
  return {
    subagentId: 'test-subagent',
    round: 1,
    timestamp: Date.now(),
    ...overrides,
  };
}

// Helper to create a mock AgentApprovalRequestEvent with required fields
function createApprovalEvent(
  overrides: Partial<AgentApprovalRequestEvent> & {
    name: string;
    callId: string;
    confirmationDetails: AgentApprovalRequestEvent['confirmationDetails'];
    respond: AgentApprovalRequestEvent['respond'];
  },
): AgentApprovalRequestEvent {
  return {
    subagentId: 'test-subagent',
    round: 1,
    timestamp: Date.now(),
    description: `Awaiting approval for ${overrides.name}`,
    ...overrides,
  };
}

// Helper to create edit confirmation details
function createEditConfirmation(
  overrides: Partial<Omit<ToolEditConfirmationDetails, 'onConfirm' | 'type'>>,
): Omit<ToolEditConfirmationDetails, 'onConfirm'> {
  return {
    type: 'edit',
    title: 'Edit file',
    fileName: '/test.ts',
    filePath: '/test.ts',
    fileDiff: '',
    originalContent: '',
    newContent: '',
    ...overrides,
  };
}

// Helper to create info confirmation details
function createInfoConfirmation(
  overrides?: Partial<Omit<ToolInfoConfirmationDetails, 'onConfirm' | 'type'>>,
): Omit<ToolInfoConfirmationDetails, 'onConfirm'> {
  return {
    type: 'info',
    title: 'Tool requires approval',
    prompt: 'Allow this action?',
    ...overrides,
  };
}

// Helper to create a mock AgentStreamTextEvent with required fields
function createStreamTextEvent(
  overrides: Partial<AgentStreamTextEvent> & { text: string },
): AgentStreamTextEvent {
  return {
    subagentId: 'test-subagent',
    round: 1,
    timestamp: Date.now(),
    ...overrides,
  };
}

describe('SubAgentTracker', () => {
  let mockContext: SessionContext;
  let mockClient: AgentSideConnection;
  let sendUpdateSpy: ReturnType<typeof vi.fn>;
  let requestPermissionSpy: ReturnType<typeof vi.fn>;
  let tracker: SubAgentTracker;
  let eventEmitter: AgentEventEmitter;
  let abortController: AbortController;

  beforeEach(() => {
    sendUpdateSpy = vi.fn().mockResolvedValue(undefined);
    requestPermissionSpy = vi.fn().mockResolvedValue({
      outcome: { optionId: ToolConfirmationOutcome.ProceedOnce },
    });

    const mockToolRegistry = {
      getTool: vi.fn().mockReturnValue(null),
    } as unknown as ToolRegistry;

    mockContext = {
      sessionId: 'test-session-id',
      config: {
        getToolRegistry: () => mockToolRegistry,
      } as unknown as Config,
      sendUpdate: sendUpdateSpy,
    };

    mockClient = {
      requestPermission: requestPermissionSpy,
    } as unknown as AgentSideConnection;

    tracker = new SubAgentTracker(
      mockContext,
      mockClient,
      'parent-call-123',
      'test-subagent',
    );
    eventEmitter = new EventEmitter() as unknown as AgentEventEmitter;
    abortController = new AbortController();
  });

  describe('setup', () => {
    it('should return cleanup function', () => {
      const cleanups = tracker.setup(eventEmitter, abortController.signal);

      expect(cleanups).toHaveLength(1);
      expect(typeof cleanups[0]).toBe('function');
    });

    it('should register event listeners', () => {
      const onSpy = vi.spyOn(eventEmitter, 'on');

      tracker.setup(eventEmitter, abortController.signal);

      expect(onSpy).toHaveBeenCalledWith(
        AgentEventType.TOOL_CALL,
        expect.any(Function),
      );
      expect(onSpy).toHaveBeenCalledWith(
        AgentEventType.TOOL_RESULT,
        expect.any(Function),
      );
      expect(onSpy).toHaveBeenCalledWith(
        AgentEventType.TOOL_WAITING_APPROVAL,
        expect.any(Function),
      );
      expect(onSpy).toHaveBeenCalledWith(
        AgentEventType.STREAM_TEXT,
        expect.any(Function),
      );
    });

    it('should remove event listeners on cleanup', () => {
      const offSpy = vi.spyOn(eventEmitter, 'off');
      const cleanups = tracker.setup(eventEmitter, abortController.signal);

      cleanups[0]();

      expect(offSpy).toHaveBeenCalledWith(
        AgentEventType.TOOL_CALL,
        expect.any(Function),
      );
      expect(offSpy).toHaveBeenCalledWith(
        AgentEventType.TOOL_RESULT,
        expect.any(Function),
      );
      expect(offSpy).toHaveBeenCalledWith(
        AgentEventType.TOOL_WAITING_APPROVAL,
        expect.any(Function),
      );
      expect(offSpy).toHaveBeenCalledWith(
        AgentEventType.STREAM_TEXT,
        expect.any(Function),
      );
    });
  });

  describe('tool call handling', () => {
    it('should emit tool_call on TOOL_CALL event', async () => {
      tracker.setup(eventEmitter, abortController.signal);

      const event = createToolCallEvent({
        name: 'read_file',
        callId: 'call-123',
        args: { path: '/test.ts' },
        description: 'Reading file',
      });

      eventEmitter.emit(AgentEventType.TOOL_CALL, event);

      // Allow async operations to complete
      await vi.waitFor(() => {
        expect(sendUpdateSpy).toHaveBeenCalled();
      });

      // ToolCallEmitter resolves metadata from registry - uses toolName when tool not found
      expect(sendUpdateSpy).toHaveBeenCalledWith(
        expect.objectContaining({
          sessionUpdate: 'tool_call',
          toolCallId: 'call-123',
          status: 'pending',
          title: 'read_file',
          content: [],
          locations: [],
          kind: 'other',
          rawInput: { path: '/test.ts' },
          _meta: expect.objectContaining({
            toolName: 'read_file',
            parentToolCallId: 'parent-call-123',
            subagentType: 'test-subagent',
          }),
        }),
      );
    });

    it('should skip tool_call for TodoWriteTool', async () => {
      tracker.setup(eventEmitter, abortController.signal);

      const event = createToolCallEvent({
        name: ToolNames.TODO_WRITE,
        callId: 'call-todo',
        args: { todos: [] },
      });

      eventEmitter.emit(AgentEventType.TOOL_CALL, event);

      // Give time for any async operation
      await new Promise((resolve) => setTimeout(resolve, 10));

      expect(sendUpdateSpy).not.toHaveBeenCalled();
    });

    it('should not emit when aborted', async () => {
      tracker.setup(eventEmitter, abortController.signal);
      abortController.abort();

      const event = createToolCallEvent({
        name: 'read_file',
        callId: 'call-123',
        args: {},
      });

      eventEmitter.emit(AgentEventType.TOOL_CALL, event);

      await new Promise((resolve) => setTimeout(resolve, 10));

      expect(sendUpdateSpy).not.toHaveBeenCalled();
    });
  });

  describe('tool result handling', () => {
    it('should emit tool_call_update on TOOL_RESULT event', async () => {
      tracker.setup(eventEmitter, abortController.signal);

      // First emit tool call to store state
      eventEmitter.emit(
        AgentEventType.TOOL_CALL,
        createToolCallEvent({
          name: 'read_file',
          callId: 'call-123',
          args: { path: '/test.ts' },
        }),
      );

      // Then emit result
      const resultEvent = createToolResultEvent({
        name: 'read_file',
        callId: 'call-123',
        success: true,
        resultDisplay: 'File contents',
      });

      eventEmitter.emit(AgentEventType.TOOL_RESULT, resultEvent);

      await vi.waitFor(() => {
        expect(sendUpdateSpy).toHaveBeenCalledWith(
          expect.objectContaining({
            sessionUpdate: 'tool_call_update',
            toolCallId: 'call-123',
            status: 'completed',
            _meta: expect.objectContaining({
              toolName: 'read_file',
              parentToolCallId: 'parent-call-123',
              subagentType: 'test-subagent',
            }),
          }),
        );
      });
    });

    it('should emit failed status on unsuccessful result', async () => {
      tracker.setup(eventEmitter, abortController.signal);

      const resultEvent = createToolResultEvent({
        name: 'read_file',
        callId: 'call-fail',
        success: false,
        resultDisplay: undefined,
      });

      eventEmitter.emit(AgentEventType.TOOL_RESULT, resultEvent);

      await vi.waitFor(() => {
        expect(sendUpdateSpy).toHaveBeenCalledWith(
          expect.objectContaining({
            sessionUpdate: 'tool_call_update',
            status: 'failed',
            _meta: expect.objectContaining({
              toolName: 'read_file',
              parentToolCallId: 'parent-call-123',
              subagentType: 'test-subagent',
            }),
          }),
        );
      });
    });

    it('should emit plan update for TodoWriteTool results', async () => {
      tracker.setup(eventEmitter, abortController.signal);

      // Store args via tool call
      eventEmitter.emit(
        AgentEventType.TOOL_CALL,
        createToolCallEvent({
          name: ToolNames.TODO_WRITE,
          callId: 'call-todo',
          args: {
            todos: [{ id: '1', content: 'Task 1', status: 'pending' }],
          },
        }),
      );

      // Emit result with todo_list display
      const resultEvent = createToolResultEvent({
        name: ToolNames.TODO_WRITE,
        callId: 'call-todo',
        success: true,
        resultDisplay: JSON.stringify({
          type: 'todo_list',
          todos: [{ id: '1', content: 'Task 1', status: 'completed' }],
        }),
      });

      eventEmitter.emit(AgentEventType.TOOL_RESULT, resultEvent);

      await vi.waitFor(() => {
        expect(sendUpdateSpy).toHaveBeenCalledWith({
          sessionUpdate: 'plan',
          entries: [
            { content: 'Task 1', priority: 'medium', status: 'completed' },
          ],
        });
      });
    });

    it('should clean up state after result', async () => {
      tracker.setup(eventEmitter, abortController.signal);

      eventEmitter.emit(
        AgentEventType.TOOL_CALL,
        createToolCallEvent({
          name: 'test_tool',
          callId: 'call-cleanup',
          args: { test: true },
        }),
      );

      eventEmitter.emit(
        AgentEventType.TOOL_RESULT,
        createToolResultEvent({
          name: 'test_tool',
          callId: 'call-cleanup',
          success: true,
        }),
      );

      // Emit another result for same callId - should not have stored args
      sendUpdateSpy.mockClear();
      eventEmitter.emit(
        AgentEventType.TOOL_RESULT,
        createToolResultEvent({
          name: 'test_tool',
          callId: 'call-cleanup',
          success: true,
        }),
      );

      await vi.waitFor(() => {
        expect(sendUpdateSpy).toHaveBeenCalled();
      });

      // Second call should not have args from first call
      // (state was cleaned up)
    });
  });

  describe('approval handling', () => {
    it('should request permission from client', async () => {
      tracker.setup(eventEmitter, abortController.signal);

      const respondSpy = vi.fn().mockResolvedValue(undefined);
      const event = createApprovalEvent({
        name: 'edit_file',
        callId: 'call-edit',
        description: 'Editing file',
        confirmationDetails: createEditConfirmation({
          fileName: '/test.ts',
          originalContent: 'old',
          newContent: 'new',
        }),
        respond: respondSpy,
      });

      eventEmitter.emit(AgentEventType.TOOL_WAITING_APPROVAL, event);

      await vi.waitFor(() => {
        expect(requestPermissionSpy).toHaveBeenCalled();
      });

      expect(requestPermissionSpy).toHaveBeenCalledWith(
        expect.objectContaining({
          sessionId: 'test-session-id',
          toolCall: expect.objectContaining({
            toolCallId: 'call-edit',
            status: 'pending',
            content: [
              {
                type: 'diff',
                path: '/test.ts',
                oldText: 'old',
                newText: 'new',
              },
            ],
          }),
        }),
      );
    });

    it('should respond to subagent with permission outcome', async () => {
      tracker.setup(eventEmitter, abortController.signal);

      const respondSpy = vi.fn().mockResolvedValue(undefined);
      const event = createApprovalEvent({
        name: 'test_tool',
        callId: 'call-123',
        confirmationDetails: createInfoConfirmation(),
        respond: respondSpy,
      });

      eventEmitter.emit(AgentEventType.TOOL_WAITING_APPROVAL, event);

      await vi.waitFor(() => {
        expect(respondSpy).toHaveBeenCalledWith(
          ToolConfirmationOutcome.ProceedOnce,
          {
            answers: undefined,
          },
        );
      });
    });

    it('should cancel on permission request failure', async () => {
      requestPermissionSpy.mockRejectedValue(new Error('Network error'));
      tracker.setup(eventEmitter, abortController.signal);

      const respondSpy = vi.fn().mockResolvedValue(undefined);
      const event = createApprovalEvent({
        name: 'test_tool',
        callId: 'call-123',
        confirmationDetails: createInfoConfirmation(),
        respond: respondSpy,
      });

      eventEmitter.emit(AgentEventType.TOOL_WAITING_APPROVAL, event);

      await vi.waitFor(() => {
        expect(respondSpy).toHaveBeenCalledWith(ToolConfirmationOutcome.Cancel);
      });
    });

    it('should handle cancelled outcome from client', async () => {
      requestPermissionSpy.mockResolvedValue({
        outcome: { outcome: 'cancelled' },
      });
      tracker.setup(eventEmitter, abortController.signal);

      const respondSpy = vi.fn().mockResolvedValue(undefined);
      const event = createApprovalEvent({
        name: 'test_tool',
        callId: 'call-123',
        confirmationDetails: createInfoConfirmation(),
        respond: respondSpy,
      });

      eventEmitter.emit(AgentEventType.TOOL_WAITING_APPROVAL, event);

      await vi.waitFor(() => {
        expect(respondSpy).toHaveBeenCalledWith(
          ToolConfirmationOutcome.Cancel,
          {
            answers: undefined,
          },
        );
      });
    });

    it('should forward answers payload from ACP permission responses', async () => {
      requestPermissionSpy.mockResolvedValue({
        outcome: {
          outcome: 'selected',
          optionId: ToolConfirmationOutcome.ProceedOnce,
        },
        answers: {
          answer: 'yes',
        },
      });
      tracker.setup(eventEmitter, abortController.signal);

      const respondSpy = vi.fn().mockResolvedValue(undefined);
      const confirmationDetails = {
        type: 'ask_user_question',
        title: 'Question',
        questions: [
          {
            question: 'Continue?',
            header: 'Question',
            options: [],
            multiSelect: false,
          },
        ],
      } as unknown as AgentApprovalRequestEvent['confirmationDetails'];
      const event = createApprovalEvent({
        name: 'ask_user_question',
        callId: 'call-ask',
        confirmationDetails,
        respond: respondSpy,
      });

      eventEmitter.emit(AgentEventType.TOOL_WAITING_APPROVAL, event);

      await vi.waitFor(() => {
        expect(respondSpy).toHaveBeenCalledWith(
          ToolConfirmationOutcome.ProceedOnce,
          {
            answers: {
              answer: 'yes',
            },
          },
        );
      });
    });

    it('should use filePath over fileName for diff content path', async () => {
      tracker.setup(eventEmitter, abortController.signal);

      const respondSpy = vi.fn().mockResolvedValue(undefined);
      const event = createApprovalEvent({
        name: 'edit_file',
        callId: 'call-path-test',
        description: 'Editing file',
        confirmationDetails: createEditConfirmation({
          fileName: 'test.ts',
          filePath: '/workspace/src/test.ts',
          originalContent: 'old content',
          newContent: 'new content',
        }),
        respond: respondSpy,
      });

      eventEmitter.emit(AgentEventType.TOOL_WAITING_APPROVAL, event);

      await vi.waitFor(() => {
        expect(requestPermissionSpy).toHaveBeenCalled();
      });

      expect(requestPermissionSpy).toHaveBeenCalledWith(
        expect.objectContaining({
          toolCall: expect.objectContaining({
            content: [
              {
                type: 'diff',
                path: '/workspace/src/test.ts',
                oldText: 'old content',
                newText: 'new content',
              },
            ],
          }),
        }),
      );
    });

    it('should fall back to fileName when filePath is not available', async () => {
      tracker.setup(eventEmitter, abortController.signal);

      const respondSpy = vi.fn().mockResolvedValue(undefined);
      const event = createApprovalEvent({
        name: 'edit_file',
        callId: 'call-fallback-test',
        description: 'Editing file',
        confirmationDetails: {
          type: 'edit' as const,
          title: 'Edit file',
          fileName: 'fallback.ts',
          fileDiff: '',
          originalContent: 'old',
          newContent: 'new',
        } as Omit<ToolEditConfirmationDetails, 'onConfirm'>,
        respond: respondSpy,
      });

      eventEmitter.emit(AgentEventType.TOOL_WAITING_APPROVAL, event);

      await vi.waitFor(() => {
        expect(requestPermissionSpy).toHaveBeenCalled();
      });

      expect(requestPermissionSpy).toHaveBeenCalledWith(
        expect.objectContaining({
          toolCall: expect.objectContaining({
            content: [
              {
                type: 'diff',
                path: 'fallback.ts',
                oldText: 'old',
                newText: 'new',
              },
            ],
          }),
        }),
      );
    });
  });

  describe('permission options', () => {
    it('should include "Allow All Edits" for edit type', async () => {
      tracker.setup(eventEmitter, abortController.signal);

      const event = createApprovalEvent({
        name: 'edit_file',
        callId: 'call-123',
        confirmationDetails: createEditConfirmation({
          fileName: '/test.ts',
          originalContent: '',
          newContent: 'new',
        }),
        respond: vi.fn(),
      });

      eventEmitter.emit(AgentEventType.TOOL_WAITING_APPROVAL, event);

      await vi.waitFor(() => {
        expect(requestPermissionSpy).toHaveBeenCalled();
      });

      const call = requestPermissionSpy.mock.calls[0][0];
      expect(call.options).toContainEqual(
        expect.objectContaining({
          optionId: ToolConfirmationOutcome.ProceedAlways,
          name: 'Allow All Edits',
        }),
      );
    });
  });

  describe('stream text handling', () => {
    it('should emit agent_message_chunk on STREAM_TEXT event', async () => {
      tracker.setup(eventEmitter, abortController.signal);

      const event = createStreamTextEvent({
        text: 'Hello, this is a response from the model.',
      });

      eventEmitter.emit(AgentEventType.STREAM_TEXT, event);

      await vi.waitFor(() => {
        expect(sendUpdateSpy).toHaveBeenCalled();
      });

      expect(sendUpdateSpy).toHaveBeenCalledWith(
        expect.objectContaining({
          sessionUpdate: 'agent_message_chunk',
          content: {
            type: 'text',
            text: 'Hello, this is a response from the model.',
          },
        }),
      );
    });

    it('should emit multiple chunks for multiple STREAM_TEXT events', async () => {
      tracker.setup(eventEmitter, abortController.signal);

      eventEmitter.emit(
        AgentEventType.STREAM_TEXT,
        createStreamTextEvent({ text: 'First chunk ' }),
      );
      eventEmitter.emit(
        AgentEventType.STREAM_TEXT,
        createStreamTextEvent({ text: 'Second chunk ' }),
      );
      eventEmitter.emit(
        AgentEventType.STREAM_TEXT,
        createStreamTextEvent({ text: 'Third chunk' }),
      );

      await vi.waitFor(() => {
        expect(sendUpdateSpy).toHaveBeenCalledTimes(3);
      });

      expect(sendUpdateSpy).toHaveBeenNthCalledWith(
        1,
        expect.objectContaining({
          sessionUpdate: 'agent_message_chunk',
          content: { type: 'text', text: 'First chunk ' },
        }),
      );
      expect(sendUpdateSpy).toHaveBeenNthCalledWith(
        2,
        expect.objectContaining({
          sessionUpdate: 'agent_message_chunk',
          content: { type: 'text', text: 'Second chunk ' },
        }),
      );
      expect(sendUpdateSpy).toHaveBeenNthCalledWith(
        3,
        expect.objectContaining({
          sessionUpdate: 'agent_message_chunk',
          content: { type: 'text', text: 'Third chunk' },
        }),
      );
    });

    it('should not emit when aborted', async () => {
      tracker.setup(eventEmitter, abortController.signal);
      abortController.abort();

      const event = createStreamTextEvent({
        text: 'This should not be emitted',
      });

      eventEmitter.emit(AgentEventType.STREAM_TEXT, event);

      await new Promise((resolve) => setTimeout(resolve, 10));

      expect(sendUpdateSpy).not.toHaveBeenCalled();
    });

    it('should emit agent_thought_chunk when thought flag is true', async () => {
      tracker.setup(eventEmitter, abortController.signal);

      const event = createStreamTextEvent({
        text: 'Let me think about this...',
        thought: true,
      });

      eventEmitter.emit(AgentEventType.STREAM_TEXT, event);

      await vi.waitFor(() => {
        expect(sendUpdateSpy).toHaveBeenCalled();
      });

      expect(sendUpdateSpy).toHaveBeenCalledWith(
        expect.objectContaining({
          sessionUpdate: 'agent_thought_chunk',
          content: {
            type: 'text',
            text: 'Let me think about this...',
          },
        }),
      );
    });

    it('should emit agent_message_chunk when thought flag is false', async () => {
      tracker.setup(eventEmitter, abortController.signal);

      const event = createStreamTextEvent({
        text: 'Here is the answer.',
        thought: false,
      });

      eventEmitter.emit(AgentEventType.STREAM_TEXT, event);

      await vi.waitFor(() => {
        expect(sendUpdateSpy).toHaveBeenCalled();
      });

      expect(sendUpdateSpy).toHaveBeenCalledWith(
        expect.objectContaining({
          sessionUpdate: 'agent_message_chunk',
          content: {
            type: 'text',
            text: 'Here is the answer.',
          },
        }),
      );
    });

    it('should emit agent_message_chunk when thought flag is undefined', async () => {
      tracker.setup(eventEmitter, abortController.signal);

      // Event without thought flag (undefined)
      const event = createStreamTextEvent({
        text: 'Default behavior text.',
      });

      eventEmitter.emit(AgentEventType.STREAM_TEXT, event);

      await vi.waitFor(() => {
        expect(sendUpdateSpy).toHaveBeenCalled();
      });

      expect(sendUpdateSpy).toHaveBeenCalledWith(
        expect.objectContaining({
          sessionUpdate: 'agent_message_chunk',
          content: {
            type: 'text',
            text: 'Default behavior text.',
          },
        }),
      );
    });
  });
});