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

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

import type { SessionEntry } from "@earendil-works/pi-coding-agent";
import {
  collectTranscriptEntries,
  estimateTokens,
  renderTranscript,
  truncateToApproximateTokens,
} from "#src/auto-review/transcript";

function entry(
  id: string,
  message: Record<string, unknown>,
): SessionEntry {
  return { type: "message", id, message } as unknown as SessionEntry;
}

function user(id: string, text: string): SessionEntry {
  return entry(id, { role: "user", content: text });
}

function toolResult(
  id: string,
  toolCallId: string,
  text: string,
  extras: Record<string, unknown> = {},
): SessionEntry {
  return entry(id, {
    role: "toolResult",
    toolName: "read",
    toolCallId,
    isError: false,
    content: [{ type: "text", text }],
    ...extras,
  });
}

const renderedContent = (line: string) => JSON.parse(line).content as string;
const renderedLabel = (line: string) => JSON.parse(line).label as string;

describe("estimateTokens", () => {
  it("counts ASCII at a quarter token per character", () => {
    expect(estimateTokens("abcd")).toBe(1);
    expect(estimateTokens("abcde")).toBe(2);
  });

  it("counts every non-ASCII code point as a full token", () => {
    expect(estimateTokens("你好")).toBe(2);
    expect(estimateTokens("ab你好cd")).toBe(1 + 2); // ceil(4/4) + 2
  });
});

describe("truncateToApproximateTokens", () => {
  it("bounds the estimated tokens of the rendered result", () => {
    const text = "汉".repeat(5_000);
    const truncated = truncateToApproximateTokens(text, 100);
    expect(estimateTokens(truncated)).toBeLessThanOrEqual(100);
    expect(truncated).toContain("...[truncated]...");
  });
});

describe("renderTranscript summary substitution", () => {
  it("replaces a tool result with a strictly shorter summary", () => {
    const original = "x".repeat(500);
    const branch = [
      user("e1", "do it"),
      toolResult("e2", "call-1", original),
    ];
    const querySummaries = vi.fn(() => [
      { toolCallId: "call-1", deterministicSummary: "short summary" },
    ]);

    const result = renderTranscript(branch, { querySummaries });

    expect(querySummaries).toHaveBeenCalledTimes(1);
    expect(querySummaries).toHaveBeenCalledWith(["call-1"]);
    const toolLine = result.entries.find((line) =>
      renderedLabel(line).includes("read"),
    );
    expect(toolLine).toBeDefined();
    expect(renderedLabel(toolLine!)).toContain("(summarized)");
    expect(renderedContent(toolLine!)).toBe("short summary\n[summarized]");
    expect(result.stats.toolResultsSummarized).toBe(1);
  });

  it("prefers the semantic summary over the deterministic one", () => {
    const branch = [user("e1", "go"), toolResult("e2", "call-1", "y".repeat(300))];
    const result = renderTranscript(branch, {
      querySummaries: () => [
        {
          toolCallId: "call-1",
          semanticSummary: "semantic",
          deterministicSummary: "deterministic",
        },
      ],
    });

    const toolLine = result.entries.find((line) =>
      renderedLabel(line).includes("summarized"),
    );
    expect(renderedContent(toolLine!)).toBe("semantic\n[summarized]");
  });

  it("keeps the original when the summary is not shorter", () => {
    const branch = [user("e1", "go"), toolResult("e2", "call-1", "tiny")];
    const result = renderTranscript(branch, {
      querySummaries: () => [
        { toolCallId: "call-1", deterministicSummary: "a much longer summary that exceeds the original text" },
      ],
    });

    const toolLine = result.entries.find((line) =>
      renderedLabel(line).includes("read"),
    );
    expect(renderedLabel(toolLine!)).not.toContain("summarized");
    expect(renderedContent(toolLine!)).toBe("tiny");
    expect(result.stats.toolResultsSummarized).toBe(0);
  });

  it("matches a record by resultTimestamp when both timestamps exist", () => {
    const branch = [
      user("e1", "go"),
      toolResult("e2", "call-1", "z".repeat(300), { timestamp: 111 }),
    ];
    const result = renderTranscript(branch, {
      querySummaries: () => [
        { toolCallId: "call-1", resultTimestamp: 999, deterministicSummary: "wrong" },
        { toolCallId: "call-1", resultTimestamp: 111, deterministicSummary: "right" },
      ],
    });

    const toolLine = result.entries.find((line) =>
      renderedLabel(line).includes("summarized"),
    );
    expect(renderedContent(toolLine!)).toBe("right\n[summarized]");
  });

  it("never summarizes a user_interaction result or an assistant toolCall argument", () => {
    const branch = [
      user("e1", "go"),
      entry("e2", {
        role: "assistant",
        content: [
          {
            type: "toolCall",
            id: "ask-1",
            name: "ask_user_question",
            arguments: { question: "proceed?" },
          },
        ],
      }),
      entry("e3", {
        role: "toolResult",
        toolName: "ask_user_question",
        toolCallId: "ask-1",
        isError: false,
        details: {
          cancelled: false,
          answers: [{ question: "proceed?", answer: "yes" }],
        },
        content: [{ type: "text", text: "answered" }],
      }),
    ];
    const querySummaries = vi.fn(() => [
      { toolCallId: "ask-1", deterministicSummary: "s" },
    ]);

    const result = renderTranscript(branch, { querySummaries });

    // The user_interaction entry carries no toolCallId, so the lookup is not
    // even consulted for it; no (summarized) label may appear.
    expect(result.entries.some((line) => renderedLabel(line).includes("summarized"))).toBe(false);
    expect(result.stats.toolResultsSummarized).toBe(0);
  });

  it("queries once for all tool results together", () => {
    const branch = [
      user("e1", "go"),
      toolResult("e2", "call-1", "a".repeat(300)),
      toolResult("e3", "call-2", "b".repeat(300)),
    ];
    const querySummaries = vi.fn(() => []);

    renderTranscript(branch, { querySummaries });

    expect(querySummaries).toHaveBeenCalledTimes(1);
    expect(querySummaries).toHaveBeenCalledWith(["call-1", "call-2"]);
  });

  it("keeps originals when the lookup throws", () => {
    const branch = [user("e1", "go"), toolResult("e2", "call-1", "c".repeat(300))];
    const result = renderTranscript(branch, {
      querySummaries: () => {
        throw new Error("service down");
      },
    });

    const toolLine = result.entries.find((line) =>
      renderedLabel(line).includes("read"),
    );
    expect(renderedContent(toolLine!)).toBe("c".repeat(300));
    expect(result.stats.toolResultsSummarized).toBe(0);
  });

  it("works without a lookup", () => {
    const branch = [user("e1", "go"), toolResult("e2", "call-1", "d".repeat(300))];

    const result = renderTranscript(branch);

    const toolLine = result.entries.find((line) =>
      renderedLabel(line).includes("read"),
    );
    expect(renderedContent(toolLine!)).toBe("d".repeat(300));
    expect(result.stats.toolResultsSummarized).toBe(0);
  });
});

describe("renderTranscript budgets", () => {
  it("honours maxUntrustedEntries", () => {
    const branch = [
      user("e1", "go"),
      toolResult("e2", "call-1", "one"),
      toolResult("e3", "call-2", "two"),
      toolResult("e4", "call-3", "three"),
    ];

    const result = renderTranscript(branch, {
      budgets: { messageTokens: 10_000, toolTokens: 10_000, maxUntrustedEntries: 2 },
    });

    const toolLines = result.entries.filter((line) =>
      renderedLabel(line).includes("read"),
    );
    expect(toolLines).toHaveLength(2);
    // Newest first: the two most recent results are retained.
    expect(result.entries.some((line) => renderedContent(line) === "three")).toBe(true);
    expect(result.entries.some((line) => renderedContent(line) === "one")).toBe(false);
  });

  it("always keeps the latest trusted user entry", () => {
    const branch = [
      user("e1", "first instruction " + "w".repeat(400)),
      toolResult("e2", "call-1", "x".repeat(2_000)),
      user("e3", "latest instruction"),
    ];

    const result = renderTranscript(branch, {
      budgets: { messageTokens: 50, toolTokens: 50, maxUntrustedEntries: 40 },
    });

    expect(result.entries.some((line) => renderedContent(line) === "latest instruction")).toBe(true);
    expect(result.stats.latestTrustedEntryRetained).toBe(true);
    expect(result.latestTrustedEntryId).toBe("e3");
  });
});

describe("collectTranscriptEntries", () => {
  it("emits user entries with their SessionEntry id", () => {
    const entries = collectTranscriptEntries([user("e1", "hello")]);

    expect(entries).toHaveLength(1);
    expect(entries[0]).toMatchObject({ kind: "user", entryId: "e1" });
  });
});