fix(workspaces): stop checkpointing inside the AI edit turn

src/features/workspaces/documents/document-session-transient.test.ts

@@ -0,0 +1,65 @@
+import { describe, expect, it } from "vitest";
+
+import {
+	isTransientDocumentSessionError,
+	retryOnTransientDocumentSessionReset,
+} from "#/features/workspaces/documents/document-session-transient";
+
+describe("isTransientDocumentSessionError", () => {
+	it("matches the mid-turn storage-timeout reset", () => {
+		const error = new Error(
+			"Durable Object storage operation exceeded timeout which caused object to be reset",
+		);
+		expect(isTransientDocumentSessionError(error)).toBe(true);
+	});
+
+	it("finds the reset through a wrapped cause chain", () => {
+		const error = new Error("RPC failed", {
+			cause: new Error("internal error ... caused object to be reset"),
+		});
+		expect(isTransientDocumentSessionError(error)).toBe(true);
+	});
+
+	it("does not match an ordinary application error", () => {
+		expect(isTransientDocumentSessionError(new Error("Document session has been deleted."))).toBe(
+			false,
+		);
+	});
+});
+
+describe("retryOnTransientDocumentSessionReset", () => {
+	it("replays after a reset and returns the later result", async () => {
+		let attempts = 0;
+		const result = await retryOnTransientDocumentSessionReset(async () => {
+			attempts += 1;
+			if (attempts === 1) {
+				throw new Error("storage operation exceeded timeout which caused object to be reset");
+			}
+			return "applied";
+		});
+		expect(attempts).toBe(2);
+		expect(result).toBe("applied");
+	});
+
+	it("rethrows a non-transient error without retrying", async () => {
+		let attempts = 0;
+		await expect(
+			retryOnTransientDocumentSessionReset(async () => {
+				attempts += 1;
+				throw new Error("content_changed");
+			}),
+		).rejects.toThrow("content_changed");
+		expect(attempts).toBe(1);
+	});
+
+	it("gives up after the retry budget and rethrows the last reset", async () => {
+		let attempts = 0;
+		await expect(
+			retryOnTransientDocumentSessionReset(async () => {
+				attempts += 1;
+				throw new Error("caused object to be reset");
+			}),
+		).rejects.toThrow("caused object to be reset");
+		expect(attempts).toBe(3);
+	});
+});

src/features/workspaces/documents/document-session-transient.ts

@@ -0,0 +1,40 @@
+// Cloudflare resets a Durable Object when a storage operation outlives its
+// deadline, and surfaces that as an RPC error whose message ends with "caused
+// object to be reset". The reset drops the in-flight turn, not the persisted
+// state, so the same call succeeds on a fresh isolate. The sandbox runtime's
+// `isPlatformTransientError` only matches the startup variant of this message,
+// so the document-session RPC path needs its own classifier.
+const durableObjectResetPattern = /caused object to be reset/i;
+const maximumDocumentSessionRetries = 3;
+
+export function isTransientDocumentSessionError(error: unknown): boolean {
+	for (let current: unknown = error, depth = 0; current != null && depth < 8; depth += 1) {
+		const message =
+			current instanceof Error ? current.message : typeof current === "string" ? current : "";
+		if (durableObjectResetPattern.test(message)) {
+			return true;
+		}
+		current = typeof current === "object" ? (current as { cause?: unknown }).cause : undefined;
+	}
+	return false;
+}
+
+/**
+ * Retry `run` when the document session Durable Object is reset mid-turn. The
+ * caller must make `run` idempotent — document edits are, because the operation
+ * id makes a repeated `applyEdits` return the first attempt's receipt.
+ */
+export async function retryOnTransientDocumentSessionReset<T>(run: () => Promise<T>): Promise<T> {
+	let lastError: unknown;
+	for (let attempt = 0; attempt < maximumDocumentSessionRetries; attempt += 1) {
+		try {
+			return await run();
+		} catch (error) {
+			if (!isTransientDocumentSessionError(error)) {
+				throw error;
+			}
+			lastError = error;
+		}
+	}
+	throw lastError;
+}

src/features/workspaces/documents/document-session.ts

@@ -74,11 +74,7 @@ export interface DocumentSessionApplyEditsResult {
 	 * receipt states what the edit did rather than what is left of it later. */
 	lineChanges?: DocumentEditLineChanges;
 	failures: {
-		code:
-			| DocumentAiEditFailureCode
-			| "content_changed"
-			| "operation_id_conflict"
-			| "path_not_found";
+		code: DocumentAiEditFailureCode | "content_changed" | "operation_id_conflict";
 		detail?: string;
 		index: number;
 	}[];
@@ -291,12 +287,10 @@ export class DocumentSession extends YServer {
 					: []),
 			]);
 		});
-		this.assertActive();
-		if (!(await this.checkpointDocument())) {
-			return rejectedDocumentEditResult("path_not_found", input.edits.length);
-		}
-		this.assertActive();
-
+		// The Postgres checkpoint takes a contended workspace row lock. Awaiting it
+		// here held these storage writes open across that round trip and reset the
+		// object under load, failing the edit. The reconcile above already schedules
+		// the debounced onSave, which checkpoints off this turn.
 		return result;
 	}
 
@@ -345,10 +339,8 @@ export class DocumentSession extends YServer {
 			}
 		});
 
-		if (!(await this.checkpointDocument())) {
-			return { status: "not_found" };
-		}
-
+		// The reconcile above schedules the debounced onSave, which checkpoints to
+		// Postgres off this turn — see applyEdits for why the checkpoint is deferred.
 		return { status: "undone" };
 	}
 

src/features/workspaces/operations/edit-item.ts

@@ -1,4 +1,5 @@
 import { getDocumentSessionFromEnv } from "#/features/workspaces/document-session-access";
+import { retryOnTransientDocumentSessionReset } from "#/features/workspaces/documents/document-session-transient";
 import {
 	authorizeWorkspaceOperation,
 	resolveWorkspaceExistingItemPath,
@@ -109,36 +110,39 @@ export async function editWorkspaceItemOperation(
 			path: resolution.path,
 		};
 	}
-	const [targets, documentSession] = await Promise.all([
-		resolveEditTargets(accessContext, input.edits, (record) =>
-			record.location.kind === "document-block" &&
-			record.location.itemId === resolution.item.id &&
-			record.revision
-				? `${record.location.blockId}.r_${record.revision}`
-				: undefined,
-		),
-		getDocumentSession({
+	const targets = await resolveEditTargets(accessContext, input.edits, (record) =>
+		record.location.kind === "document-block" &&
+		record.location.itemId === resolution.item.id &&
+		record.revision
+			? `${record.location.blockId}.r_${record.revision}`
+			: undefined,
+	);
+	const resolvedEdits: DocumentAiEdit[] = await Promise.all(
+		input.edits.map(async (edit) => {
+			const contentEdit = replacePublicRefs(edit, targets);
+			return "html" in contentEdit
+				? {
+						...contentEdit,
+						html: await resolveDocumentCitations({
+							context: accessContext,
+							html: contentEdit.html,
+						}),
+					}
+				: contentEdit;
+		}),
+	);
+
+	// A reset object drops the turn, not the persisted document, so a fresh stub
+	// replays the edit; the operation id keeps the replay from applying twice.
+	const result = await retryOnTransientDocumentSessionReset(async () => {
+		const documentSession = await getDocumentSession({
 			itemId: resolution.item.id,
 			workspaceId: accessContext.workspaceId,
-		}),
-	]);
-
-	const result = await documentSession.applyEdits({
-		edits: await Promise.all(
-			input.edits.map(async (edit) => {
-				const contentEdit = replacePublicRefs(edit, targets);
-				return "html" in contentEdit
-					? {
-							...contentEdit,
-							html: await resolveDocumentCitations({
-								context: accessContext,
-								html: contentEdit.html,
-							}),
-						}
-					: contentEdit;
-			}),
-		),
-		operationId: accessContext.operationId,
+		});
+		return documentSession.applyEdits({
+			edits: resolvedEdits,
+			operationId: accessContext.operationId,
+		});
 	});
 
 	return {