Restore archive leaf entries after failed compaction

Author: TheGreatAxios

src/session/optimized-context-store.test.ts

@@ -1,4 +1,4 @@
-import { describe, test, expect } from "bun:test";
+import { describe, test, expect, spyOn } from "bun:test";
 import fs from "node:fs";
 import os from "node:os";
 import path from "node:path";
@@ -616,6 +616,188 @@ async function gitLsTree(dir: string): Promise<string[]> {
 }
 
 describe("createOptimizedContextStore unpublished rewrite", () => {
+  test("partial archive staging and failed rollback fence shared audit publication until reconciliation succeeds", async () => {
+    const dir = tempDir();
+    const { storage } = await createSessionStores(dir);
+    const archiveDir = path.join(dir, "evidence-archive");
+    fs.mkdirSync(archiveDir, { recursive: true });
+    fs.writeFileSync(path.join(archiveDir, "index.jsonl"), "original");
+    await storage.writeTurns([turn("old-a"), turn("old-b")]);
+    await storage.writeMetadata(EMPTY_CHECKPOINT_METADATA);
+    const published = await storage.commit({ message: "original" });
+    const { audit } = await createSessionStores(dir);
+    fs.writeFileSync(path.join(archiveDir, "index.jsonl"), "changed");
+    fs.writeFileSync(path.join(archiveDir, "added.jsonl"), "new");
+    await storage.writeTurns([turn("[Compacted prior context]")]);
+    const originalAdd = git.add;
+    const originalReset = git.resetIndex;
+    let resetFails = true;
+    const staged: string[] = [];
+    const addSpy = spyOn(git, "add").mockImplementation(async (args) => {
+      if (args.dir === dir && args.filepath === "evidence-archive/index.jsonl")
+        throw new Error("partial staging failure");
+      await originalAdd(args);
+      if (args.dir === dir)
+        staged.push(...(Array.isArray(args.filepath) ? args.filepath : [args.filepath]));
+    });
+    const resetSpy = spyOn(git, "resetIndex").mockImplementation(async (args) => {
+      if (args.dir === dir && args.filepath.startsWith("evidence-archive/") && resetFails)
+        throw new Error("injected reset failure");
+      return originalReset(args);
+    });
+    const errorRecord = {
+      source: "reactor" as const,
+      category: "test",
+      message: "error",
+      fatal: false,
+      timestamp: new Date().toISOString(),
+      sessionId: "s",
+      seq: 1,
+    };
+    try {
+      await expect(storage.commit({ message: "partial" })).rejects.toThrow(
+        "rollback remains pending",
+      );
+      expect(staged).toContain("evidence-archive/added.jsonl");
+      await expect(audit.commitErrors([errorRecord])).rejects.toThrow("injected reset failure");
+      await expect(
+        audit.commitAudit([
+          {
+            callId: "a",
+            tool: "read_file",
+            arguments: {},
+            authz: null,
+            result: { content: "ok", isError: false },
+            timestamp: new Date().toISOString(),
+            sessionId: "s",
+            seq: 2,
+          },
+        ]),
+      ).rejects.toThrow("injected reset failure");
+      expect(await git.resolveRef({ fs, dir, ref: "HEAD" })).toBe(published.hash);
+      resetFails = false;
+      await audit.commitErrors([errorRecord]);
+      expect(
+        (await git.listFiles({ fs, dir, ref: "HEAD" })).filter((name) =>
+          name.startsWith("evidence-archive"),
+        ),
+      ).toEqual(["evidence-archive/index.jsonl"]);
+      const { blob } = await git.readBlob({
+        fs,
+        dir,
+        oid: await git.resolveRef({ fs, dir, ref: "HEAD" }),
+        filepath: "evidence-archive/index.jsonl",
+      });
+      expect(new TextDecoder().decode(blob)).toBe("original");
+      expect(fs.readFileSync(path.join(archiveDir, "index.jsonl"), "utf8")).toBe("changed");
+    } finally {
+      addSpy.mockRestore();
+      resetSpy.mockRestore();
+    }
+    await storage.commit({ message: "retry" });
+    expect(
+      (await git.listFiles({ fs, dir, ref: "HEAD" })).filter((name) =>
+        name.startsWith("evidence-archive"),
+      ),
+    ).toEqual(["evidence-archive/added.jsonl", "evidence-archive/index.jsonl"]);
+  });
+  test.each([true, false])(
+    "failed compact restores concrete archive index leaves (already tracked: %s)",
+    async (tracked) => {
+      const dir = tempDir();
+      const { loadOrCreateCommitSigner } = await import("./commit-signer.js");
+      const sign = await loadOrCreateCommitSigner(dir);
+      let fail = false;
+      const { storage, audit } = await createSessionStores(dir, {
+        signer: (payload) => {
+          if (fail) throw new Error("injected signing failure");
+          return sign(payload);
+        },
+      });
+      const archiveDir = path.join(dir, "evidence-archive");
+      fs.mkdirSync(archiveDir, { recursive: true });
+      if (tracked) {
+        fs.writeFileSync(path.join(archiveDir, "index.jsonl"), "original evidence");
+        fs.writeFileSync(path.join(archiveDir, "deleted.jsonl"), "old evidence");
+      }
+      await storage.writeTurns([turn("old-a"), turn("old-b")]);
+      await storage.writeMetadata(EMPTY_CHECKPOINT_METADATA);
+      const original = await storage.commit({ message: "original" });
+      const originalLeaves = (await git.listFiles({ fs, dir })).filter((name) =>
+        name.startsWith("evidence-archive"),
+      );
+      fs.writeFileSync(path.join(archiveDir, "index.jsonl"), "new evidence");
+      fs.writeFileSync(path.join(archiveDir, "added.jsonl"), "added evidence");
+      if (tracked) fs.unlinkSync(path.join(archiveDir, "deleted.jsonl"));
+      await storage.writeTurns([turn("[Compacted prior context]")]);
+      fail = true;
+      await expect(storage.commit({ message: "failed compact" })).rejects.toThrow(
+        "injected signing failure",
+      );
+      expect(
+        (await git.listFiles({ fs, dir })).filter((name) => name.startsWith("evidence-archive")),
+      ).toEqual(originalLeaves);
+      expect(await git.resolveRef({ fs, dir, ref: "HEAD" })).toBe(original.hash);
+      for (const [, head, , stage] of await git.statusMatrix({
+        fs,
+        dir,
+        filepaths: ["evidence-archive"],
+      }))
+        expect(stage).toBe(head);
+      expect(fs.readFileSync(path.join(archiveDir, "index.jsonl"), "utf8")).toBe("new evidence");
+      fail = false;
+      await audit.commitAudit([
+        {
+          callId: "audit",
+          tool: "read_file",
+          arguments: {},
+          authz: null,
+          result: { content: "ok", isError: false },
+          timestamp: new Date().toISOString(),
+          sessionId: "s",
+          seq: 1,
+        },
+      ]);
+      await audit.commitErrors([
+        {
+          source: "reactor",
+          category: "test",
+          message: "test error",
+          fatal: false,
+          timestamp: new Date().toISOString(),
+          sessionId: "s",
+          seq: 2,
+        },
+      ]);
+      expect(
+        (await git.listFiles({ fs, dir, ref: "HEAD" })).filter((name) =>
+          name.startsWith("evidence-archive"),
+        ),
+      ).toEqual(originalLeaves);
+      if (tracked) {
+        const { blob } = await git.readBlob({
+          fs,
+          dir,
+          oid: await git.resolveRef({ fs, dir, ref: "HEAD" }),
+          filepath: "evidence-archive/index.jsonl",
+        });
+        expect(new TextDecoder().decode(blob)).toBe("original evidence");
+      }
+      await storage.commit({ message: "retry compact" });
+      expect(
+        (await git.listFiles({ fs, dir, ref: "HEAD" })).filter((name) =>
+          name.startsWith("evidence-archive"),
+        ),
+      ).toEqual(["evidence-archive/added.jsonl", "evidence-archive/index.jsonl"]);
+      const { blob } = await git.readBlob({
+        fs,
+        dir,
+        oid: await git.resolveRef({ fs, dir, ref: "HEAD" }),
+        filepath: "evidence-archive/index.jsonl",
+      });
+      expect(new TextDecoder().decode(blob)).toBe("new evidence");
+    },
+  );
   test("rewrite writeTurns stays off the live generation until commit", async () => {
     const dir = tempDir();
     const store = await createOptimizedContextStore(dir);

src/session/optimized-context-store.ts

@@ -330,14 +330,29 @@ async function blobTextAtCommit(dir: string, hash: string, filepath: string): Pr
 
 async function resetIndexPaths(dir: string, filepaths: readonly string[]): Promise<void> {
   for (const filepath of filepaths) {
-    try {
-      await git.resetIndex({ fs, dir, filepath });
-    } catch {
-      // Not in the index; vendor restore already covers its own paths.
-    }
+    await git.resetIndex({ fs, dir, filepath });
   }
 }
 
+const pendingCheckpointRollbacks = new Map<string, () => Promise<void>>();
+
+async function reconcilePendingCheckpoint(dir: string): Promise<void> {
+  const key = path.resolve(dir);
+  const rollback = pendingCheckpointRollbacks.get(key);
+  if (rollback === undefined) return;
+  await rollback();
+  pendingCheckpointRollbacks.delete(key);
+}
+
+function withReconciledDirLock<T>(dir: string, operation: () => Promise<T>): Promise<T> {
+  return withResolvedDirLock(dir, async () => {
+    // All store instances sharing this index must finish a failed rollback
+    // before an audit, error, or checkpoint commit can publish it.
+    await reconcilePendingCheckpoint(dir);
+    return operation();
+  });
+}
+
 async function restorePublishedTurnFiles(dir: string): Promise<void> {
   const hash = await git.resolveRef({ fs, dir, ref: "HEAD" });
   const present = new Set(await git.listFiles({ fs, dir, ref: hash }));
@@ -589,7 +604,7 @@ export async function createSessionStores(
       pendingBlobFilepaths.add(`${TOOL_OUTPUT_DIR}/${filename}`);
     },
     async commit(options, signal) {
-      return withResolvedDirLock(dir, async () => {
+      return withReconciledDirLock(dir, async () => {
         const stagedRewrite = unpublishedRewrite;
         const extraPaths: string[] = [];
         try {
@@ -602,8 +617,14 @@ export async function createSessionStores(
             if (await pathExists(path.join(dir, filepath))) toAdd.push(filepath);
             else toRemove.push(filepath);
           }
-          if (await pathExists(path.join(dir, EVIDENCE_ARCHIVE_DIR))) {
-            toAdd.push(EVIDENCE_ARCHIVE_DIR);
+          const archiveRows = await git.statusMatrix({
+            fs,
+            dir,
+            filepaths: [EVIDENCE_ARCHIVE_DIR],
+          });
+          for (const [filepath, , worktree] of archiveRows) {
+            if (worktree === 0) toRemove.push(filepath);
+            else toAdd.push(filepath);
           }
           // Reconcile disk segments even if a prior process lost its pending set.
           await reconcileSegmentStaging(dir, TURNS_FILE, toAdd, toRemove);
@@ -615,11 +636,7 @@ export async function createSessionStores(
             await git.add({ fs, dir, filepath });
           }
           for (const filepath of remove) {
-            try {
-              await git.remove({ fs, dir, filepath });
-            } catch {
-              // Already absent from the index.
-            }
+            await git.remove({ fs, dir, filepath });
           }
           const committed = await base.commit(options, signal);
           pendingBlobFilepaths.clear();
@@ -630,20 +647,30 @@ export async function createSessionStores(
           }
           return committed;
         } catch (cause) {
-          await resetIndexPaths(dir, extraPaths);
-          if (stagedRewrite !== null) {
-            await restorePublishedTurnFiles(dir);
-            writeTurnsSegmented = createSegmentedJSONLWriter(dir, TURNS_FILE);
-            liveTurnRefs = null;
+          pendingCheckpointRollbacks.set(path.resolve(dir), async () => {
+            await resetIndexPaths(dir, extraPaths);
+            if (stagedRewrite !== null) {
+              await restorePublishedTurnFiles(dir);
+              writeTurnsSegmented = createSegmentedJSONLWriter(dir, TURNS_FILE);
+              liveTurnRefs = null;
+            }
+          });
+          try {
+            await reconcilePendingCheckpoint(dir);
+          } catch (rollbackCause) {
+            throw new AggregateError(
+              [cause, rollbackCause],
+              "Checkpoint failed; index rollback remains pending",
+            );
           }
           throw cause;
         }
       });
     },
     commitAudit: (records, signal) =>
-      withResolvedDirLock(dir, () => base.commitAudit(records, signal)),
+      withReconciledDirLock(dir, () => base.commitAudit(records, signal)),
     commitErrors: (records, signal) =>
-      withResolvedDirLock(dir, () => base.commitErrors(records, signal)),
+      withReconciledDirLock(dir, () => base.commitErrors(records, signal)),
     loadAudit: (sessionId, signal) => base.loadAudit(sessionId, signal),
   };