211 lines · 8.5 KB · TypeScript
RawDownload
import { afterEach, beforeEach, describe, expect, test } from "bun:test";import { appendFileSync, mkdtempSync, rmSync, writeFileSync, truncateSync } from "node:fs";import { tmpdir } from "node:os";import path from "node:path";import { readSessionUsageTotals, type SessionUsageTotals } from "../src/session-usage"; const temporaryRoots: string[] = [];const originalDataDir = process.env.LOCAL_STUDIO_DATA_DIR; // The usage cache is backed by a directory under the resolved data dir. Without// this the suite writes cache entries into the developer's real ~/.local-studio.beforeEach(() => {  const dataDir = mkdtempSync(path.join(tmpdir(), "session-usage-data-"));  temporaryRoots.push(dataDir);  process.env.LOCAL_STUDIO_DATA_DIR = dataDir;}); afterEach(() => {  if (originalDataDir === undefined) delete process.env.LOCAL_STUDIO_DATA_DIR;  else process.env.LOCAL_STUDIO_DATA_DIR = originalDataDir;  for (const root of temporaryRoots.splice(0)) rmSync(root, { recursive: true, force: true });}); function fixture(): string {  const root = mkdtempSync(path.join(tmpdir(), "session-usage-"));  temporaryRoots.push(root);  return path.join(root, "rollout.jsonl");} const header = () => JSON.stringify({ type: "session", id: "s1", cwd: "/tmp" }); /** One assistant turn that cost `input`/`output` tokens. */const turn = (input: number, output: number) =>  JSON.stringify({    type: "message",    message: { role: "assistant", usage: { input, output, cost: { total: 0.5 } } },  }); /** A line with no usage block — the pre-filter should skip it entirely. */const chatter = (text: string) =>  JSON.stringify({ type: "message", message: { role: "user", content: [{ type: "text", text }] } }); // Resuming a scan means trusting a byte offset. Everything below is about that// offset being exactly right, because being one byte off does not throw — it// silently reports a wrong lifetime spend.describe("incremental usage scan", () => {  test("totals a whole rollout on the first read", async () => {    const file = fixture();    writeFileSync(file, `${header()}\n${turn(100, 20)}\n${turn(200, 30)}\n`);     const totals = await readSessionUsageTotals(file);     expect(totals.input).toBe(300);    expect(totals.output).toBe(50);    expect(totals.calls).toBe(2);  });   test("an appended turn adds to the total instead of restarting it", async () => {    const file = fixture();    writeFileSync(file, `${header()}\n${turn(100, 20)}\n`);    expect((await readSessionUsageTotals(file)).input).toBe(100);     appendFileSync(file, `${turn(400, 60)}\n`);     const totals = await readSessionUsageTotals(file);    expect(totals.input).toBe(500);    expect(totals.output).toBe(80);    expect(totals.calls).toBe(2);  });   test("a half-written final line is re-read once it is complete", async () => {    // The writer is mid-append when we scan: the last line has no newline yet.    // Counting it as scanned would lose that turn forever.    const file = fixture();    const partial = turn(700, 90);    writeFileSync(file, `${header()}\n${turn(100, 20)}\n${partial.slice(0, 30)}`);    expect((await readSessionUsageTotals(file)).input).toBe(100);     writeFileSync(file, `${header()}\n${turn(100, 20)}\n${partial}\n`);     const totals = await readSessionUsageTotals(file);    expect(totals.input).toBe(800);    expect(totals.calls).toBe(2);  });   test("multi-byte turns do not drift the resume offset", async () => {    // Character offsets and byte offsets diverge the moment a turn contains    // anything non-ASCII, and the resume point is a byte offset.    const file = fixture();    writeFileSync(file, `${header()}\n${chatter("绿茶 — ☕️ naïve")}\n${turn(100, 20)}\n`);    expect((await readSessionUsageTotals(file)).input).toBe(100);     appendFileSync(file, `${chatter("مرحبا 🌍")}\n${turn(250, 25)}\n`);     const totals = await readSessionUsageTotals(file);    expect(totals.input).toBe(350);    expect(totals.output).toBe(45);    expect(totals.calls).toBe(2);  });   test("a rewritten file is rescanned rather than resumed", async () => {    // The resume is only sound while the file is append-only. If it is replaced,    // the cached prefix belongs to a different session.    const file = fixture();    writeFileSync(file, `${header()}\n${turn(100, 20)}\n`);    expect((await readSessionUsageTotals(file)).input).toBe(100);     const replacement = JSON.stringify({ type: "session", id: "s2", cwd: "/tmp/other" });    writeFileSync(file, `${replacement}\n${turn(7, 3)}\n`);     const totals = await readSessionUsageTotals(file);    expect(totals.input).toBe(7);    expect(totals.calls).toBe(1);  });   test("a truncated file is rescanned rather than resumed", async () => {    const file = fixture();    writeFileSync(file, `${header()}\n${turn(100, 20)}\n${turn(200, 30)}\n`);    expect((await readSessionUsageTotals(file)).input).toBe(300);     const shorter = `${header()}\n${turn(100, 20)}\n`;    truncateSync(file, Buffer.byteLength(shorter, "utf-8"));     expect((await readSessionUsageTotals(file)).input).toBe(100);  });   test("compactions keep counting across a resume", async () => {    const file = fixture();    writeFileSync(file, `${header()}\n${JSON.stringify({ type: "compaction" })}\n`);    expect((await readSessionUsageTotals(file)).compactions).toBe(1);     appendFileSync(file, `${JSON.stringify({ type: "compaction" })}\n${turn(10, 5)}\n`);     const totals = await readSessionUsageTotals(file);    expect(totals.compactions).toBe(2);    expect(totals.input).toBe(10);  });}); // --- surviving a restart ----------------------------------------------------- /** * The in-process memo dies with the controller. Backing it on disk is only * worth anything if a *fresh process* can pick up where the last one stopped, * so this actually spawns one rather than clearing a module-level Map. */async function scanInFreshProcess(file: string, dataDir: string): Promise<SessionUsageTotals> {  const script = `    const { readSessionUsageTotals } = await import(${JSON.stringify(      path.resolve(import.meta.dir, "../src/session-usage.ts"),    )});    const started = performance.now();    const totals = await readSessionUsageTotals(${JSON.stringify(file)});    console.log(JSON.stringify({ ...totals, ms: performance.now() - started }));  `;  const proc = Bun.spawn(["bun", "-e", script], {    env: { ...process.env, LOCAL_STUDIO_DATA_DIR: dataDir },    stdout: "pipe",    stderr: "pipe",  });  const out = await new Response(proc.stdout).text();  const code = await proc.exited;  if (code !== 0) throw new Error(await new Response(proc.stderr).text());  // The runtime may print its own notices before ours; the result is the last  // line we wrote.  const lastLine = out.trim().split("\n").at(-1) ?? "";  return JSON.parse(lastLine) as SessionUsageTotals;} describe("usage cache across processes", () => {  test("a second process reuses the first one's scan", async () => {    const file = fixture();    const dataDir = process.env.LOCAL_STUDIO_DATA_DIR as string;    writeFileSync(file, `${header()}\n${turn(100, 20)}\n${turn(200, 30)}\n`);     const first = await scanInFreshProcess(file, dataDir);    expect(first.input).toBe(300);     const second = await scanInFreshProcess(file, dataDir);    expect(second.input).toBe(300);    expect(second.calls).toBe(2);  });   test("a second process resumes rather than rescanning a grown file", async () => {    const file = fixture();    const dataDir = process.env.LOCAL_STUDIO_DATA_DIR as string;    writeFileSync(file, `${header()}\n${turn(100, 20)}\n`);    expect((await scanInFreshProcess(file, dataDir)).input).toBe(100);     appendFileSync(file, `${turn(400, 60)}\n`);     // If the restored prefix were ignored the total would still be 500 — the    // number that proves resumption is the byte offset, so assert the total is    // right AND that the first turn was not double-counted.    const totals = await scanInFreshProcess(file, dataDir);    expect(totals.input).toBe(500);    expect(totals.calls).toBe(2);  });   test("a rewritten file is not resumed across processes either", async () => {    const file = fixture();    const dataDir = process.env.LOCAL_STUDIO_DATA_DIR as string;    writeFileSync(file, `${header()}\n${turn(100, 20)}\n`);    expect((await scanInFreshProcess(file, dataDir)).input).toBe(100);     const replacement = JSON.stringify({ type: "session", id: "s2", cwd: "/tmp/other" });    writeFileSync(file, `${replacement}\n${turn(7, 3)}\n`);     expect((await scanInFreshProcess(file, dataDir)).input).toBe(7);  });});