my-pi/pi-permission-system/test/auto-review-reviewer.test.ts

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

import type { Api, AssistantMessage, Model } from "@earendil-works/pi-ai";
import type { SessionEntry } from "@earendil-works/pi-coding-agent";
import {
  DEFAULT_AUTO_REVIEW_CONFIG,
  type AutoReviewConfig,
} from "#src/auto-review/config";
import { REVIEW_SYSTEM_PROMPT, REVIEW_TOOLS } from "#src/auto-review/policy";
import {
  createApprovalCache,
  createPermissionReviewer,
  type ReviewContext,
} from "#src/auto-review/reviewer";
import type { PermissionQuery } from "#src/service";
import { makeAuthorizerLog } from "#test/helpers/authorizer-log-fixtures";
import { makePromptDetails } from "#test/helpers/prompt-details-fixtures";

// ── Fixtures ──────────────────────────────────────────────────────────────

function userEntry(id: string, text: string): SessionEntry {
  return {
    type: "message",
    id,
    message: { role: "user", content: text },
  } as unknown as SessionEntry;
}

/** A branch holding one trusted user instruction. */
const TRUSTED_BRANCH: SessionEntry[] = [userEntry("e1", "please read the readme")];

function makeModel(
  id: string,
  overrides: { provider?: string; reasoning?: boolean } = {},
): Model<Api> {
  return {
    id,
    provider: overrides.provider ?? "test-provider",
    reasoning: overrides.reasoning ?? false,
    api: "openai-completions",
    name: id,
    baseUrl: "https://example.invalid/v1/",
    input: ["text"],
    contextWindow: 128_000,
    maxTokens: 8_000,
    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
  } as unknown as Model<Api>;
}

function verdictMessage(
  name: "approve" | "reject" | string,
  reason = "looks fine",
): AssistantMessage {
  return {
    role: "assistant",
    content: [
      {
        type: "toolCall",
        id: "call-1",
        name,
        arguments: { reason },
      },
    ],
    stopReason: "toolUse",
    usage: { input: 11, output: 3 },
  } as unknown as AssistantMessage;
}

/**
 * Which registry completion surface the fake exposes — the real Pi 1.3.0
 * extension-facing registry has `streamSimple` only, so that is the default
 * path under test.
 */
type CompletionPath = "stream" | "complete" | "none";

interface HarnessOptions {
  readonly config?: Partial<AutoReviewConfig>;
  readonly branch?: SessionEntry[];
  readonly produce?: ReturnType<typeof vi.fn>;
  readonly sessionModel?: Model<Api> | undefined;
  readonly configuredModel?: Model<Api> | undefined;
  readonly completion?: CompletionPath;
  readonly sessionSignal?: AbortSignal;
  readonly ctxSignal?: AbortSignal;
  readonly registryFind?: ReturnType<typeof vi.fn>;
}

function harness(options: HarnessOptions = {}) {
  const {
    branch = TRUSTED_BRANCH,
    configuredModel = makeModel("configured-model", { provider: "cfg" }),
    completion = "stream",
    sessionSignal,
    ctxSignal,
  } = options;
  const sessionModel = "sessionModel" in options
    ? options.sessionModel
    : makeModel("session-model");
  const config: AutoReviewConfig = {
    ...DEFAULT_AUTO_REVIEW_CONFIG,
    ...options.config,
  };
  const produce =
    options.produce ?? vi.fn(async () => verdictMessage("approve"));
  // Mirror `streamSimple(model, context, options)` returning a stream whose
  // `result()` resolves the final AssistantMessage — the surface Pi 1.3.0's
  // extension-facing ModelRegistry actually publishes.
  const streamSimple = vi.fn(function (
    this: unknown,
    model: unknown,
    context: unknown,
    streamOptions?: unknown,
  ) {
    return {
      result: () =>
        (produce as (this: unknown, ...args: unknown[]) => Promise<unknown>).call(
          this,
          model,
          context,
          streamOptions,
        ),
    };
  });
  const find =
    options.registryFind ??
    vi.fn((provider: string, modelId: string) =>
      provider === "cfg" && modelId === configuredModel.id
        ? configuredModel
        : undefined,
    );
  const registry = {
    find,
    ...(completion === "stream" ? { streamSimple } : {}),
    ...(completion === "complete" ? { completeSimple: produce } : {}),
  };
  const ctx = {
    model: sessionModel,
    modelRegistry: registry,
    sessionManager: {
      getBranch: () => branch,
      getSessionId: () => "session-1",
    },
    signal: ctxSignal,
  };
  const approvalCache = createApprovalCache();
  const log = makeAuthorizerLog();

  const authorize = createPermissionReviewer({
    config,
    context: ctx as unknown as ReviewContext,
    approvalCache,
    ...(sessionSignal !== undefined ? { sessionSignal } : {}),
  });
  return { authorize, produce, streamSimple, registry, ctx, log };
}

const query = {
  checkPermission: vi.fn(),
  getToolPermission: vi.fn(),
} as unknown as PermissionQuery;

function failureCategory(log: ReturnType<typeof makeAuthorizerLog>): unknown {
  const call = log.debug.mock.calls.find(
    ([event]) => event === "auto_review.failure",
  );
  return call?.[1]?.errorCategory;
}

function decisionFields(log: ReturnType<typeof makeAuthorizerLog>) {
  const call = log.review.mock.calls.find(
    ([event]) => event === "auto_review.decision",
  );
  return call?.[1] as Record<string, unknown> | undefined;
}

// ── Decision path ─────────────────────────────────────────────────────────

describe("createPermissionReviewer", () => {
  it("follows the live session model in session-model mode", async () => {
    const first = makeModel("model-a");
    const second = makeModel("model-b");
    const { authorize, produce, ctx } = harness({
      sessionModel: first,
      config: { provider: undefined, model: undefined },
    });

    await authorize(makePromptDetails({ requestId: "r1", command: "ls a" }), query, makeAuthorizerLog());
    ctx.model = second;
    await authorize(makePromptDetails({ requestId: "r2", command: "ls b" }), query, makeAuthorizerLog());

    expect(produce).toHaveBeenCalledTimes(2);
    expect(produce.mock.calls[0]?.[0]).toBe(first);
    expect(produce.mock.calls[1]?.[0]).toBe(second);
  });

  it("uses the configured provider/model pair when both are set", async () => {
    const configured = makeModel("pinned", { provider: "cfg" });
    const { authorize, produce, registry } = harness({
      config: { provider: "cfg", model: "pinned" },
      configuredModel: configured,
    });

    const verdict = await authorize(makePromptDetails(), query, makeAuthorizerLog());

    expect(verdict).toEqual({ kind: "allow" });
    expect(registry.find).toHaveBeenCalledWith("cfg", "pinned");
    expect(produce.mock.calls[0]?.[0]).toBe(configured);
  });

  it("defers with the pairing note when exactly one of provider/model is set", async () => {
    const { authorize, produce, log } = harness({
      config: { provider: "cfg", model: undefined },
    });

    const verdict = await authorize(makePromptDetails(), query, log);

    expect(verdict.kind).toBe("defer");
    if (verdict.kind === "defer") {
      expect(verdict.note?.origin).toBe("system");
      expect(verdict.note?.text).toContain("provider 与 autoReview.model 需同时设置或同时留空");
    }
    expect(produce).not.toHaveBeenCalled();
    expect(failureCategory(log)).toBe("model-unconfigured");
  });

  it("defers with the missing-model note when the named model does not resolve", async () => {
    const { authorize, produce, log } = harness({
      config: { provider: "cfg", model: "nonexistent" },
    });

    const verdict = await authorize(makePromptDetails(), query, log);

    expect(verdict.kind).toBe("defer");
    if (verdict.kind === "defer") {
      expect(verdict.note?.text).toContain("cfg/nonexistent");
    }
    expect(produce).not.toHaveBeenCalled();
    expect(failureCategory(log)).toBe("model-unresolved");
  });

  it("defers with the no-session-model note in follow mode when ctx.model is unset", async () => {
    const { authorize, produce, log } = harness({ sessionModel: undefined });

    const verdict = await authorize(makePromptDetails(), query, log);

    expect(verdict.kind).toBe("defer");
    if (verdict.kind === "defer") {
      expect(verdict.note?.text).toContain("当前会话没有选中模型");
    }
    expect(produce).not.toHaveBeenCalled();
    expect(failureCategory(log)).toBe("no-session-model");
  });

  it("defers with the unsupported note when the registry cannot complete", async () => {
    const { authorize, log } = harness({ completion: "none" });

    const verdict = await authorize(makePromptDetails(), query, log);

    expect(verdict.kind).toBe("defer");
    if (verdict.kind === "defer") {
      expect(verdict.note?.text).toContain("不支持独立模型调用");
    }
    expect(failureCategory(log)).toBe("completion-unsupported");
  });

  it("returns allow when the model calls approve", async () => {
    const { authorize } = harness();

    const verdict = await authorize(makePromptDetails(), query, makeAuthorizerLog());

    expect(verdict).toEqual({ kind: "allow" });
  });

  it("defers with a model-origin note when the model calls reject", async () => {
    const { authorize } = harness({
      produce: vi.fn(async () => verdictMessage("reject", "too risky")),
    });

    const verdict = await authorize(makePromptDetails(), query, makeAuthorizerLog());

    expect(verdict.kind).toBe("defer");
    if (verdict.kind === "defer") {
      expect(verdict.note).toEqual({ origin: "model", text: "too risky" });
    }
  });

  it("treats a reply with no tool call as no-verdict, even valid JSON text", async () => {
    const message = {
      role: "assistant",
      content: [{ type: "text", text: '{"outcome":"allow"}' }],
      stopReason: "stop",
    } as unknown as AssistantMessage;
    const { authorize, produce, log } = harness({
      produce: vi.fn(async () => message),
    });

    const verdict = await authorize(makePromptDetails(), query, log);

    expect(verdict.kind).toBe("defer");
    if (verdict.kind === "defer") {
      expect(verdict.note?.origin).toBe("system");
      expect(verdict.note?.text).toContain("未调用 approve/reject");
    }
    expect(failureCategory(log)).toBe("no-verdict");
    expect(produce).toHaveBeenCalledTimes(1);
  });

  it("treats two tool calls as an invalid response", async () => {
    const message = {
      role: "assistant",
      content: [
        { type: "toolCall", id: "1", name: "approve", arguments: { reason: "a" } },
        { type: "toolCall", id: "2", name: "reject", arguments: { reason: "b" } },
      ],
      stopReason: "toolUse",
    } as unknown as AssistantMessage;
    const { authorize, produce, log } = harness({
      produce: vi.fn(async () => message),
    });

    const verdict = await authorize(makePromptDetails(), query, log);

    expect(verdict.kind).toBe("defer");
    if (verdict.kind === "defer") {
      expect(verdict.note?.text).toContain("结论格式无效");
    }
    expect(failureCategory(log)).toBe("invalid-response");
    expect(produce).toHaveBeenCalledTimes(1);
  });

  it("treats an unknown verdict tool as an invalid response", async () => {
    const { authorize, log } = harness({
      produce: vi.fn(async () => verdictMessage("maybe")),
    });

    const verdict = await authorize(makePromptDetails(), query, log);

    expect(verdict.kind).toBe("defer");
    expect(failureCategory(log)).toBe("invalid-response");
  });

  it("treats a missing or empty reason as an invalid response", async () => {
    const message = {
      role: "assistant",
      content: [{ type: "toolCall", id: "1", name: "approve", arguments: {} }],
      stopReason: "toolUse",
    } as unknown as AssistantMessage;
    const { authorize, log } = harness({
      produce: vi.fn(async () => message),
    });

    const verdict = await authorize(makePromptDetails(), query, log);

    expect(verdict.kind).toBe("defer");
    expect(failureCategory(log)).toBe("invalid-response");
  });

  it("does not retry a truncated response", async () => {
    const message = {
      role: "assistant",
      content: [{ type: "text", text: "partial" }],
      stopReason: "length",
    } as unknown as AssistantMessage;
    const { authorize, produce, log } = harness({
      produce: vi.fn(async () => message),
    });

    const verdict = await authorize(makePromptDetails(), query, log);

    expect(verdict.kind).toBe("defer");
    if (verdict.kind === "defer") {
      expect(verdict.note?.text).toContain("输出被截断");
    }
    expect(failureCategory(log)).toBe("truncated");
    expect(produce).toHaveBeenCalledTimes(1);
  });

  it("retries a provider error and allows on the second attempt", async () => {
    const errorMessage = {
      role: "assistant",
      content: [],
      stopReason: "error",
      errorMessage: "upstream 500",
    } as unknown as AssistantMessage;
    const produce = vi
      .fn()
      .mockResolvedValueOnce(errorMessage)
      .mockResolvedValue(verdictMessage("approve"));
    const { authorize, log } = harness({ produce, config: { maxAttempts: 2 } });

    const verdict = await authorize(makePromptDetails(), query, log);

    expect(verdict).toEqual({ kind: "allow" });
    expect(produce).toHaveBeenCalledTimes(2);
    expect(decisionFields(log)?.attempts).toBe(2);
  });

  it("defers with a timeout note when the call outlasts the budget", async () => {
    const produce = vi.fn(
      (_model: unknown, _context: unknown, options?: { signal?: AbortSignal }) =>
        new Promise<AssistantMessage>((resolve) => {
          options?.signal?.addEventListener(
            "abort",
            () =>
              resolve({
                role: "assistant",
                content: [],
                stopReason: "aborted",
              } as unknown as AssistantMessage),
            { once: true },
          );
        }),
    );
    const { authorize, log } = harness({
      produce,
      config: { timeoutMs: 20 },
    });

    const verdict = await authorize(makePromptDetails(), query, log);

    expect(verdict.kind).toBe("defer");
    if (verdict.kind === "defer") {
      expect(verdict.note?.text).toContain("超时");
    }
    expect(failureCategory(log)).toBe("timeout");
  });

  it("denies outright when the user aborts via ctx.signal", async () => {
    const controller = new AbortController();
    const produce = vi.fn(
      (_model: unknown, _context: unknown, options?: { signal?: AbortSignal }) =>
        new Promise<AssistantMessage>((resolve) => {
          options?.signal?.addEventListener(
            "abort",
            () =>
              resolve({
                role: "assistant",
                content: [],
                stopReason: "aborted",
              } as unknown as AssistantMessage),
            { once: true },
          );
        }),
    );
    const { authorize, log } = harness({
      produce,
      ctxSignal: controller.signal,
      config: { timeoutMs: 60_000 },
    });

    const pending = authorize(makePromptDetails(), query, log);
    controller.abort();
    const verdict = await pending;

    expect(verdict).toEqual({
      kind: "deny",
      reason: "Operation aborted by the user.",
    });
  });

  it("defers with the cancelled note when the session shuts down mid-call", async () => {
    const controller = new AbortController();
    const produce = vi.fn(
      (_model: unknown, _context: unknown, options?: { signal?: AbortSignal }) =>
        new Promise<AssistantMessage>((resolve) => {
          options?.signal?.addEventListener(
            "abort",
            () =>
              resolve({
                role: "assistant",
                content: [],
                stopReason: "aborted",
              } as unknown as AssistantMessage),
            { once: true },
          );
        }),
    );
    const { authorize, log } = harness({
      produce,
      sessionSignal: controller.signal,
      config: { timeoutMs: 60_000 },
    });

    const pending = authorize(makePromptDetails(), query, log);
    controller.abort();
    const verdict = await pending;

    expect(verdict.kind).toBe("defer");
    if (verdict.kind === "defer") {
      expect(verdict.note?.text).toBe("自动审查已取消，已转人工确认。");
    }
    expect(failureCategory(log)).toBe("cancelled");
  });

  it("passes reasoning only for reasoning-capable models, and never 'off'", async () => {
    const reasoningModel = makeModel("thinker", { reasoning: true });
    const { authorize, produce } = harness({
      sessionModel: reasoningModel,
      config: { reasoning: "low" },
    });
    await authorize(makePromptDetails({ command: "a" }), query, makeAuthorizerLog());
    expect(produce.mock.calls[0]?.[2]?.reasoning).toBe("low");

    const plain = makeModel("plain", { reasoning: false });
    const second = harness({ sessionModel: plain, config: { reasoning: "low" } });
    await second.authorize(makePromptDetails({ command: "b" }), query, makeAuthorizerLog());
    expect(second.produce.mock.calls[0]?.[2]).not.toHaveProperty("reasoning");

    const off = harness({ sessionModel: reasoningModel, config: { reasoning: "off" } });
    await off.authorize(makePromptDetails({ command: "c" }), query, makeAuthorizerLog());
    expect(off.produce.mock.calls[0]?.[2]).not.toHaveProperty("reasoning");
  });

  it("sends the review tools and no toolChoice", async () => {
    const { authorize, streamSimple, registry } = harness();

    await authorize(makePromptDetails(), query, makeAuthorizerLog());

    // The reviewer reaches the model through streamSimple(...).result() on
    // Pi 1.3.0's extension-facing registry — assert the call sees the same
    // model/context/options and `this` bound to the registry itself.
    expect(streamSimple).toHaveBeenCalledTimes(1);
    expect(streamSimple.mock.contexts[0]).toBe(registry);
    const [model, context, options] = streamSimple.mock.calls[0]! as [
      unknown,
      { systemPrompt: string; tools: unknown },
      Record<string, unknown>,
    ];
    expect(model).toBeDefined();
    expect(context.systemPrompt).toBe(REVIEW_SYSTEM_PROMPT);
    expect(context.tools).toBe(REVIEW_TOOLS);
    expect(options).not.toHaveProperty("toolChoice");
    expect(options.maxRetries).toBe(0);
    expect(options.maxTokens).toBe(4_000);
  });

  it("uses completeSimple directly when the registry exposes it", async () => {
    const { authorize, produce, streamSimple } = harness({
      completion: "complete",
    });

    const verdict = await authorize(makePromptDetails(), query, makeAuthorizerLog());

    expect(verdict).toEqual({ kind: "allow" });
    expect(produce).toHaveBeenCalledTimes(1);
    expect(streamSimple).not.toHaveBeenCalled();
  });

  it("caches approvals keyed to the action and latest trusted entry", async () => {
    const { authorize, produce, log } = harness();
    const details = makePromptDetails();

    await authorize(details, query, makeAuthorizerLog());
    await authorize(details, query, log);

    expect(produce).toHaveBeenCalledTimes(1);
    expect(decisionFields(log)?.cached).toBe(true);
  });

  it("invalidates the cache when a new trusted entry arrives", async () => {
    const branch = [...TRUSTED_BRANCH];
    const { authorize, produce } = harness({ branch });
    const details = makePromptDetails();

    await authorize(details, query, makeAuthorizerLog());
    branch.push(userEntry("e2", "now do something else"));
    await authorize(details, query, makeAuthorizerLog());

    expect(produce).toHaveBeenCalledTimes(2);
  });

  it("never caches a reject", async () => {
    const { authorize, produce } = harness({
      produce: vi.fn(async () => verdictMessage("reject", "no")),
    });
    const details = makePromptDetails();

    await authorize(details, query, makeAuthorizerLog());
    await authorize(details, query, makeAuthorizerLog());

    expect(produce).toHaveBeenCalledTimes(2);
  });
});