diff --git a/src/session/hooks.ts b/src/session/hooks.ts index 429aaa7dc..d9a19b4aa 100644 --- a/src/session/hooks.ts +++ b/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 }; diff --git a/src/session/run-sink.test.ts b/src/session/run-sink.test.ts index 9d0328cf5..fbdaa16db 100644 --- a/src/session/run-sink.test.ts +++ b/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(), diff --git a/src/session/run-sink.ts b/src/session/run-sink.ts index 76b3feacd..6f2f63566 100644 --- a/src/session/run-sink.ts +++ b/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(); diff --git a/src/session/state.test.ts b/src/session/state.test.ts new file mode 100644 index 000000000..ed28ff369 --- /dev/null +++ b/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>; + +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 { + 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"); +}); diff --git a/src/session/state.ts b/src/session/state.ts index 4da31c630..d9cc2ce7e 100644 --- a/src/session/state.ts +++ 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>(); + export async function saveState( cwd: string, sessionId: string, state: RunState, home?: string, ): Promise { - 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; } diff --git a/src/tui/resume-seed.test.ts b/src/tui/resume-seed.test.ts new file mode 100644 index 000000000..0fb870a1a --- /dev/null +++ b/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 { + 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([]); + }); +}); diff --git a/src/tui/runner.ts b/src/tui/runner.ts index 7291fbe30..ab56aa11e 100644 --- a/src/tui/runner.ts +++ b/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 = { session: "This session", project: "This project", @@ -406,6 +429,9 @@ export async function runTUI(initialConfig: Config): Promise { 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 }); @@ -413,6 +439,7 @@ export async function runTUI(initialConfig: Config): Promise { 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,11 +462,11 @@ export async function runTUI(initialConfig: Config): Promise { // 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 @@ -447,8 +474,16 @@ export async function runTUI(initialConfig: Config): Promise { // 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 { 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 { // 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.