Merge run.json write serialization and resume seeding

Author: TheGreatAxios

src/session/hooks.ts

@@ -174,6 +174,9 @@ export type TurnContextCollectorOptions = {
   // standing copy of recent history, so callers with nothing to hand it to
   // (no lifecycle hook) can opt out of retaining it.
   retainHistory?: boolean;
+  // Resuming a session should continue the persisted run.json turn count
+  // rather than restart it at zero.
+  initialTurnCount?: number;
 };
 
 export function createTurnContextCollector(
@@ -193,7 +196,7 @@ export function createTurnContextCollector(
 } {
   const retainHistory = options.retainHistory ?? true;
   const turns: TurnContext[] = [];
-  let turnCount = 0;
+  let turnCount = options.initialTurnCount ?? 0;
   let pending: PendingTurn | null = null;
   let cycleStartedAt = now();
   let tokenUsage: TokenUsage = { ...emptyUsage };

src/session/run-sink.test.ts

@@ -77,6 +77,24 @@ describe("createRunSink", () => {
     expect(runSink.getTokenUsage()).toEqual({ input: 1, output: 1, cacheRead: 0, cacheWrite: 0, thinking: 0 });
   });
 
+  test("seeds the turn count from a resumed session's prior turnsUsed", () => {
+    const runSink = createRunSink({
+      emitter: new EventEmitter(),
+      hookManager: stubHookManager([]),
+      initialTurnCount: 7,
+    });
+
+    expect(runSink.getTurnCount()).toBe(7);
+
+    runSink.sink(event("inference.done", {
+      turn: { role: "assistant", content: [], model: "test", timestamp: 0 },
+      usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0, thinking: 0 },
+      source: { provider: "test", model: "test" },
+    }));
+
+    expect(runSink.getTurnCount()).toBe(8);
+  });
+
   test("getLastTurnUsage reports the latest turn alone, not the running sum", () => {
     const runSink = createRunSink({
       emitter: new EventEmitter(),

src/session/run-sink.ts

@@ -19,6 +19,9 @@ export type RunSinkArgs = {
   // turn actually ran against, so consumers report per-turn provider/model
   // even if the live selection changed mid-run.
   onTurnComplete?: (ctx: import("./hooks.js").TurnContext) => void;
+  // Continues a resumed session's persisted run.json turn count instead of
+  // restarting the collector at zero.
+  initialTurnCount?: number;
 };
 
 export type RunSink = {
@@ -66,7 +69,7 @@ export function resolveExecRunStatus(args: {
 }
 
 export function createRunSink(args: RunSinkArgs): RunSink {
-  const { emitter, hookManager, onTurnComplete } = args;
+  const { emitter, hookManager, onTurnComplete, initialTurnCount } = args;
 
   function hasConfiguredHooks(): boolean {
     return hookManager.getStatuses().length > 0;
@@ -82,15 +85,19 @@ export function createRunSink(args: RunSinkArgs): RunSink {
     onTurnComplete?.(ctx);
   };
 
-  function createCollector(): TurnCollector {
+  // The initial seed only applies to the run's first collector (a resumed
+  // session's prior turnsUsed); a later reset() starts a fresh sub-session
+  // and should count from zero, not re-seed.
+  function createCollector(seedTurnCount?: number): TurnCollector {
     return createTurnContextCollector(handleTurn, Date.now, {
       retainHistory: hasConfiguredHooks(),
+      ...(seedTurnCount !== undefined ? { initialTurnCount: seedTurnCount } : {}),
     });
   }
 
   let runCompleted = false;
   let runError: string | undefined;
-  let turnCollector = createCollector();
+  let turnCollector = createCollector(initialTurnCount);
   // Always-on local PerfTrace: not gated by lifecycle hooks.
   let perfObserver = createPerfReactorObserver();
 

src/session/state.test.ts

@@ -0,0 +1,76 @@
+import { afterEach, beforeEach, expect, mock, test } from "bun:test";
+import * as realFs from "node:fs/promises";
+import { mkdir, rm } from "node:fs/promises";
+import { join } from "node:path";
+import { tmpdir } from "node:os";
+
+// Simulates the straggler write's real await point (e.g. cycleRecorder.dispose
+// during the terminal path) landing its writeFile after a later-issued
+// terminal write's writeFile, so rename-order alone would let it win.
+const realWriteFile = realFs.writeFile;
+let delayNextWrite = false;
+mock.module("node:fs/promises", () => ({
+  ...realFs,
+  writeFile: async (path: string, data: string) => {
+    if (delayNextWrite) {
+      delayNextWrite = false;
+      await new Promise((resolve) => setTimeout(resolve, 30));
+    }
+    return realWriteFile(path, data);
+  },
+}));
+
+const { loadState, saveState } = await import("./state.js");
+type RunState = Awaited<ReturnType<typeof loadState>>;
+
+let cwd = "";
+let home = "";
+
+beforeEach(async () => {
+  const stamp = `${Date.now()}-${Math.random().toString(16).slice(2)}`;
+  cwd = join(tmpdir(), `corbits-state-${stamp}`);
+  home = join(tmpdir(), `corbits-state-home-${stamp}`);
+  await mkdir(cwd, { recursive: true });
+  await mkdir(home, { recursive: true });
+});
+
+afterEach(async () => {
+  await rm(cwd, { recursive: true, force: true });
+  await rm(home, { recursive: true, force: true });
+});
+
+function state(overrides: Partial<NonNullable<RunState>>): NonNullable<RunState> {
+  return {
+    status: "running",
+    turnsUsed: 0,
+    task: "task",
+    startedAt: 1,
+    ...overrides,
+  };
+}
+
+test("a straggler snapshot started before a terminal write does not overwrite it", async () => {
+  const sessionId = "sess-race";
+
+  delayNextWrite = true;
+  const straggler = saveState(cwd, sessionId, state({ status: "running" }), home);
+  const terminal = saveState(cwd, sessionId, state({ status: "done", finishedAt: 999 }), home);
+
+  await Promise.all([straggler, terminal]);
+
+  const final = await loadState(cwd, sessionId, home);
+  expect(final?.status).toBe("done");
+  expect(final?.finishedAt).toBe(999);
+});
+
+test("saveState calls for different sessions do not block each other", async () => {
+  await Promise.all([
+    saveState(cwd, "session-a", state({ task: "a" }), home),
+    saveState(cwd, "session-b", state({ task: "b" }), home),
+  ]);
+
+  const a = await loadState(cwd, "session-a", home);
+  const b = await loadState(cwd, "session-b", home);
+  expect(a?.task).toBe("a");
+  expect(b?.task).toBe("b");
+});

src/session/state.ts

@@ -55,13 +55,40 @@ export function warnUnreadableState(path: string, reason: string): void {
   process.stderr.write(`${COMMAND_NAME}: ignoring unreadable state at ${path} (${reason}); starting fresh\n`);
 }
 
+// Concurrent saveState calls for the same session (a straggler progress
+// snapshot racing a terminal finalize write) have no ordering guarantee
+// between their underlying rename()s — the later call could still finish
+// first and resurrect a closed run.json as "running". Chaining each session's
+// writes onto the previous one forces them to apply in call order, so a
+// write issued after another always lands after it regardless of how long
+// either write's fs calls take. Keyed by sessionId, not path, since callers
+// only ever address one file per session.
+const writeChains = new Map<string, Promise<void>>();
+
 export async function saveState(
   cwd: string,
   sessionId: string,
   state: RunState,
   home?: string,
 ): Promise<void> {
-  await atomicWrite(statePath(cwd, sessionId, home), JSON.stringify(state, null, 2));
+  const path = statePath(cwd, sessionId, home);
+  const content = JSON.stringify(state, null, 2);
+  const previous = writeChains.get(sessionId) ?? Promise.resolve();
+  const write = previous.then(
+    () => atomicWrite(path, content),
+    () => atomicWrite(path, content),
+  );
+  // Swallow the error in the chain tail (not in `write`, which still rejects
+  // for this caller) so one failed save doesn't permanently wedge later
+  // saves for the same session.
+  const tail = write.catch(() => {});
+  writeChains.set(sessionId, tail);
+  // Once this is the last write for the session, drop the entry so a
+  // long-lived process doesn't retain a chain per session forever.
+  void tail.then(() => {
+    if (writeChains.get(sessionId) === tail) writeChains.delete(sessionId);
+  });
+  return write;
 }
 
 

src/tui/resume-seed.test.ts

@@ -0,0 +1,41 @@
+import { describe, test, expect } from "bun:test";
+import { resolveResumeSeed } from "./runner.js";
+import type { RunState } from "../session/state.js";
+
+function pickedState(overrides: Partial<RunState>): RunState {
+  return {
+    status: "running",
+    turnsUsed: 0,
+    task: "task",
+    startedAt: 1,
+    ...overrides,
+  };
+}
+
+describe("resolveResumeSeed", () => {
+  test("a fresh (non-resumed) run seeds zero turns and no servers", () => {
+    expect(resolveResumeSeed(null)).toEqual({ turnsUsed: 0, mcpServers: [] });
+  });
+
+  test("carries forward a resumed session's non-zero turnsUsed and non-empty mcpServers", () => {
+    const seed = resolveResumeSeed(
+      pickedState({
+        turnsUsed: 12,
+        mcpServers: [{ name: "filesystem", toolCount: 5 }, { name: "search", toolCount: 2 }],
+      }),
+    );
+
+    expect(seed.turnsUsed).toBe(12);
+    expect(seed.mcpServers).toEqual([
+      { name: "filesystem", toolCount: 5 },
+      { name: "search", toolCount: 2 },
+    ]);
+  });
+
+  test("defaults mcpServers to empty when a resumed record predates that field", () => {
+    const seed = resolveResumeSeed(pickedState({ turnsUsed: 3 }));
+
+    expect(seed.turnsUsed).toBe(3);
+    expect(seed.mcpServers).toEqual([]);
+  });
+});

src/tui/runner.ts

@@ -211,6 +211,29 @@ export function resumeTranscriptLoadErrorBlock(err: unknown): {
   return { type: "error", message: `Could not load prior session transcript: ${message}` };
 }
 
+export type ResumeSeed = {
+  turnsUsed: number;
+  mcpServers: ConnectedMcpServer[];
+};
+
+const FRESH_RESUME_SEED: ResumeSeed = { turnsUsed: 0, mcpServers: [] };
+
+/**
+ * Fold a resumed session's run.json into a concrete seed once, at the
+ * resume boundary, so every downstream reader (the run sink, the
+ * connected-servers list, the immediate post-resume saveState) trusts a
+ * fully-populated value instead of each repeating its own `?? 0` / `?? []`
+ * default. A fresh (non-resumed) run gets the same shape via
+ * FRESH_RESUME_SEED, so callers never branch on "was this a resume."
+ */
+export function resolveResumeSeed(pickedState: RunState | null): ResumeSeed {
+  if (pickedState === null) return FRESH_RESUME_SEED;
+  return {
+    turnsUsed: pickedState.turnsUsed,
+    mcpServers: pickedState.mcpServers ?? [],
+  };
+}
+
 const GRANT_SCOPE_LABEL: Record<GrantScope, string> = {
   session: "This session",
   project: "This project",
@@ -406,13 +429,17 @@ export async function runTUI(initialConfig: Config): Promise<number> {
   let resumeSkipInitialTask = config.skipInitialTask === true;
   let startedAt = Date.now();
   let runTaskTitle = config.task;
+  // Resolved once at the resume boundary so turnsUsed/mcpServers reads
+  // downstream never repeat their own omission-handling default.
+  let resumeSeed: ResumeSeed = FRESH_RESUME_SEED;
 
   if (config.resumePicker) {
     const picked = await pickSession(config.cwd, { includeCompleted: config.force });
     if (picked === null) return 0;
     sessionId = picked.sessionId;
     resumeSkipInitialTask = true;
     const pickedState = await loadState(config.cwd, sessionId);
+    resumeSeed = resolveResumeSeed(pickedState);
     if (pickedState !== null) {
       startedAt = pickedState.startedAt;
       runTaskTitle = pickedState.task;
@@ -435,20 +462,28 @@ export async function runTUI(initialConfig: Config): Promise<number> {
   // run.json at all.
   await saveState(config.cwd, sessionId, {
     status: "running",
-    turnsUsed: 0,
+    turnsUsed: resumeSeed.turnsUsed,
     task: runTaskTitle.trim().length > 0 ? runTaskTitle.trim() : "(conversation)",
     startedAt,
     model: `${config.providerName}:${config.model}`,
-    mcpServers: [],
+    mcpServers: resumeSeed.mcpServers,
   });
 
   // Crash guard: if anything from setup onward throws all the way out of
   // runTUI instead of reaching the normal finalize block, this still closes
   // out run.json so status and finishedAt never disagree. Declared before the
   // try so every fallible step after the minimal write above is covered.
   // `finalized` is set by the normal finalize path so this never double-writes
-  // on a clean exit; it also gates straggler snapshot writes (see
-  // persistRunSnapshot) from resurrecting a closed record.
+  // on a clean exit. It also gates persistRunSnapshot (below) from *issuing*
+  // a straggler write at all once the run is closed — a different job from
+  // saveState's per-session write ordering in state.ts. That ordering only
+  // decides which already-issued write lands last; it has no way to know a
+  // "running" snapshot fired after finalize is stale and should never be
+  // written in the first place. Without this flag such a snapshot would
+  // still queue behind the terminal write and legitimately "win" the
+  // ordering, resurrecting a closed run.json. Two different constraints
+  // (don't issue a stale write vs. order the writes you do issue), each
+  // owned by its own layer — not a duplicate check.
   let finalized = false;
   // Bound after the cycle recorder exists (it needs the session workdir); the
   // crash guard is declared first so it covers every fallible step below.
@@ -1305,6 +1340,7 @@ export async function runTUI(initialConfig: Config): Promise<number> {
   const runSink = createRunSink({
     emitter,
     hookManager,
+    initialTurnCount: resumeSeed.turnsUsed,
     onTurnComplete: (ctx) => {
       // provider_id is the canonical provider kind, never ctx.source.sourceId:
       // sourceId is the user-typed label from onboarding/settings, and free
@@ -1324,7 +1360,7 @@ export async function runTUI(initialConfig: Config): Promise<number> {
 
   // MCP servers connected so far, keyed by name so a reconnect after a failure
   // replaces rather than duplicates the entry.
-  let connectedMcpServers: ConnectedMcpServer[] = [];
+  let connectedMcpServers: ConnectedMcpServer[] = resumeSeed.mcpServers;
   // Every configured server's latest state, for the /mcp surface. Unlike
   // `connectedMcpServers` (persisted run metadata) this keeps the ones that
   // failed or are still waiting on authorization.