Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
5.5 kB · 129 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130import { existsSync, readFileSync, rmSync, statSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import { fauxAssistantMessage } from "@earendil-works/pi-ai";import { describe, expect, test } from "vitest";import { collectMarkedPastes, evictPastesToBudget, formatPastedMarker, formatPastedStub, type PastedEntry, reapSessionPastes, remapPastedMarkers, resolvePastesInText, toggledPastedMode, writePastedFile,} from "../src/modes/interactive/paste-files.js";import { createHarness, getUserTexts } from "./suite/harness.js";
describe("pasted-file markers pure logic", () => { test("collectMarkedPastes resolves only markers still present, in paste order (delete drops payload)", () => { const pending = new Map<number, string>([ [1, "first"], [2, "second"], ]); expect(collectMarkedPastes(pending, "[pasted #2] then [pasted #1]")).toEqual(["first", "second"]); expect(collectMarkedPastes(pending, "kept [pasted #1] only")).toEqual(["first"]); expect(collectMarkedPastes(pending, "deleted marker drops payload")).toEqual([]); });
test("a restored marker still resolves its pasted file (undo-safe)", () => { const pending = new Map<number, string>([[1, "payload"]]); expect(collectMarkedPastes(pending, "no marker")).toEqual([]); expect(collectMarkedPastes(pending, "back [pasted #1]")).toEqual(["payload"]); });
test("toggledPastedMode cycles a live pasted file so alt+k can reach secret mode", () => { expect(toggledPastedMode("blob")).toBe("secret"); expect(toggledPastedMode("secret")).toBe("blob"); });
test("remapPastedMarkers updates matching marker ids", () => { expect(remapPastedMarkers("[pasted #1] [pasted #2]", new Map([[1, 10]]))).toBe("[pasted #10] [pasted #2]"); });
test("writePastedFile writes atomic file with mode 0600 and directory 0700", () => { const tempDir = join(tmpdir(), `pasted-direct-${Date.now()}`); const entry: PastedEntry = { id: 42, content: "direct payload", mode: "blob" }; const res = writePastedFile(entry, { sessionId: "s1", agentDir: tempDir }); expect(existsSync(res.path)).toBe(true); expect(readFileSync(res.path, "utf-8")).toBe("direct payload"); expect(statSync(res.path).mode & 0o777).toBe(0o600); rmSync(tempDir, { recursive: true, force: true }); });
test("formatPastedStub is one line with the path and size, tagging secret", () => { const info = { path: "/tmp/test.txt", bytes: 1234, lines: 45, sha256: "abc123", mode: "blob" as const, }; expect(formatPastedStub(info)).toBe("[pasted: /tmp/test.txt · 1234 B]\n"); expect(formatPastedStub({ ...info, mode: "secret" })).toBe("[pasted secret: /tmp/test.txt · 1234 B]\n"); });
const size = (s: string) => s.length;
test.each([ { name: "drops oldest entries until within budget", budget: 8, kept: [] as number[], expected: [2, 3] }, { name: "never evicts kept ids, even past budget", budget: 8, kept: [1], expected: [1, 3] }, { name: "keeps everything when all ids are kept", budget: 1, kept: [1, 2, 3], expected: [1, 2, 3] }, { name: "is a no-op within budget", budget: 100, kept: [], expected: [1, 2, 3] }, ])("evictPastesToBudget $name", ({ budget, kept, expected }) => { const pastes = new Map([ [1, "aaaa"], [2, "bbbb"], [3, "cccc"], ]); evictPastesToBudget(pastes, size, budget, new Set(kept)); expect([...pastes.keys()]).toEqual(expected); });});
describe("pasted files and session transcript exclusion", () => { test("paste resolution writes 0600 file, emits stub, excludes payload from session jsonl, and reaps", async () => { const harness = await createHarness({ persistSession: true }); try { const sessionId = harness.session.sessionId; const agentDir = harness.tempDir; const payload = "SECRET_TOKEN_XYZ_987654321_DO_NOT_LEAK"; const pastes = new Map<number, PastedEntry>([[1, { id: 1, content: payload, mode: "secret" }]]);
const input = `check secret: ${formatPastedMarker(1)}`; const resolved = resolvePastesInText(input, pastes, { sessionId, agentDir });
// (a) file exists at expected path with byte-identical content and mode 0600 const written = resolved.writtenFiles[0]!; expect(written).toBeDefined(); expect(existsSync(written.path)).toBe(true); expect(readFileSync(written.path, "utf-8")).toBe(payload); expect(statSync(written.path).mode & 0o777).toBe(0o600);
// (b) user text session receives contains stub and not payload harness.setResponses([fauxAssistantMessage("acknowledged")]); await harness.session.prompt(resolved.text); const userTexts = getUserTexts(harness); expect(userTexts[0]).toContain(written.path); expect(userTexts[0]).not.toContain(payload);
// (c) persisted session jsonl contains payload string nowhere (the regression test) const sessionJsonl = readFileSync(harness.session.sessionFile!, "utf-8"); expect(sessionJsonl.includes(payload)).toBe(false);
// (d) the secret tag is what the model sees, and the reading rules live in the system prompt expect(userTexts[0]).toContain("pasted secret:"); expect(userTexts[0].trimEnd().split("\n")).toHaveLength(1);
// Reaper: mid-session file is retained; reap removes directory expect(existsSync(written.path)).toBe(true); const reaped = reapSessionPastes(sessionId, agentDir); expect(reaped.length).toBe(1); expect(existsSync(written.path)).toBe(false); } finally { harness.cleanup(); } });});