Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion src/session/hooks.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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 };
Expand Down
18 changes: 18 additions & 0 deletions src/session/run-sink.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down
13 changes: 10 additions & 3 deletions src/session/run-sink.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand Down Expand Up @@ -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;
Expand All @@ -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();

Expand Down
76 changes: 76 additions & 0 deletions src/session/state.test.ts
Original file line number Diff line number Diff line change
@@ -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");
});
29 changes: 28 additions & 1 deletion src/session/state.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}


Expand Down
41 changes: 41 additions & 0 deletions src/tui/resume-seed.test.ts
Original file line number Diff line number Diff line change
@@ -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([]);
});
});
46 changes: 41 additions & 5 deletions src/tui/runner.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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;
Expand All @@ -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.
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down
Loading