diff --git a/packages/agent/src/harness/session/jsonl/repo.ts b/packages/agent/src/harness/session/jsonl/repo.ts index 143b23651..ee7482c21 100644 --- a/packages/agent/src/harness/session/jsonl/repo.ts +++ b/packages/agent/src/harness/session/jsonl/repo.ts @@ -51,6 +51,7 @@ export class JsonlSessionRepo private readonly fileSystem: FileSystem; private readonly sessionsRootInput: string; private readonly now: () => number; + private readonly activeCreateDestinations = new Set(); private rootPromise: Promise | undefined; private closed = false; private closePromise: Promise | undefined; @@ -64,45 +65,43 @@ export class JsonlSessionRepo async create(options: JsonlSessionCreateOptions): Promise> { this.assertOpen(); const createdAt = this.now(); - const id = options.id ?? uuidv7(createdAt); - const cwd = fileValue( - await this.fileSystem.absolutePath(options.cwd), - `Failed to resolve session cwd ${options.cwd}`, - ); - // TODO: Prevent concurrent create calls from passing this asynchronous check for the same cwd and id. - if (await this.sessionIdExists(cwd, id)) throw new Error(`Session already exists: ${id}`); - const path = await this.createPath(cwd, createdAt, id); - const header: JsonlStorageHeader = { - v: JSONL_FORMAT_VERSION, - kind: "header", - id, - storageVersion: JSONL_STORAGE_VERSION, - createdAt, - cwd, - ...(options.parentSessionId === undefined ? {} : { parentSessionId: options.parentSessionId }), - }; - // TODO: do we want to make session creation atomic? - const storage = await JsonlStorage.create({ fileSystem: this.fileSystem, path, now: this.now }, header); - try { - await storage.commit({ - writes: [ - { kind: "register", op: "set", namespace: "lane.leaf", key: "main", value: null }, - { - kind: "register", - op: "set", - namespace: "lane.state", - key: "main", - value: { currentOperationId: null, pendingNextRun: [] }, - }, - ], - }); - const info = fileValue(await this.fileSystem.fileInfo(path), `Failed to read session ${path}`); - return new StorageBackedSession(metadataFromHeader(header, path, info.mtimeMs), storage); - } catch (error) { - await storage.close().catch(() => undefined); - await this.fileSystem.remove(path, { force: true }); - throw error; - } + const destination = await this.resolveCreateDestination(options.cwd, options.id, createdAt); + return this.claimCreateDestination(destination, async () => { + const { cwd, id } = destination; + if (await this.sessionIdExists(cwd, id)) throw new Error(`Session already exists: ${id}`); + const path = await this.createPath(cwd, createdAt, id); + const header: JsonlStorageHeader = { + v: JSONL_FORMAT_VERSION, + kind: "header", + id, + storageVersion: JSONL_STORAGE_VERSION, + createdAt, + cwd, + ...(options.parentSessionId === undefined ? {} : { parentSessionId: options.parentSessionId }), + }; + // TODO: do we want to make session creation atomic? + const storage = await JsonlStorage.create({ fileSystem: this.fileSystem, path, now: this.now }, header); + try { + await storage.commit({ + writes: [ + { kind: "register", op: "set", namespace: "lane.leaf", key: "main", value: null }, + { + kind: "register", + op: "set", + namespace: "lane.state", + key: "main", + value: { currentOperationId: null, pendingNextRun: [] }, + }, + ], + }); + const info = fileValue(await this.fileSystem.fileInfo(path), `Failed to read session ${path}`); + return new StorageBackedSession(metadataFromHeader(header, path, info.mtimeMs), storage); + } catch (error) { + await storage.close().catch(() => undefined); + await this.fileSystem.remove(path, { force: true }); + throw error; + } + }); } async open(metadata: JsonlSessionMetadata): Promise> { @@ -142,36 +141,39 @@ export class JsonlSessionRepo async fork(source: JsonlSessionMetadata, options: ForkOptions): Promise> { this.assertOpen(); const createdAt = this.now(); - const id = options.id ?? uuidv7(createdAt); - if (await this.sessionIdExists(source.cwd, id)) throw new Error(`Session already exists: ${id}`); + const destination = await this.resolveCreateDestination(source.cwd, options.id, createdAt); + return this.claimCreateDestination(destination, async () => { + const { cwd, id } = destination; + if (await this.sessionIdExists(cwd, id)) throw new Error(`Session already exists: ${id}`); - const sourceStorage = await this.loadStorage(source); - const sourceSnapshot = await sourceStorage.snapshot().finally(() => sourceStorage.close()); - const snapshot = createForkSnapshot(sourceSnapshot, options); - const path = await this.createPath(source.cwd, createdAt, id); - const header: JsonlStorageHeader = { - v: JSONL_FORMAT_VERSION, - kind: "header", - id, - storageVersion: JSONL_STORAGE_VERSION, - createdAt, - cwd: source.cwd, - parentSessionId: source.id, - }; - let storage: JsonlStorage | undefined; - try { - storage = await JsonlStorage.createFromSnapshot( - { fileSystem: this.fileSystem, path, now: this.now }, - header, - snapshot, - ); - const info = fileValue(await this.fileSystem.fileInfo(path), `Failed to read session ${path}`); - return new StorageBackedSession(metadataFromHeader(header, path, info.mtimeMs), storage); - } catch (error) { - await storage?.close().catch(() => undefined); - await this.fileSystem.remove(path, { force: true }); - throw error; - } + const sourceStorage = await this.loadStorage(source); + const sourceSnapshot = await sourceStorage.snapshot().finally(() => sourceStorage.close()); + const snapshot = createForkSnapshot(sourceSnapshot, options); + const path = await this.createPath(cwd, createdAt, id); + const header: JsonlStorageHeader = { + v: JSONL_FORMAT_VERSION, + kind: "header", + id, + storageVersion: JSONL_STORAGE_VERSION, + createdAt, + cwd, + parentSessionId: source.id, + }; + let storage: JsonlStorage | undefined; + try { + storage = await JsonlStorage.createFromSnapshot( + { fileSystem: this.fileSystem, path, now: this.now }, + header, + snapshot, + ); + const info = fileValue(await this.fileSystem.fileInfo(path), `Failed to read session ${path}`); + return new StorageBackedSession(metadataFromHeader(header, path, info.mtimeMs), storage); + } catch (error) { + await storage?.close().catch(() => undefined); + await this.fileSystem.remove(path, { force: true }); + throw error; + } + }); } close(): Promise { @@ -182,6 +184,30 @@ export class JsonlSessionRepo return this.closePromise; } + private async resolveCreateDestination( + cwdInput: string, + id: string | undefined, + createdAt: number, + ): Promise<{ cwd: string; id: string }> { + const destinationId = id ?? uuidv7(createdAt); + const cwd = fileValue(await this.fileSystem.absolutePath(cwdInput), `Failed to resolve session cwd ${cwdInput}`); + return { cwd, id: destinationId }; + } + + private async claimCreateDestination( + destination: { cwd: string; id: string }, + operation: () => Promise, + ): Promise { + const key = `${destination.cwd}\0${destination.id}`; + if (this.activeCreateDestinations.has(key)) throw new Error(`Session already exists: ${destination.id}`); + this.activeCreateDestinations.add(key); + try { + return await operation(); + } finally { + this.activeCreateDestinations.delete(key); + } + } + private async listDirectory(directory: string, cwd?: string): Promise { if (!fileValue(await this.fileSystem.exists(directory), `Failed to check sessions directory ${directory}`)) { return []; diff --git a/packages/agent/src/harness/session/testing/conformance/session-repo.ts b/packages/agent/src/harness/session/testing/conformance/session-repo.ts index db89f5e04..20dc56f50 100644 --- a/packages/agent/src/harness/session/testing/conformance/session-repo.ts +++ b/packages/agent/src/harness/session/testing/conformance/session-repo.ts @@ -476,8 +476,8 @@ export function createSessionRepoForkBehaviorConformance( +/** Creates fork cases that require destination reservation across create and fork. */ +export function createSessionRepoForkDestinationReservationConformance( backendFactory: () => Promise, "create" | "fork">>, onClose?: () => void | Promise, ): readonly ConformanceCase[] { @@ -515,6 +515,16 @@ export function createSessionRepoForkCoordinationConformance( + backendFactory: () => Promise, "create" | "fork">>, + onClose?: () => void | Promise, +): readonly ConformanceCase[] { + const factory = prepareRepoCaseFactory(backendFactory, onClose); + return [ createCase( factory, "fork coordination", @@ -555,6 +565,17 @@ export function createSessionRepoForkCoordinationConformance( + backendFactory: () => Promise, "create" | "fork">>, + onClose?: () => void | Promise, +): readonly ConformanceCase[] { + return [ + ...createSessionRepoForkDestinationReservationConformance(backendFactory, onClose), + ...createSessionRepoForkSourceSnapshotConformance(backendFactory, onClose), + ]; +} + /** Creates every fork conformance case. */ export function createSessionRepoForkConformance( backendFactory: () => Promise, "create" | "list" | "fork">>, diff --git a/packages/agent/src/harness/session/testing/index.ts b/packages/agent/src/harness/session/testing/index.ts index 839adc6fa..e0e32dffb 100644 --- a/packages/agent/src/harness/session/testing/index.ts +++ b/packages/agent/src/harness/session/testing/index.ts @@ -21,6 +21,8 @@ export { createSessionRepoForkBehaviorConformance, createSessionRepoForkConformance, createSessionRepoForkCoordinationConformance, + createSessionRepoForkDestinationReservationConformance, + createSessionRepoForkSourceSnapshotConformance, createSessionRepoLifecycleConformance, createSessionRepoMessageConformance, createSessionRepoOwnershipConformance, diff --git a/packages/agent/test/harness/jsonl-session-repo-conformance.test.ts b/packages/agent/test/harness/jsonl-session-repo-conformance.test.ts index 7e54f89ed..8fbda2d46 100644 --- a/packages/agent/test/harness/jsonl-session-repo-conformance.test.ts +++ b/packages/agent/test/harness/jsonl-session-repo-conformance.test.ts @@ -4,6 +4,7 @@ import { type JsonlSessionMetadata, JsonlSessionRepo } from "../../src/harness/s import { type ConformanceCase, createSessionRepoForkBehaviorConformance, + createSessionRepoForkDestinationReservationConformance, createSessionRepoLifecycleConformance, createSessionRepoMessageConformance, } from "../../src/harness/session/testing/index.ts"; @@ -43,4 +44,7 @@ registerConformance("JsonlSessionRepo conformance", [ ...createSessionRepoLifecycleConformance(createConformanceRepo, () => jsonlRepo.close()), ...createSessionRepoMessageConformance(createConformanceRepo, () => jsonlRepo.close()), ...createSessionRepoForkBehaviorConformance(createConformanceRepo, () => jsonlRepo.close()), + ...createSessionRepoForkDestinationReservationConformance(createConformanceRepo, () => + jsonlRepo.close(), + ), ]);