packages/agent-core-v2/test/features/sessionInit/sessionInit.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';

import { SyncDescriptor } from '#/_base/di/descriptors';
import { DisposableStore } from '#/_base/di/lifecycle';
import { TestInstantiationService } from '#/_base/di/test';
import { UserCancellationError } from '#/_base/utils/abort';
import { IBootstrapService } from '#/app/bootstrap/bootstrap';
import { IEventBus } from '#/app/event/eventBus';
import { ITelemetryService } from '#/app/telemetry/telemetry';
import { IHostEnvironment } from '#/os/interface/hostEnvironment';
import { IHostFileSystem, type HostFileStat } from '#/os/interface/hostFileSystem';
import { IAgentPermissionModeService } from '#/agent/permissionMode/permissionMode';
import { IAgentProfileService } from '#/agent/profile/profile';
import { IAgentAgentsMdReminderService } from '#/agent/agentsMdReminder/agentsMdReminder';
import { IEventDispatcher } from '#/state/eventDispatcher';
import { ErrorCodes, Error2 } from '#/errors';
import type { AgentContext } from '#/agent/agentContext/agentContext';
import { IAgentScopeContext } from '#/agent/scopeContext/scopeContext';
import { IAgentLifecycleService } from '#/session/agentLifecycle/agentLifecycle';
import { ISessionContext } from '#/session/sessionContext/sessionContext';
import { ISessionInitService } from '#/features/sessionInit/sessionInit';
import { SessionInitService } from '#/features/sessionInit/sessionInitService';
import { IAgentReminderService } from '#/features/reminder/reminderService';
import { ISessionSubagentService } from '#/session/subagent/subagent';
import { stubAgentContext } from '../../agent/agentContext/stubs';

const WORK_DIR = '/project';
const AGENTS_MD = 'latest project instructions';
const AGENTS_MD_PATH = `${WORK_DIR}/AGENTS.md`;
const GIT_DIR_PATH = `${WORK_DIR}/.git`;

describe('SessionInitService', () => {
  let disposables: DisposableStore;
  let ix: TestInstantiationService;
  let events: unknown[];
  let appendReminder: ReturnType<typeof vi.fn>;
  let seedInjected: ReturnType<typeof vi.fn>;
  let flush: ReturnType<typeof vi.fn>;
  let republishStatus: ReturnType<typeof vi.fn>;
  let create: ReturnType<typeof vi.fn>;
  let run: ReturnType<typeof vi.fn>;
  let runCompletion: Promise<{ summary: string; usage?: undefined }>;

  beforeEach(() => {
    disposables = new DisposableStore();
    ix = disposables.add(new TestInstantiationService());
    events = [];
    appendReminder = vi.fn(() => 'reminder-id');
    seedInjected = vi.fn();
    flush = vi.fn(async () => {});
    republishStatus = vi.fn(() => {
      events.push({ type: 'agent.status.updated', model: 'mock-model' });
    });
    runCompletion = Promise.resolve({ summary: 'Explored and wrote AGENTS.md', usage: undefined });

    const handles: Record<string, { id: string; accessor: { get: (id: unknown) => unknown } }> = {};
    const lifecycle = {
      _serviceBrand: undefined,
      hooks: {
        onWillStartAgentTask: { run: vi.fn(async () => {}) },
      },
      notifyAgentTaskStopped: vi.fn(),
      handleOf: vi.fn((agentId: string) => handles[agentId]),
      create: vi.fn(async () => stubAgentContext('agent-0', 1)),
      run: vi.fn(async (agent: AgentContext) => ({
        agentId: agent.agentId,
        turn: {},
        completion: runCompletion,
      })),
    };
    create = lifecycle.create;
    run = lifecycle.run;

    const eventBus = { publish: vi.fn((event: unknown) => events.push(event)) };
    const telemetry = { track: vi.fn(), track2: vi.fn() };
    const profile = {
      data: () => ({ modelAlias: 'mock-model', thinkingLevel: 'off' }),
    };
    const permissionMode = { mode: 'auto', setMode: vi.fn() };

    handles['main'] = {
      id: 'main',
      accessor: {
        get: (id: unknown) => {
          if (id === IAgentLifecycleService) return lifecycle;
          if (id === ISessionSubagentService) return lifecycle;
          if (id === IAgentScopeContext) {
            return { agentContext: stubAgentContext('main', 1) };
          }
          if (id === IAgentProfileService) return profile;
          if (id === IAgentPermissionModeService) return permissionMode;
          if (id === IAgentAgentsMdReminderService) return { seedInjected };
          if (id === IAgentReminderService) return { notify: appendReminder };
          if (id === IEventDispatcher) {
            return {
              flush,
              dispatch: async (event: unknown) => {
                eventBus.publish(event);
              },
            };
          }
          if (id === IEventBus) return eventBus;
          if (id === ITelemetryService) return telemetry;
          return undefined;
        },
      },
    };
    handles['agent-0'] = {
      id: 'agent-0',
      accessor: {
        get: (id: unknown) => {
          if (id === IAgentScopeContext) {
            return {
              agentId: 'agent-0',
              agentContext: stubAgentContext('agent-0', 1),
            };
          }
          if (id === IAgentPermissionModeService) return permissionMode;
          if (id === IAgentProfileService)
            return { republishStatus, getEffectiveThinkingLevel: () => 'off' };
          return undefined;
        },
      },
    };

    ix.stub(IAgentLifecycleService, lifecycle as unknown as IAgentLifecycleService);
    ix.stub(ISessionSubagentService, lifecycle as unknown as ISessionSubagentService);
    ix.stub(IHostFileSystem, {
      _serviceBrand: undefined,
      stat: vi.fn(async (path: string): Promise<HostFileStat> => {
        if (path === GIT_DIR_PATH) return { isFile: false, isDirectory: true, size: 0 };
        if (path === AGENTS_MD_PATH)
          return { isFile: true, isDirectory: false, size: AGENTS_MD.length };
        throw new Error(`ENOENT: ${path}`);
      }),
      readText: vi.fn(async (path: string) => {
        if (path === AGENTS_MD_PATH) return AGENTS_MD;
        throw new Error(`ENOENT: ${path}`);
      }),
    } as unknown as IHostFileSystem);
    ix.stub(IHostEnvironment, {
      _serviceBrand: undefined,
      homeDir: '/home',
    } as unknown as IHostEnvironment);
    ix.stub(IBootstrapService, {
      _serviceBrand: undefined,
      homeDir: '/home/brand',
    } as unknown as IBootstrapService);
    ix.stub(ISessionContext, {
      _serviceBrand: undefined,
      cwd: WORK_DIR,
    } as unknown as ISessionContext);
    ix.set(ISessionInitService, new SyncDescriptor(SessionInitService));
  });

  afterEach(() => disposables.dispose());

  it('spawns a coder subagent, reloads AGENTS.md, and reminds the main agent', async () => {
    const svc = ix.get(ISessionInitService);
    await svc.generateAgentsMd();

    expect(create).toHaveBeenCalledTimes(1);
    expect(create.mock.calls[0]![0]).toMatchObject({
      binding: { profile: 'coder', model: 'mock-model', thinking: 'off' },
    });

    expect(run).toHaveBeenCalledTimes(1);
    const runArgs = run.mock.calls[0]!;
    expect(runArgs[0]).toMatchObject({ agentId: 'agent-0', generation: 1 });
    expect(runArgs[1]).toMatchObject({ kind: 'prompt' });
    expect((runArgs[1] as { prompt: string }).prompt).toContain('Task requirements:');

    expect(appendReminder).toHaveBeenCalledTimes(1);
    const [content, notification] = appendReminder.mock.calls[0] as [
      string,
      { variant: string },
    ];
    expect(notification).toEqual({ variant: 'init' });
    expect(content).toContain('The user just ran `/init` slash command.');
    expect(content).toContain('Latest AGENTS.md file content:');
    expect(content).toContain(AGENTS_MD);

    expect(seedInjected).toHaveBeenCalledWith([AGENTS_MD_PATH], WORK_DIR);

    expect(flush).toHaveBeenCalledTimes(1);

    expect(events).toContainEqual(
      expect.objectContaining({
        type: 'subagent.spawned',
        subagentId: 'agent-0',
        subagentName: 'coder',
        parentToolCallId: 'generate-agents-md',
        callerAgentId: 'main',
        model: 'mock-model',
        thinkingEffort: 'off',
      }),
    );
    expect(republishStatus).toHaveBeenCalledTimes(1);
    const eventTypes = events.map((event) => (event as { type?: string }).type);
    const spawnedIndex = eventTypes.indexOf('subagent.spawned');
    expect(eventTypes[spawnedIndex + 1]).toBe('agent.status.updated');
    expect(events).toContainEqual(
      expect.objectContaining({ type: 'subagent.completed', subagentId: 'agent-0' }),
    );
  });

  it('wraps a subagent failure in SESSION_INIT_FAILED', async () => {
    run.mockImplementationOnce((agentId: string) => ({
      agentId,
      turn: {},
      completion: Promise.reject(new Error('coder exploded')),
    }));
    const svc = ix.get(ISessionInitService);

    const error = await svc.generateAgentsMd().catch((e) => e);
    expect(error).toBeInstanceOf(Error2);
    expect((error as Error2).code).toBe(ErrorCodes.SESSION_INIT_FAILED);
    expect((error as Error2).message).toContain('coder exploded');
  });

  it('throws AGENT_NOT_FOUND when the main agent is missing', async () => {
    const lifecycle = ix.get(IAgentLifecycleService) as unknown as {
      handleOf: ReturnType<typeof vi.fn>;
    };
    lifecycle.handleOf.mockReturnValue(undefined);
    const svc = ix.get(ISessionInitService);

    const error = await svc.generateAgentsMd().catch((e) => e);
    expect(error).toBeInstanceOf(Error2);
    expect((error as Error2).code).toBe(ErrorCodes.AGENT_NOT_FOUND);
  });

  it('cancelInit aborts the in-flight run without wrapping the cancellation', async () => {
    run.mockImplementationOnce((agentId: string, _req: unknown, opts: { signal: AbortSignal }) => ({
      agentId,
      turn: {},
      completion: new Promise<{ summary: string }>((_resolve, reject) => {
        opts.signal.addEventListener('abort', () => reject(opts.signal.reason));
      }),
    }));
    const svc = ix.get(ISessionInitService);

    const pending = svc.generateAgentsMd();
    await vi.waitFor(() => expect(run).toHaveBeenCalled());
    svc.cancelInit();

    const error = await pending.catch((e) => e);
    expect(error).toBeInstanceOf(UserCancellationError);
    expect(events).not.toContainEqual(
      expect.objectContaining({ type: 'subagent.failed', subagentId: 'agent-0' }),
    );
  });

  it('cancelInit is a no-op when no init run is in flight', () => {
    const svc = ix.get(ISessionInitService);
    expect(() => svc.cancelInit()).not.toThrow();
  });
});