File size: 18,637 Bytes

7a1ad33

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

import {
  describe,
  it,
  expect,
  vi,
  beforeEach,
  afterEach,
  type Mocked,
} from 'vitest';
import {
  AgentTerminateMode,
  type LocalAgentDefinition,
  type SubagentActivityEvent,
  type AgentInputs,
  type SubagentProgress,
  SubagentActivityErrorType,
  SUBAGENT_REJECTED_ERROR_PREFIX,
  SubagentState,
} from './types.js';
import { LocalSubagentInvocation } from './local-invocation.js';
import { LocalAgentExecutor } from './local-executor.js';
import { makeFakeConfig } from '../test-utils/config.js';
import type { Config } from '../config/config.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
import { type z } from 'zod';
import { createMockMessageBus } from '../test-utils/mock-message-bus.js';

vi.mock('./local-executor.js');

const MockLocalAgentExecutor = vi.mocked(LocalAgentExecutor);

let mockConfig: Config;

const testDefinition: LocalAgentDefinition<z.ZodUnknown> = {
  kind: 'local',
  name: 'MockAgent',
  displayName: 'Mock Agent',
  description: 'A mock agent.',
  inputConfig: {
    inputSchema: {
      type: 'object',
      properties: {
        task: { type: 'string', description: 'task' },
        priority: { type: 'number', description: 'prio' },
      },
      required: ['task'],
    },
  },
  modelConfig: {
    model: 'test',
    generateContentConfig: {
      temperature: 0,
      topP: 1,
    },
  },
  runConfig: { maxTimeMinutes: 1 },
  promptConfig: { systemPrompt: 'test' },
};

describe('LocalSubagentInvocation', () => {
  let mockExecutorInstance: Mocked<LocalAgentExecutor<z.ZodUnknown>>;
  let mockMessageBus: MessageBus;

  beforeEach(() => {
    vi.clearAllMocks();
    mockConfig = makeFakeConfig();
    // .config is already set correctly by the getter on the instance.
    Object.defineProperty(mockConfig, 'promptId', {
      get: () => 'test-prompt-id',
      configurable: true,
    });
    mockMessageBus = createMockMessageBus();

    mockExecutorInstance = {
      run: vi.fn(),
      definition: testDefinition,
      agentId: 'test-agent-id',
    } as unknown as Mocked<LocalAgentExecutor<z.ZodUnknown>>;

    MockLocalAgentExecutor.create.mockResolvedValue(
      mockExecutorInstance as unknown as LocalAgentExecutor<z.ZodTypeAny>,
    );
  });

  afterEach(() => {
    vi.restoreAllMocks();
  });

  it('should pass the messageBus to the parent constructor', () => {
    const params = { task: 'Analyze data' };
    const invocation = new LocalSubagentInvocation(
      testDefinition,
      mockConfig,
      params,
      mockMessageBus,
    );

    // Access the protected messageBus property by casting to any
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    expect((invocation as any).messageBus).toBe(mockMessageBus);
  });

  describe('getDescription', () => {
    it('should format the description with inputs', () => {
      const params = { task: 'Analyze data', priority: 5 };
      const invocation = new LocalSubagentInvocation(
        testDefinition,
        mockConfig,
        params,
        mockMessageBus,
      );
      const description = invocation.getDescription();
      expect(description).toBe(
        "Running subagent 'MockAgent' with inputs: { task: Analyze data, priority: 5 }",
      );
    });

    it('should not truncate long input values', () => {
      const longTask = 'A'.repeat(100);
      const params = { task: longTask };
      const invocation = new LocalSubagentInvocation(
        testDefinition,
        mockConfig,
        params,
        mockMessageBus,
      );
      const description = invocation.getDescription();
      expect(description).toBe(
        `Running subagent 'MockAgent' with inputs: { task: ${'A'.repeat(100)} }`,
      );
    });

    it('should not truncate the overall description', () => {
      // Create a definition and inputs that result in a very long description
      const longNameDef: LocalAgentDefinition = {
        ...testDefinition,
        name: 'VeryLongAgentNameThatTakesUpSpace',
      };
      const params: AgentInputs = {};
      for (let i = 0; i < 20; i++) {
        params[`input${i}`] = `value${i}`;
      }
      const invocation = new LocalSubagentInvocation(
        longNameDef,
        mockConfig,
        params,
        mockMessageBus,
      );
      const description = invocation.getDescription();
      expect(description.length).toBeGreaterThan(300);
      expect(
        description.startsWith(
          "Running subagent 'VeryLongAgentNameThatTakesUpSpace'",
        ),
      ).toBe(true);
    });
  });

  describe('execute', () => {
    let signal: AbortSignal;
    let updateOutput: ReturnType<typeof vi.fn>;
    const params = { task: 'Execute task' };
    let invocation: LocalSubagentInvocation;

    beforeEach(() => {
      signal = new AbortController().signal;
      updateOutput = vi.fn();
      invocation = new LocalSubagentInvocation(
        testDefinition,
        mockConfig,
        params,
        mockMessageBus,
      );
    });

    it('should initialize and run the executor successfully', async () => {
      const mockOutput = {
        result: 'Analysis complete.',
        terminate_reason: AgentTerminateMode.GOAL,
      };
      mockExecutorInstance.run.mockResolvedValue(mockOutput);

      const result = await invocation.execute({
        abortSignal: signal,
        updateOutput,
      });

      expect(MockLocalAgentExecutor.create).toHaveBeenCalledWith(
        testDefinition,
        mockConfig,
        expect.any(Function),
      );
      expect(updateOutput).toHaveBeenCalledWith(
        expect.objectContaining({
          isSubagentProgress: true,
          agentName: 'MockAgent',
        }),
      );

      expect(mockExecutorInstance.run).toHaveBeenCalledWith(params, signal);

      expect(result.llmContent).toEqual([
        {
          text: expect.stringContaining(
            "Subagent 'MockAgent' finished.\nTermination Reason: GOAL\nResult:\nAnalysis complete.",
          ),
        },
      ]);
      const display = result.returnDisplay as SubagentProgress;
      expect(display.isSubagentProgress).toBe(true);
      expect(display.state).toBe(SubagentState.COMPLETED);
      expect(display.result).toBe('Analysis complete.');
      expect(display.terminateReason).toBe(AgentTerminateMode.GOAL);
    });

    it('should show detailed UI for non-goal terminations (e.g., TIMEOUT)', async () => {
      const mockOutput = {
        result: 'Partial progress...',
        terminate_reason: AgentTerminateMode.TIMEOUT,
      };
      mockExecutorInstance.run.mockResolvedValue(mockOutput);

      const result = await invocation.execute({
        abortSignal: signal,
        updateOutput,
      });

      const display = result.returnDisplay as SubagentProgress;
      expect(display.isSubagentProgress).toBe(true);
      expect(display.state).toBe(SubagentState.COMPLETED);
      expect(display.result).toBe('Partial progress...');
      expect(display.terminateReason).toBe(AgentTerminateMode.TIMEOUT);
    });

    it('should stream THOUGHT_CHUNK activities from the executor, replacing the last running thought', async () => {
      mockExecutorInstance.run.mockImplementation(async () => {
        const onActivity = MockLocalAgentExecutor.create.mock.calls[0][2];

        if (onActivity) {
          onActivity({
            isSubagentActivityEvent: true,
            agentName: 'MockAgent',
            type: 'THOUGHT_CHUNK',
            data: { text: 'Analyzing...' },
          } as SubagentActivityEvent);
          onActivity({
            isSubagentActivityEvent: true,
            agentName: 'MockAgent',
            type: 'THOUGHT_CHUNK',
            data: { text: 'Thinking about next steps.' },
          } as SubagentActivityEvent);
        }
        return { result: 'Done', terminate_reason: AgentTerminateMode.GOAL };
      });

      await invocation.execute({ abortSignal: signal, updateOutput });

      expect(updateOutput).toHaveBeenCalledTimes(4); // Initial + 2 updates + Final completion
      const lastCall = updateOutput.mock.calls[3][0] as SubagentProgress;
      expect(lastCall.recentActivity).toContainEqual(
        expect.objectContaining({
          type: 'thought',
          content: 'Thinking about next steps.',
        }),
      );
      expect(lastCall.recentActivity).not.toContainEqual(
        expect.objectContaining({
          type: 'thought',
          content: 'Analyzing...',
        }),
      );
    });

    it('should overwrite the thought content with new THOUGHT_CHUNK activity', async () => {
      mockExecutorInstance.run.mockImplementation(async () => {
        const onActivity = MockLocalAgentExecutor.create.mock.calls[0][2];

        if (onActivity) {
          onActivity({
            isSubagentActivityEvent: true,
            agentName: 'MockAgent',
            type: 'THOUGHT_CHUNK',
            data: { text: 'I am thinking.' },
          } as SubagentActivityEvent);
          onActivity({
            isSubagentActivityEvent: true,
            agentName: 'MockAgent',
            type: 'THOUGHT_CHUNK',
            data: { text: 'Now I will act.' },
          } as SubagentActivityEvent);
        }
        return { result: 'Done', terminate_reason: AgentTerminateMode.GOAL };
      });

      await invocation.execute({ abortSignal: signal, updateOutput });

      const calls = updateOutput.mock.calls;
      const lastCall = calls[calls.length - 1][0] as SubagentProgress;
      expect(lastCall.recentActivity).toContainEqual(
        expect.objectContaining({
          type: 'thought',
          content: 'Now I will act.',
        }),
      );
    });

    it('should stream other activities (e.g., TOOL_CALL_START, ERROR)', async () => {
      mockExecutorInstance.run.mockImplementation(async () => {
        const onActivity = MockLocalAgentExecutor.create.mock.calls[0][2];

        if (onActivity) {
          onActivity({
            isSubagentActivityEvent: true,
            agentName: 'MockAgent',
            type: 'TOOL_CALL_START',
            data: { name: 'ls', args: {} },
          } as SubagentActivityEvent);
          onActivity({
            isSubagentActivityEvent: true,
            agentName: 'MockAgent',
            type: 'ERROR',
            data: { error: 'Failed' },
          } as SubagentActivityEvent);
        }
        return { result: 'Done', terminate_reason: AgentTerminateMode.GOAL };
      });

      await invocation.execute({ abortSignal: signal, updateOutput });

      expect(updateOutput).toHaveBeenCalledTimes(4); // Initial + 2 updates + Final completion
      const lastCall = updateOutput.mock.calls[3][0] as SubagentProgress;
      expect(lastCall.recentActivity).toContainEqual(
        expect.objectContaining({
          type: 'thought',
          content: 'Error: Failed',
          status: SubagentState.ERROR,
        }),
      );
    });

    it('should mark tool call as error when TOOL_CALL_END contains isError: true', async () => {
      mockExecutorInstance.run.mockImplementation(async () => {
        const onActivity = MockLocalAgentExecutor.create.mock.calls[0][2];

        if (onActivity) {
          onActivity({
            isSubagentActivityEvent: true,
            agentName: 'MockAgent',
            type: 'TOOL_CALL_START',
            data: { name: 'ls', args: {}, callId: 'call1' },
          } as SubagentActivityEvent);
          onActivity({
            isSubagentActivityEvent: true,
            agentName: 'MockAgent',
            type: 'TOOL_CALL_END',
            data: { name: 'ls', id: 'call1', data: { isError: true } },
          } as SubagentActivityEvent);
        }
        return { result: 'Done', terminate_reason: AgentTerminateMode.GOAL };
      });

      await invocation.execute({ abortSignal: signal, updateOutput });

      expect(updateOutput).toHaveBeenCalled();
      const lastCall = updateOutput.mock.calls[
        updateOutput.mock.calls.length - 1
      ][0] as SubagentProgress;
      expect(lastCall.recentActivity).toContainEqual(
        expect.objectContaining({
          type: 'tool_call',
          content: 'ls',
          status: SubagentState.ERROR,
        }),
      );
    });

    it('should reflect tool rejections in the activity stream as cancelled but not abort the agent', async () => {
      mockExecutorInstance.run.mockImplementation(async () => {
        const onActivity = MockLocalAgentExecutor.create.mock.calls[0][2];

        if (onActivity) {
          onActivity({
            isSubagentActivityEvent: true,
            agentName: 'MockAgent',
            type: 'TOOL_CALL_START',
            data: { name: 'ls', args: {}, callId: 'call1' },
          } as SubagentActivityEvent);
          onActivity({
            isSubagentActivityEvent: true,
            agentName: 'MockAgent',
            type: 'ERROR',
            data: {
              name: 'ls',
              callId: 'call1',
              error: `${SUBAGENT_REJECTED_ERROR_PREFIX} Please acknowledge this, rethink your strategy, and try a different approach. If you cannot proceed without the rejected operation, summarize the issue and use \`complete_task\` to report your findings and the blocker.`,
              errorType: SubagentActivityErrorType.REJECTED,
            },
          } as SubagentActivityEvent);
        }
        return {
          result: 'Rethinking...',
          terminate_reason: AgentTerminateMode.GOAL,
        };
      });

      await invocation.execute({ abortSignal: signal, updateOutput });

      expect(updateOutput).toHaveBeenCalledTimes(4);
      const lastCall = updateOutput.mock.calls[3][0] as SubagentProgress;
      expect(lastCall.recentActivity).toContainEqual(
        expect.objectContaining({
          type: 'tool_call',
          content: 'ls',
          status: SubagentState.CANCELLED,
        }),
      );
    });

    it('should run successfully without an updateOutput callback', async () => {
      mockExecutorInstance.run.mockImplementation(async () => {
        const onActivity = MockLocalAgentExecutor.create.mock.calls[0][2];
        if (onActivity) {
          // Ensure calling activity doesn't crash when updateOutput is undefined
          onActivity({
            isSubagentActivityEvent: true,
            agentName: 'testAgent',
            type: 'THOUGHT_CHUNK',
            data: { text: 'Thinking silently.' },
          } as SubagentActivityEvent);
        }
        return { result: 'Done', terminate_reason: AgentTerminateMode.GOAL };
      });

      // Execute without the optional callback
      const result = await invocation.execute({ abortSignal: signal });
      expect(result.error).toBeUndefined();
      const display = result.returnDisplay as SubagentProgress;
      expect(display.isSubagentProgress).toBe(true);
      expect(display.state).toBe(SubagentState.COMPLETED);
      expect(display.result).toBe('Done');
    });

    it('should handle executor run failure', async () => {
      const error = new Error('Model failed during execution.');
      mockExecutorInstance.run.mockRejectedValue(error);

      const result = await invocation.execute({
        abortSignal: signal,
        updateOutput,
      });

      expect(result.error).toBeUndefined();
      expect(result.llmContent).toBe(
        `Subagent 'MockAgent' failed. Error: ${error.message}`,
      );
      const display = result.returnDisplay as SubagentProgress;
      expect(display.isSubagentProgress).toBe(true);
      expect(display.recentActivity).toContainEqual(
        expect.objectContaining({
          type: 'thought',
          content: `Error: ${error.message}`,
          status: SubagentState.ERROR,
        }),
      );
    });

    it('should handle executor creation failure', async () => {
      const creationError = new Error('Failed to initialize tools.');
      MockLocalAgentExecutor.create.mockRejectedValue(creationError);

      const result = await invocation.execute({
        abortSignal: signal,
        updateOutput,
      });

      expect(mockExecutorInstance.run).not.toHaveBeenCalled();
      expect(result.error).toBeUndefined();
      expect(result.llmContent).toContain(creationError.message);

      const display = result.returnDisplay as SubagentProgress;
      expect(display.recentActivity).toContainEqual(
        expect.objectContaining({
          content: `Error: ${creationError.message}`,
          status: SubagentState.ERROR,
        }),
      );
    });

    it('should handle abortion signal during execution', async () => {
      const abortError = new Error('Aborted');
      abortError.name = 'AbortError';
      mockExecutorInstance.run.mockRejectedValue(abortError);

      const controller = new AbortController();
      const executePromise = invocation.execute({
        abortSignal: controller.signal,
        updateOutput,
      });
      controller.abort();
      await expect(executePromise).rejects.toThrow('Aborted');

      expect(mockExecutorInstance.run).toHaveBeenCalledWith(
        params,
        controller.signal,
      );
    });

    it('should throw an error and bubble cancellation when execution returns ABORTED', async () => {
      const mockOutput = {
        result: 'Cancelled by user',
        terminate_reason: AgentTerminateMode.ABORTED,
      };
      mockExecutorInstance.run.mockResolvedValue(mockOutput);

      await expect(
        invocation.execute({ abortSignal: signal, updateOutput }),
      ).rejects.toThrow('Operation cancelled by user');
    });

    it('should publish SUBAGENT_ACTIVITY events to the MessageBus', async () => {
      const { MessageBusType } = await import('../confirmation-bus/types.js');

      mockExecutorInstance.run.mockImplementation(async () => {
        const onActivity = MockLocalAgentExecutor.create.mock.calls[0][2];

        if (onActivity) {
          onActivity({
            isSubagentActivityEvent: true,
            agentName: 'MockAgent',
            type: 'THOUGHT_CHUNK',
            data: { text: 'Thinking...' },
          } as SubagentActivityEvent);
        }
        return { result: 'Done', terminate_reason: AgentTerminateMode.GOAL };
      });

      await invocation.execute({ abortSignal: signal, updateOutput });

      expect(mockMessageBus.publish).toHaveBeenCalledWith(
        expect.objectContaining({
          type: MessageBusType.SUBAGENT_ACTIVITY,
          subagentName: 'Mock Agent',
          activity: expect.objectContaining({
            type: 'thought',
            content: 'Thinking...',
          }),
        }),
      );
    });
  });
});