fix(studio-server): undo keeps an outside change even when Studio's history claim arrives late (#4631)

* test(studio-server): a write landing between Studio's read and its patch survives undo

* test(studio-server): the outside write also survives a claim that arrives after receipts expire

* fix(studio-server): undo keeps an outside change even when Studio's history claim arrives late

* fix(studio-server): history forgets replaced bytes on delete and frees unused bytes past budget

Also recreates the history folder before storing replaced bytes, hears writes inside the queued start, and pins the agent-turn claim that arrives after the receipt expired.

Author: miguel-heygen

packages/studio-server/src/helpers/fileVersion.ts

@@ -9,12 +9,13 @@ export interface FileWriteReceipt {
 
 interface StoredReceipt extends FileWriteReceipt {
   recordedAt: number;
-  /** The bytes this write replaced, so the project history can keep a save that landed just before it. */
-  overwrote?: string | Uint8Array;
 }
 
+type OverwriteListener = (absPath: string, version: string, overwrote: string | Uint8Array) => void;
+
 const RECEIPT_TTL_MS = 10_000;
 const receipts = new Map<string, StoredReceipt[]>();
+const overwriteListeners = new Set<OverwriteListener>();
 
 /** Strong content version used as both the JSON version and HTTP ETag. */
 export function fileContentVersion(content: string | Uint8Array): string {
@@ -50,13 +51,20 @@ export function createWriteToken(requestToken?: string): string {
   return token && token.length <= 200 ? token : randomUUID();
 }
 
+/** Hears the bytes each API write replaced, so a project history can keep a save it never saw. */
+export function onFileOverwritten(listener: OverwriteListener): () => void {
+  overwriteListeners.add(listener);
+  return () => overwriteListeners.delete(listener);
+}
+
 export function recordFileWriteReceipt(
   filePath: string,
-  receipt: Omit<StoredReceipt, "recordedAt">,
+  { overwrote, ...receipt }: FileWriteReceipt & { overwrote?: string | Uint8Array },
 ): void {
   const absPath = realFilePath(filePath);
+  if (overwrote !== undefined)
+    for (const listener of overwriteListeners) listener(absPath, receipt.version, overwrote);
   const now = Date.now();
-  // Every path's expired receipts go, not just this one's: a receipt can hold a whole file's bytes.
   for (const [path, list] of receipts) {
     const live = list.filter((entry) => now - entry.recordedAt < RECEIPT_TTL_MS);
     if (live.length > 0) receipts.set(path, live);
@@ -92,20 +100,6 @@ export function identifyFileWrite(
   return { path, version, writeToken };
 }
 
-/** The bytes the API write of `version` replaced, while its receipt lives. */
-export function bytesOverwrittenBy(
-  filePath: string,
-  version: string,
-): string | Uint8Array | undefined {
-  return newestReceipt(realFilePath(filePath), version)?.overwrote;
-}
-
-/** Drops the replaced bytes a claim walked through, so a later claim can't walk back through them. */
-export function forgetOverwrittenBytes(filePath: string, versions: ReadonlySet<string>): void {
-  for (const receipt of receipts.get(realFilePath(filePath)) ?? [])
-    if (versions.has(receipt.version)) delete receipt.overwrote;
-}
-
 function newestReceipt(absPath: string, expectedVersion: string): StoredReceipt | undefined {
   const now = Date.now();
   const current = (receipts.get(absPath) ?? []).filter(

packages/studio-server/src/history/projectHistory.test.ts

@@ -18,7 +18,7 @@ import { spawnSync } from "node:child_process";
 import { tmpdir } from "node:os";
 import { basename, dirname, join, relative } from "node:path";
 import { afterEach, describe, expect, it, vi } from "vitest";
-import { fileContentVersion, recordFileWriteReceipt } from "../helpers/fileVersion";
+import { fileContentVersion, hashOfVersion, recordFileWriteReceipt } from "../helpers/fileVersion";
 import { HistoryBusyError } from "./ownerLock";
 import { HistoryIdError } from "./historyId";
 import { HistoryClosedError, openProjectHistory, type ProjectHistory } from "./projectHistory";
@@ -1099,7 +1099,10 @@ describe("openProjectHistory", () => {
     expect((await window.close())?.id).toBe(window.id);
   });
 
-  it("keeps the bytes a claim cut at through a budget fold that runs before they are logged", async () => {
+  it.each([
+    ["while the claim runs", Infinity],
+    ["before the claim comes", 1],
+  ])("keeps the bytes a claim cut at through a budget fold %s", async (_when, dragIdleMs) => {
     const saved = "B".repeat(3000);
     const { history, write, read, projectDir } = await project(
       { "index.html": "a", "other.html": "o" },
@@ -1108,7 +1111,7 @@ describe("openProjectHistory", () => {
     await change(history, you, "Old", () => write("other.html", "o2"));
     history.pin((await change(history, you, "Pinned", () => write("other.html", "o3"))).id, true);
     write("other.html", "o4");
-    await history.claim(you, "Drag", ["other.html"], { coalesceKey: "drag" });
+    await history.claim(you, "Drag", ["other.html"], { coalesceKey: "drag", idleMs: dragIdleMs });
     // Studio's write of `edited` replaced a save it never read.
     const edited = `${saved}!`;
     write("index.html", edited);
@@ -1118,6 +1121,7 @@ describe("openProjectHistory", () => {
       writeToken: "studio",
       overwrote: saved,
     });
+    await new Promise((settle) => setTimeout(settle, 20));
 
     const edit = await history.claim(you, "Edit", ["index.html"], {
       overwrote: { "index.html": fileContentVersion("a") },
@@ -1126,6 +1130,66 @@ describe("openProjectHistory", () => {
     expect(read("index.html")).toBe(saved);
   });
 
+  it("keeps the bytes an API write replaced when its history folder was removed while open", async () => {
+    const { history, write, read, projectDir, historyRoot } = await project({ "index.html": "a" });
+    rmSync(historyRoot, { recursive: true, force: true });
+    write("index.html", "E");
+    recordFileWriteReceipt(join(projectDir, "index.html"), {
+      path: "index.html",
+      version: fileContentVersion("E"),
+      writeToken: "studio",
+      overwrote: "S",
+    });
+
+    const edit = await history.claim(you, "Edit", ["index.html"], {
+      overwrote: { "index.html": fileContentVersion("a") },
+    });
+    expect((await history.undo(edit!.id, { who: you })).ok).toBe(true);
+    expect(read("index.html")).toBe("S");
+  });
+
+  it("forgets what an API write replaced once the file is removed", async () => {
+    const { history, write, read, projectDir } = await project({ "index.html": "a" });
+    write("index.html", "E");
+    recordFileWriteReceipt(join(projectDir, "index.html"), {
+      path: "index.html",
+      version: fileContentVersion("E"),
+      writeToken: "studio",
+      overwrote: "S",
+    });
+    rmSync(join(projectDir, "index.html"));
+    await history.claim(you, "sweep", []);
+    // Re-created outside Studio with the bytes the API wrote: nothing Studio wrote is on disk now.
+    write("index.html", "E");
+
+    const edit = await history.claim(you, "Edit", ["index.html"], {
+      overwrote: { "index.html": fileContentVersion("a") },
+    });
+    expect((await history.undo(edit!.id, { who: you })).ok).toBe(true);
+    expect(read("index.html")).toBe("a");
+  });
+
+  it("frees unreferenced bytes past its budget even when nothing is left to fold", async () => {
+    const saved = "B".repeat(3000);
+    const { history, write, projectDir } = await project(
+      { "index.html": "a", "other.html": "o" },
+      { budgetBytes: 2000 },
+    );
+    history.pin((await change(history, you, "Pinned", () => write("other.html", "o2"))).id, true);
+    const edited = `${saved}!`;
+    write("index.html", edited);
+    recordFileWriteReceipt(join(projectDir, "index.html"), {
+      path: "index.html",
+      version: fileContentVersion(edited),
+      writeToken: "agent",
+      overwrote: saved,
+    });
+    write("index.html", "x");
+    await history.flush();
+
+    await expect(history.readBlob(hashOfVersion(fileContentVersion(saved))!)).rejects.toThrow();
+  });
+
   it("cuts a claim at a restored save, not at bytes an earlier claim already used", async () => {
     const { history, write, read, projectDir } = await project({ "index.html": "W" });
     const studioWrites = (content: string, overwrote: string) => {
@@ -1153,7 +1217,7 @@ describe("openProjectHistory", () => {
     expect(read("index.html")).toBe("X");
   });
 
-  it("keeps what a Studio write landing during a claim replaced, for the next claim", async () => {
+  it("keeps what a Studio write replaced over an unseen save, for the next claim", async () => {
     const { history, write, read, projectDir } = await project({ "index.html": "W" });
     const receipt = (content: string, overwrote: string) =>
       recordFileWriteReceipt(join(projectDir, "index.html"), {
@@ -1167,16 +1231,91 @@ describe("openProjectHistory", () => {
     });
     write("index.html", "X");
     receipt("X", "W");
-    // Studio's next write, over an editor's save B, lands after this claim's scan.
-    receipt("Y", "B");
     await history.claim(you, "First", ["index.html"], told("W"));
+    // Studio's next write lands over an editor's save B that the history never saw.
     write("index.html", "Y");
+    receipt("Y", "B");
 
     const second = await history.claim(you, "Second", ["index.html"], told("X"));
     expect((await history.undo(second!.id, { who: you })).ok).toBe(true);
     expect(read("index.html")).toBe("B");
   });
 
+  it.each([
+    ["once its change is committed", true],
+    ["once an editor wrote over it", false],
+  ])(
+    "forgets what an API write replaced %s, so a later claim cannot cut at it",
+    async (_when, commit) => {
+      const { history, write, read, projectDir } = await project({ "index.html": "W" });
+      write("index.html", "X");
+      recordFileWriteReceipt(join(projectDir, "index.html"), {
+        path: "index.html",
+        version: fileContentVersion("X"),
+        writeToken: "agent",
+        overwrote: "W",
+      });
+      if (commit) await history.flush();
+      write("index.html", "C");
+      await history.flush();
+      // An editor, not the API, brings X back.
+      write("index.html", "X");
+
+      const edit = await history.claim(you, "Edit", ["index.html"], {
+        overwrote: { "index.html": fileContentVersion("C") },
+      });
+      expect((await history.undo(edit!.id, { who: you })).ok).toBe(true);
+      expect(read("index.html")).toBe("C");
+    },
+  );
+
+  it("cuts at an editor's save, not at what a rolled-back API write replaced", async () => {
+    const { history, write, read, projectDir } = await project({ "index.html": "W" });
+    const apiWrites = (content: string, overwrote?: string) => {
+      write("index.html", content);
+      recordFileWriteReceipt(join(projectDir, "index.html"), {
+        path: "index.html",
+        version: fileContentVersion(content),
+        writeToken: "api",
+        overwrote,
+      });
+    };
+    apiWrites("X", "W");
+    apiWrites("W");
+    apiWrites("K", "W");
+    await history.claim(you, "First", ["index.html"], {
+      overwrote: { "index.html": fileContentVersion("W") },
+    });
+    write("index.html", "X");
+    apiWrites("P", "X");
+
+    const second = await history.claim(you, "Second", ["index.html"], {
+      overwrote: { "index.html": fileContentVersion("K") },
+    });
+    expect((await history.undo(second!.id, { who: you })).ok).toBe(true);
+    expect(read("index.html")).toBe("X");
+  });
+
+  it("lets the budget free the bytes a committed API write replaced", async () => {
+    const saved = "B".repeat(3000);
+    const { history, write, projectDir } = await project(
+      { "index.html": "a", "other.html": "o" },
+      { budgetBytes: 2000 },
+    );
+    const edited = `${saved}!`;
+    write("index.html", edited);
+    recordFileWriteReceipt(join(projectDir, "index.html"), {
+      path: "index.html",
+      version: fileContentVersion(edited),
+      writeToken: "agent",
+      overwrote: saved,
+    });
+    await history.flush();
+
+    await change(history, you, "Other", () => write("other.html", "o2"));
+    expect(history.list().map((entry) => entry.label)).toContain("Other");
+  });
+
   it("keeps the history of a small edit in a project larger than its budget", async () => {
     const { history, write } = await project(
       { "media.bin": Buffer.alloc(4096, 1), "index.html": "a" },