my-pi/pi-subagents/test/child-session-shutdown.test.ts

/**
 * #242 — `runAgent` calls `session.bindExtensions()` so `session_start` fires and
 * extensions can set up per-session state, but nothing ever closed that lifecycle:
 * both eviction (`removeRecord`) and quit (`dispose`) called `session.dispose()`,
 * which in pi only calls `ExtensionRunner.invalidate()` — it does NOT emit
 * `session_shutdown`. So anything an extension armed in `session_start` (timers, fs
 * watchers, sockets) leaked once per spawn, and its next tick threw `assertActive()`
 * from a bare `Timeout._onTimeout` — an uncaughtException that killed interactive pi.
 *
 * pi emits the event itself in `AgentSessionRuntime.dispose()` before disposing; these
 * tests pin that we do the same, that quit actually waits for the handlers, and that a
 * hung handler can't strand the user at a dead terminal.
 */
import { afterEach, describe, expect, it, vi } from "vitest";
import { AgentManager } from "../src/agent-manager.js";
import { registerPermissionChildSession } from "../src/permission-system-bridge.js";

vi.mock("../src/agent-runner.js", () => ({
  runAgent: vi.fn(),
  resumeAgent: vi.fn(),
}));

vi.mock("../src/worktree.js", () => ({
  createWorktree: vi.fn(),
  cleanupWorktree: vi.fn(() => ({ hasChanges: false })),
  pruneWorktrees: vi.fn(),
  isWorktreeIsolationEnabled: vi.fn(() => true),
}));

import { runAgent } from "../src/agent-runner.js";
import { pruneWorktrees } from "../src/worktree.js";

const mockPi = {} as any;
const mockCtx = { cwd: "/tmp" } as any;

/** A child session as `runAgent` leaves it: extensions bound, so a runner with handlers. */
function boundSession(emit: (...args: any[]) => any = vi.fn(async () => {})) {
  return {
    dispose: vi.fn(),
    extensionRunner: {
      hasHandlers: vi.fn((event: string) => event === "session_shutdown"),
      emit: vi.fn(emit),
    },
  } as any;
}

/** Spawn one background agent that resolves with `session`, and wait for it to complete. */
async function spawnCompleted(manager: AgentManager, session: any) {
  vi.mocked(runAgent).mockResolvedValue({
    responseText: "done",
    session,
    aborted: false,
    steered: false,
  } as any);
  const id = manager.spawn(mockPi, mockCtx, "general-purpose", "test", {
    description: "test",
    isBackground: true,
  });
  await manager.getRecord(id)!.promise;
  return id;
}

describe("child session shutdown (#242)", () => {
  let manager: AgentManager;

  afterEach(async () => {
    await manager?.dispose();
    vi.useRealTimers();
    vi.clearAllMocks();
  });

  it("emits session_shutdown before disposing an evicted session", async () => {
    manager = new AgentManager();
    const session = boundSession();
    await spawnCompleted(manager, session);

    manager.clearCompleted();
    await vi.waitFor(() => expect(session.dispose).toHaveBeenCalled());

    expect(session.extensionRunner.emit).toHaveBeenCalledWith({
      type: "session_shutdown",
      reason: "quit",
    });
    // Order is the whole point: after `dispose()` the runner is invalidated and
    // every `ctx` getter throws, so a handler emitted afterwards is useless.
    expect(session.extensionRunner.emit.mock.invocationCallOrder[0])
      .toBeLessThan(session.dispose.mock.invocationCallOrder[0]);
  });

  it("publishes child disposal only after the session is closed", async () => {
    manager = new AgentManager();
    const order: string[] = [];
    const pi = {
      events: { emit: vi.fn((channel: string) => { order.push(channel); }) },
    } as any;
    const session = boundSession();
    session.sessionManager = { getSessionId: () => "child-session" };
    session.dispose.mockImplementation(() => { order.push("session-disposed"); });
    registerPermissionChildSession(pi, session, "root-session");
    await spawnCompleted(manager, session);

    manager.clearCompleted();
    await vi.waitFor(() => expect(session.dispose).toHaveBeenCalled());

    expect(order).toEqual([
      "subagents:child:session-created",
      "session-disposed",
      "subagents:child:disposed",
    ]);
  });

  it("quit waits for the child's shutdown handlers to finish", async () => {
    manager = new AgentManager();
    let releaseHandler!: () => void;
    const session = boundSession(() => new Promise<void>(r => { releaseHandler = r; }));
    await spawnCompleted(manager, session);

    const disposed = manager.dispose();
    let settled = false;
    void disposed.then(() => { settled = true; });
    // Several microtask turns: enough for a fire-and-forget implementation to have
    // resolved, not enough for a correctly awaited one.
    for (let i = 0; i < 5; i++) await Promise.resolve();

    expect(settled).toBe(false);
    expect(session.dispose).not.toHaveBeenCalled();

    releaseHandler();
    await disposed;
    expect(settled).toBe(true);
    expect(session.dispose).toHaveBeenCalledOnce();
  });

  it("quit is not hostage to a handler that never resolves", async () => {
    manager = new AgentManager();
    const session = boundSession(() => new Promise<void>(() => {}));
    await spawnCompleted(manager, session);

    vi.useFakeTimers();
    const disposed = manager.dispose();
    // Past the internal ceiling. Without it the TUI is already torn down and the
    // user is left at a dead terminal with only Ctrl-C.
    await vi.advanceTimersByTimeAsync(5_000);
    await disposed;

    expect(session.dispose).toHaveBeenCalledOnce();
    // Teardown continues past the timeout rather than unwinding.
    expect(pruneWorktrees).toHaveBeenCalled();
  });

  it("skips the emit when no extension handles session_shutdown", async () => {
    manager = new AgentManager();
    const session = boundSession();
    session.extensionRunner.hasHandlers = vi.fn(() => false);
    await spawnCompleted(manager, session);

    manager.clearCompleted();
    await vi.waitFor(() => expect(session.dispose).toHaveBeenCalled());

    expect(session.extensionRunner.emit).not.toHaveBeenCalled();
  });

  it("degrades on a stubbed session instead of throwing", async () => {
    const rejections: unknown[] = [];
    const onRejection = (err: unknown) => rejections.push(err);
    process.on("unhandledRejection", onRejection);
    try {
      manager = new AgentManager();
      // No extensionRunner at all — an older pi, or a partial `onSessionCreated` stub.
      const noRunner = { dispose: vi.fn() } as any;
      await spawnCompleted(manager, noRunner);
      manager.clearCompleted();
      await vi.waitFor(() => expect(noRunner.dispose).toHaveBeenCalled());

      // Nothing at all — not even dispose.
      const empty = {} as any;
      await spawnCompleted(manager, empty);
      manager.clearCompleted();
      await expect(manager.dispose()).resolves.toBeUndefined();

      await new Promise(r => setImmediate(r));
      expect(rejections).toEqual([]);
    } finally {
      process.off("unhandledRejection", onRejection);
    }
  });
});