my-pi/pi-permission-system/test/handlers/before-agent-start.test.ts

import { describe, expect, it, vi } from "vitest";

import { AgentPrepHandler } from "#src/handlers/before-agent-start";

import { makeCheckResult, makeCtx } from "#test/helpers/handler-fixtures";
import {
  makeRealResolver,
  makeRealSession,
} from "#test/helpers/session-fixtures";

function makeSetup() {
  const { session, permissionManager, sessionRules, configStore, forwarding } =
    makeRealSession();
  const { resolver } = makeRealResolver(permissionManager, sessionRules);
  vi.mocked(permissionManager.check).mockReturnValue(makeCheckResult());
  const warmParser = vi.fn();
  const handler = new AgentPrepHandler(session, resolver, warmParser);
  return {
    handler,
    session,
    permissionManager,
    configStore,
    forwarding,
    warmParser,
  };
}

describe("AgentPrepHandler", () => {
  it("warms the parser and activates the permission session", async () => {
    const { handler, forwarding, warmParser } = makeSetup();
    const ctx = makeCtx();

    await handler.handle({ systemPrompt: "You are an assistant." }, ctx);

    expect(warmParser).toHaveBeenCalledTimes(1);
    expect(forwarding.start).toHaveBeenCalledWith(ctx);
  });

  it("refreshes global permission configuration", async () => {
    const { handler, configStore } = makeSetup();
    const ctx = makeCtx({ isProjectTrusted: () => false });

    await handler.handle({ systemPrompt: "You are an assistant." }, ctx);

    expect(configStore.refresh).toHaveBeenCalledWith(ctx, false);
  });

  it("does not rewrite a prompt without skill entries", async () => {
    const { handler, session } = makeSetup();

    const result = await handler.handle(
      { systemPrompt: "You are an assistant." },
      makeCtx(),
    );

    expect(result).toEqual({});
    expect(session.getActiveSkillEntries()).toEqual([]);
  });

  it("filters denied skills out of structured systemPromptOptions.skills", async () => {
    const { handler, session, permissionManager } = makeSetup();
    vi.mocked(permissionManager.check).mockImplementation((intent) =>
      makeCheckResult({
        state:
          intent.kind === "tool" &&
          (intent.input as { name?: string })?.name === "dangerous"
            ? "deny"
            : "allow",
      }),
    );

    const options = {
      skills: [
        { name: "librarian", description: "d", filePath: "/skills/librarian/SKILL.md", baseDir: "/skills/librarian" },
        { name: "dangerous", description: "d", filePath: "/skills/dangerous/SKILL.md", baseDir: "/skills/dangerous" },
        { name: "manual", description: "d", filePath: "/skills/manual/SKILL.md", baseDir: "/skills/manual", disableModelInvocation: true },
      ],
    };
    const result = await handler.handle(
      { systemPrompt: "You are an assistant.", systemPromptOptions: options },
      makeCtx(),
    );

    expect(result).toEqual({});
    expect(options.skills.map((skill) => skill.name)).toEqual(["librarian", "manual"]);
    expect(session.getActiveSkillEntries().map((entry) => entry.name)).toEqual([
      "librarian",
    ]);
  });
});