popagent

Latest change 2c7138e7ddf8cb6b5ac0dc1abbe6d76743366c19 - Checkpoint workspace and secret management work by AkurAI Build

import { describe, expect, test } from "bun:test";
import type { ChatMessage } from "./api-types";
import { latestUserText, replaceLatestUserText, sanitizeSensitiveToolMessages } from "./chat-messages";

describe("chat message transforms", () => {
  test("replaces every text fragment in the latest user message and preserves attachments", () => {
    const attachment = { type: "file", mediaType: "text/plain", filename: "notes.txt" };
    const toolPart = { type: "tool-result", toolCallId: "tool-1", output: "kept" };
    const messages: ChatMessage[] = [
      { id: "old", role: "user", parts: [{ type: "text", text: "older prompt" }] },
      {
        id: "latest",
        role: "user",
        parts: [
          attachment,
          { type: "text", text: "unsafe first fragment", providerMetadata: { source: "client" } },
          toolPart,
          { type: "text", text: "unsafe second fragment" },
        ],
      },
    ];

    expect(latestUserText(messages)).toBe("unsafe first fragment\nunsafe second fragment");
    replaceLatestUserText(messages, "safe replacement");

    expect(messages[0]?.parts).toEqual([{ type: "text", text: "older prompt" }]);
    expect(messages[1]?.parts).toEqual([
      attachment,
      {
        type: "text",
        text: "safe replacement",
        providerMetadata: { source: "client" },
      },
      toolPart,
    ]);
    expect(latestUserText(messages)).toBe("safe replacement");
  });

  test("joins all text parts to exactly the replacement", () => {
    const messages: ChatMessage[] = [{
      id: "multi", role: "user",
      parts: [{ type: "text", text: "one" }, { type: "text", text: "two" }],
    }];
    replaceLatestUserText(messages, "replacement");
    expect(latestUserText(messages)).toBe("replacement");
  });
});

describe("sensitive tool message sanitization", () => {
  test("redacts current and legacy plaintext secret tool payloads without mutating other tools", () => {
    const marker = `credential-${crypto.randomUUID()}`;
    const messages: ChatMessage[] = [{
      id: "assistant",
      role: "assistant",
      parts: [
        {
          type: "tool-recallSecret",
          toolCallId: "secret-current",
          state: "output-available",
          input: { name: "login/password" },
          output: { name: "login/password", value: marker },
        },
        {
          type: "tool-invocation",
          toolInvocation: {
            toolName: "store-secret",
            toolCallId: "secret-legacy",
            state: "result",
            args: { name: "login/password", value: marker },
            result: { stored: true },
          },
        },
        {
          type: "tool-webSearch",
          toolCallId: "search",
          state: "output-available",
          input: { query: "safe query" },
          output: { results: [] },
        },
      ],
    }];

    const sanitized = sanitizeSensitiveToolMessages(messages);

    expect(JSON.stringify(sanitized)).not.toContain(marker);
    expect(sanitized[0]?.parts[0]).toEqual(expect.objectContaining({
      type: "tool-recallSecret",
      input: "[redacted]",
      output: "[redacted]",
    }));
    expect(sanitized[0]?.parts[1]).toEqual(expect.objectContaining({
      toolInvocation: expect.objectContaining({
        toolName: "store-secret",
        args: "[redacted]",
        result: "[redacted]",
      }),
    }));
    expect(sanitized[0]?.parts[2]).toEqual(messages[0]?.parts[2]);
    expect(JSON.stringify(messages)).toContain(marker);
  });
});