@tanstack/ai-sandbox 0.3.3 → 0.4.0

package/dist/esm/middleware.js

@@ -1,13 +1,16 @@
 import { SandboxCapability, provideSandbox, provideSandboxPolicy } from "./capabilities.js";
 import { SandboxInstanceStoreCapability } from "./instance-store.js";
+import { SandboxCheckpointError } from "./checkpoint-store.js";
+import { captureSandboxArtifacts, captureSandboxFiles, resolveSandboxSnapshotPolicy, restoreSandboxFiles } from "./snapshots.js";
+import { resolveAllSecrets, resolveSecret } from "./secrets.js";
+import { computeWorkspaceHash } from "./key.js";
+import { resolveHarnessCwd } from "./harness-cwd.js";
+import "./bootstrap.js";
+import { ensureSandboxWithOutcome } from "./sandbox.js";
 import { ProjectionCapability, provideWorkspaceProjection } from "./projection.js";
 import { provideSandboxDurability, resolveSandboxDurability } from "./durability.js";
-import { computeWorkspaceHash } from "./key.js";
 import { buildFileHookEvent, resolveFileEvents } from "./file-diff.js";
-import { resolveSecret } from "./secrets.js";
 import { createToolHistoryRecorder, stripObservedToolCalls } from "./tool-history.js";
-import { resolveHarnessCwd } from "./harness-cwd.js";
-import "./bootstrap.js";
 import { watchWorkspace } from "./watch.js";
 import { defineChatMiddleware, provideDetachableRun, provideRunDetached, wasCancelRequested } from "@tanstack/ai";
 import { InMemoryLockStore, LocksCapability } from "@tanstack/ai/locks";
@@ -32,6 +35,52 @@ import { getPendingTurn, getRunDisconnect, getSandboxRuntime } from "@tanstack/a
 * chunks), not from here — middleware setup runs before streaming begins.
 */
 var runState = /* @__PURE__ */ new WeakMap();
+function stopSnapshotLease(state, options = {}) {
+	if (options.closePortable) state.snapshotClosed = true;
+	if (state.snapshotStop) return state.snapshotStop;
+	if (state.snapshotCleaned) return Promise.resolve();
+	state.snapshotCleaned = true;
+	state.snapshotRenewalGeneration++;
+	if (state.snapshotRenewal !== void 0) clearTimeout(state.snapshotRenewal);
+	state.snapshotRenewal = void 0;
+	const renewTask = state.snapshotRenewTask;
+	const captureTask = state.snapshotCaptureTask;
+	const lease = state.snapshotLease;
+	state.snapshotLease = void 0;
+	state.snapshotStop = (async () => {
+		await renewTask?.catch(() => {});
+		await captureTask?.catch(() => {});
+		await lease?.release();
+	})();
+	return state.snapshotStop;
+}
+function startSnapshotRenewal(state) {
+	const lease = state.snapshotLease;
+	if (!lease) return;
+	const schedule = () => {
+		const generation = state.snapshotRenewalGeneration;
+		state.snapshotRenewal = setTimeout(() => {
+			(async () => {
+				state.snapshotRenewal = void 0;
+				if (state.snapshotCleaned || generation !== state.snapshotRenewalGeneration) return;
+				const renewal = Promise.resolve().then(async () => {
+					await lease.renew();
+				});
+				state.snapshotRenewTask = renewal;
+				try {
+					await renewal;
+				} catch (error) {
+					state.snapshotLost = error instanceof Error ? error : new Error(String(error));
+				} finally {
+					if (state.snapshotRenewTask === renewal) state.snapshotRenewTask = void 0;
+				}
+				if (state.snapshotLost) await stopSnapshotLease(state).catch(() => {});
+				else if (!state.snapshotCleaned && generation === state.snapshotRenewalGeneration) schedule();
+			})();
+		}, lease.renewAfterMs);
+	};
+	schedule();
+}
 /**
 * Stop the watcher and drain any in-flight `diff()` promises before teardown,
 * so the final file's diff isn't dropped when a run finishes/aborts/errors
@@ -50,6 +99,10 @@ async function drainWatcher(state, phase) {
 	await Promise.allSettled(state.pendingDiffs);
 	if (state.watcher) state.logger?.sandbox("sandbox watcher stopped", { phase });
 }
+function canPublishPortableSnapshot(state, lease) {
+	if (state.snapshotLost) throw state.snapshotLost;
+	return !state.snapshotClosed && !state.snapshotCleaned && state.snapshotLease === lease;
+}
 /**
 * Record the two facts a later attach and the reaper both need, then publish the
 * detach verdict core reads.
@@ -161,6 +214,25 @@ function withSandbox(definition, options) {
 		optionalRequires: [SandboxInstanceStoreCapability, LocksCapability],
 		async setup(ctx) {
 			const ensureCtx = buildEnsureCtx(ctx, options);
+			const snapshotConfig = options?.snapshots;
+			const snapshotWorkspaceHash = definition.workspace ? computeWorkspaceHash(definition.workspace) : void 0;
+			const snapshotPolicy = snapshotConfig ? resolveSandboxSnapshotPolicy(snapshotConfig.policy, snapshotWorkspaceHash) : void 0;
+			let snapshotRuntime;
+			let snapshotLease;
+			if (snapshotConfig) {
+				if (!snapshotConfig.persistence?.stores?.messages || !snapshotConfig.persistence.stores.artifacts || !snapshotConfig.persistence.stores.blobs) throw new Error("Sandbox snapshots require persistence stores.messages, stores.artifacts, and stores.blobs");
+				const persistenceModule = await import("@tanstack/ai-persistence");
+				const persistence = ctx.getOptional(persistenceModule.PersistenceCapability);
+				if (persistence === void 0) throw new Error("Sandbox snapshots require withPersistence(snapshots.persistence) before withSandbox");
+				if (persistence !== snapshotConfig.persistence) throw new Error("Sandbox snapshots require the same persistence instance passed to withPersistence");
+				const completion = ctx.getOptional(persistenceModule.PersistenceCompletionCapability);
+				if (!completion) throw new Error("Sandbox snapshots require withPersistence before withSandbox");
+				snapshotRuntime = {
+					persistence: snapshotConfig.persistence,
+					completion
+				};
+				snapshotLease = await snapshotConfig.checkpoints.acquireWriter(ctx.threadId);
+			}
 			const durability = resolveSandboxDurability(options);
 			if (durability !== void 0) {
 				provideSandboxDurability(ctx, durability);
@@ -170,12 +242,17 @@ function withSandbox(definition, options) {
 			const logger = runtime?.logger;
 			const state = {
 				ensureCtx,
+				snapshotRenewalGeneration: 0,
 				pendingDiffs: [],
 				toolHistory: createToolHistoryRecorder(),
 				...logger ? { logger } : {},
 				...durability ? { durability } : {}
 			};
 			runState.set(ctx, state);
+			if (snapshotLease) {
+				state.snapshotLease = snapshotLease;
+				startSnapshotRenewal(state);
+			}
 			if (durability !== void 0) {
 				try {
 					await durability.runs.createOrResume({
@@ -199,86 +276,148 @@ function withSandbox(definition, options) {
 				}
 			}
 			if (durability !== void 0 && durability.detachOnDisconnect) getRunDisconnect(ctx, { optional: true })?.subscribe(async () => {
-				if (await cancelIntent(durability, ctx.runId, false)) return;
-				if (await recordDetach(definition, state, durability, ctx, "disconnect")) state.logger?.sandbox("sandbox run detached on disconnect; the run continues", { runId: ctx.runId });
+				const snapshotStop = stopSnapshotLease(state, { closePortable: true });
+				snapshotStop.catch(() => {});
+				if (await cancelIntent(durability, ctx.runId, false)) {
+					await snapshotStop.catch((error) => {
+						state.logger?.warn("sandbox snapshot writer release failed", {
+							runId: ctx.runId,
+							phase: "disconnect",
+							error
+						});
+					});
+					return;
+				}
+				if (await recordDetach(definition, state, durability, ctx, "disconnect")) {
+					try {
+						await snapshotStop;
+					} catch (error) {
+						state.logger?.warn("sandbox snapshot writer release failed", {
+							runId: ctx.runId,
+							phase: "disconnect",
+							error
+						});
+					}
+					state.logger?.sandbox("sandbox run detached on disconnect; the run continues", { runId: ctx.runId });
+				} else await snapshotStop.catch((error) => {
+					state.logger?.warn("sandbox snapshot writer release failed", {
+						runId: ctx.runId,
+						phase: "disconnect",
+						error
+					});
+				});
 			});
-			const handle = await definition.ensure(ensureCtx);
+			let outcome = "created";
+			let handle;
+			try {
+				if (snapshotConfig) ({handle, outcome} = await ensureSandboxWithOutcome(definition, ensureCtx));
+				else handle = await definition.ensure(ensureCtx);
+				state.handle = handle;
+				state.privateHandle = snapshotConfig ? outcome !== "resumed" : true;
+				if (snapshotConfig && outcome !== "resumed") {
+					const head = await snapshotConfig.checkpoints.getHead(ctx.threadId);
+					if (head) {
+						const checkpoint = await snapshotConfig.checkpoints.get(head);
+						if (!checkpoint) throw new SandboxCheckpointError("SANDBOX_SNAPSHOT_CHECKPOINT_NOT_FOUND", `Checkpoint '${head}' was not found`);
+						await restoreSandboxFiles(handle, {
+							blobs: snapshotConfig.persistence.stores.blobs,
+							workspaceRoot: definition.workspace?.root ?? "/workspace"
+						}, checkpoint, snapshotPolicy);
+					}
+				}
+			} catch (error) {
+				await stopSnapshotLease(state).catch(() => {});
+				if (state.handle && state.privateHandle) await definition.destroy(ensureCtx).catch(() => {});
+				throw error;
+			}
 			state.handle = handle;
-			provideSandbox(ctx, handle);
-			if (definition.policy) provideSandboxPolicy(ctx, definition.policy);
-			if (durability !== void 0 && (ensureCtx.locks === void 0 || ensureCtx.locks instanceof InMemoryLockStore)) logger?.warn("sandbox durability is wired over an InMemoryLockStore: run claims are serialized within this process only and the lease never signals loss, so two hosts can drive one run and duplicate its event log. Use a distributed LockStore via withLocks for any multi-replica deploy.", { runId: ctx.runId });
-			const watchRoot = definition.workspace?.root ?? "/workspace";
-			let baseSha = "";
+			if (snapshotConfig) {
+				state.snapshotConfig = snapshotConfig;
+				state.snapshotPolicy = snapshotPolicy;
+				state.snapshotRuntime = snapshotRuntime;
+			}
 			try {
-				const shaRes = await handle.process.exec("git rev-parse HEAD", { cwd: watchRoot });
-				if (shaRes.exitCode === 0) {
-					baseSha = shaRes.stdout.trim();
-					logger?.sandbox("sandbox git baseline captured", {
+				provideSandbox(ctx, handle);
+				if (definition.policy) provideSandboxPolicy(ctx, definition.policy);
+				if (durability !== void 0 && (ensureCtx.locks === void 0 || ensureCtx.locks instanceof InMemoryLockStore)) logger?.warn("sandbox durability is wired over an InMemoryLockStore: run claims are serialized within this process only and the lease never signals loss, so two hosts can drive one run and duplicate its event log. Use a distributed LockStore via withLocks for any multi-replica deploy.", { runId: ctx.runId });
+				const watchRoot = definition.workspace?.root ?? "/workspace";
+				let baseSha = "";
+				try {
+					const shaRes = await handle.process.exec("git rev-parse HEAD", { cwd: watchRoot });
+					if (shaRes.exitCode === 0) {
+						baseSha = shaRes.stdout.trim();
+						logger?.sandbox("sandbox git baseline captured", {
+							root: watchRoot,
+							baseSha
+						});
+					} else logger?.sandbox("sandbox git baseline unavailable (non-zero exit)", {
 						root: watchRoot,
-						baseSha
+						exitCode: shaRes.exitCode,
+						stderr: shaRes.stderr
 					});
-				} else logger?.sandbox("sandbox git baseline unavailable (non-zero exit)", {
-					root: watchRoot,
-					exitCode: shaRes.exitCode,
-					stderr: shaRes.stderr
-				});
+				} catch (error) {
+					logger?.warn("sandbox git baseline capture failed", {
+						root: watchRoot,
+						error
+					});
+				}
+				const workspace = definition.workspace;
+				if (workspace !== void 0) {
+					const virtualRoot = workspace.root ?? "/workspace";
+					const root = resolveHarnessCwd(handle, virtualRoot);
+					const workspaceHash = computeWorkspaceHash(workspace);
+					const secrets = workspace.secrets;
+					provideWorkspaceProjection(ctx, {
+						skills: workspace.skills ?? [],
+						plugins: workspace.plugins ?? [],
+						resolveSecret: (ref) => {
+							if (secrets === void 0) throw new Error(`resolveSecret: no secrets defined on this workspace (ref: "${ref.__secretName}")`);
+							return resolveSecret(secrets, ref);
+						},
+						markerPath: `${root}/.tanstack-projected-${workspaceHash}`,
+						root,
+						...workspace.scripts !== void 0 ? { scripts: workspace.scripts } : {}
+					});
+				}
+				const hooks = definition.hooks;
+				await hooks?.onReady?.(handle);
+				const fe = resolveFileEvents(definition.fileEvents);
+				const pendingDiffs = state.pendingDiffs;
+				let watcher;
+				if (fe.enabled) {
+					watcher = await watchWorkspace(handle, {
+						onEvent: (event) => {
+							const enriched = buildFileHookEvent(handle, watchRoot, baseSha, event, logger);
+							dispatchDefinitionHooks(hooks, enriched, logger);
+							runtime?.emit(enriched);
+							if (fe.diff) pendingDiffs.push(enriched.diff().then((diff) => {
+								runtime?.emitFileDiff({
+									path: event.path,
+									diff
+								});
+							}).catch((error) => {
+								logger?.warn("sandbox file diff emit failed", {
+									path: event.path,
+									error
+								});
+							}));
+						},
+						root: watchRoot,
+						...ctx.signal !== void 0 ? { signal: ctx.signal } : {},
+						...logger !== void 0 ? { logger } : {}
+					});
+					logger?.sandbox("sandbox watcher started", {
+						root: watchRoot,
+						diff: fe.diff
+					});
+				}
+				if (watcher) state.watcher = watcher;
 			} catch (error) {
-				logger?.warn("sandbox git baseline capture failed", {
-					root: watchRoot,
-					error
-				});
-			}
-			const workspace = definition.workspace;
-			if (workspace !== void 0) {
-				const virtualRoot = workspace.root ?? "/workspace";
-				const root = resolveHarnessCwd(handle, virtualRoot);
-				const workspaceHash = computeWorkspaceHash(workspace);
-				const secrets = workspace.secrets;
-				provideWorkspaceProjection(ctx, {
-					skills: workspace.skills ?? [],
-					plugins: workspace.plugins ?? [],
-					resolveSecret: (ref) => {
-						if (secrets === void 0) throw new Error(`resolveSecret: no secrets defined on this workspace (ref: "${ref.__secretName}")`);
-						return resolveSecret(secrets, ref);
-					},
-					markerPath: `${root}/.tanstack-projected-${workspaceHash}`,
-					root,
-					...workspace.scripts !== void 0 ? { scripts: workspace.scripts } : {}
-				});
+				await drainWatcher(state, "error");
+				await stopSnapshotLease(state).catch(() => {});
+				if (state.privateHandle) await definition.destroy(ensureCtx).catch(() => {});
+				throw error;
 			}
-			const hooks = definition.hooks;
-			await hooks?.onReady?.(handle);
-			const fe = resolveFileEvents(definition.fileEvents);
-			const pendingDiffs = state.pendingDiffs;
-			let watcher;
-			if (fe.enabled) {
-				watcher = await watchWorkspace(handle, {
-					onEvent: (event) => {
-						const enriched = buildFileHookEvent(handle, watchRoot, baseSha, event, logger);
-						dispatchDefinitionHooks(hooks, enriched, logger);
-						runtime?.emit(enriched);
-						if (fe.diff) pendingDiffs.push(enriched.diff().then((diff) => {
-							runtime?.emitFileDiff({
-								path: event.path,
-								diff
-							});
-						}).catch((error) => {
-							logger?.warn("sandbox file diff emit failed", {
-								path: event.path,
-								error
-							});
-						}));
-					},
-					root: watchRoot,
-					...ctx.signal !== void 0 ? { signal: ctx.signal } : {},
-					...logger !== void 0 ? { logger } : {}
-				});
-				logger?.sandbox("sandbox watcher started", {
-					root: watchRoot,
-					diff: fe.diff
-				});
-			}
-			if (watcher) state.watcher = watcher;
 		},
 		onConfig(_ctx, config) {
 			const messages = stripObservedToolCalls(config.messages);
@@ -288,8 +427,13 @@ function withSandbox(definition, options) {
 		onIteration(ctx) {
 			runState.get(ctx)?.toolHistory.reconcile(ctx);
 		},
-		onChunk(ctx, chunk) {
-			runState.get(ctx)?.toolHistory.observe(chunk, ctx);
+		async onChunk(ctx, chunk) {
+			const state = runState.get(ctx);
+			state?.toolHistory.observe(chunk, ctx);
+			if (state && chunk.type === "RUN_FINISHED" && chunk.outcome?.type === "interrupt") {
+				await drainWatcher(state, "pause");
+				await stopSnapshotLease(state, { closePortable: true });
+			}
 		},
 		async onFinish(ctx) {
 			const state = runState.get(ctx);
@@ -297,49 +441,150 @@ function withSandbox(definition, options) {
 			const { handle, ensureCtx } = state;
 			state.toolHistory.reconcile(ctx);
 			await drainWatcher(state, "finish");
-			const lifecycle = definition.lifecycle;
-			if (lifecycle?.snapshot === "after-run" && handle?.capabilities.snapshots && handle.snapshot) {
-				const snapshot = await handle.snapshot(`after-run-${ctx.runId}`);
-				const store = ensureCtx.store;
-				if (store) {
-					const key = definition.key(ensureCtx);
-					const existing = await store.get(key);
-					if (existing) await store.upsert({
-						...existing,
-						latestSnapshotId: snapshot.id,
-						updatedAt: Date.now()
+			let primaryError;
+			try {
+				const snapshotCaptureTask = Promise.resolve().then(async () => {
+					const config = state.snapshotConfig;
+					const runtime = state.snapshotRuntime;
+					const lease = state.snapshotLease;
+					if (!config || !runtime || !handle || !lease) {
+						if (state.snapshotLost) throw state.snapshotLost;
+						return;
+					}
+					if (!canPublishPortableSnapshot(state, lease)) return;
+					await runtime.completion.waitForRunCompletion();
+					if (!canPublishPortableSnapshot(state, lease)) return;
+					const conversation = await runtime.persistence.stores.messages.loadThread(ctx.threadId);
+					if (!canPublishPortableSnapshot(state, lease)) return;
+					const files = await captureSandboxFiles(handle, {
+						blobs: config.persistence.stores.blobs,
+						workspaceRoot: definition.workspace?.root ?? "/workspace"
+					}, state.snapshotPolicy, definition.workspace?.secrets !== void 0 ? resolveAllSecrets(definition.workspace.secrets) : {});
+					if (!canPublishPortableSnapshot(state, lease)) return;
+					const artifacts = await captureSandboxArtifacts({
+						blobs: config.persistence.stores.blobs,
+						artifacts: config.persistence.stores.artifacts
+					}, ctx.threadId, definition.workspace?.secrets !== void 0 ? resolveAllSecrets(definition.workspace.secrets) : {});
+					if (!canPublishPortableSnapshot(state, lease)) return;
+					const parentCheckpointId = await config.checkpoints.getHead(ctx.threadId);
+					if (!canPublishPortableSnapshot(state, lease)) return;
+					try {
+						await config.checkpoints.append({
+							checkpoint: {
+								id: `checkpoint-${ctx.runId}`,
+								threadId: ctx.threadId,
+								parentCheckpointId,
+								createdAt: Date.now(),
+								reason: "automatic",
+								sourceRunId: ctx.runId,
+								files: files.files,
+								conversation,
+								artifacts
+							},
+							expectedHeadId: parentCheckpointId,
+							writer: lease
+						});
+					} catch (error) {
+						if (state.snapshotLost) throw state.snapshotLost;
+						throw error;
+					}
+					if (state.snapshotLost) throw state.snapshotLost;
+				});
+				state.snapshotCaptureTask = snapshotCaptureTask;
+				try {
+					await snapshotCaptureTask;
+				} finally {
+					if (state.snapshotCaptureTask === snapshotCaptureTask) state.snapshotCaptureTask = void 0;
+				}
+				const lifecycle = definition.lifecycle;
+				if (lifecycle?.snapshot === "after-run" && handle?.capabilities.snapshots && handle.snapshot) {
+					const snapshot = await handle.snapshot(`after-run-${ctx.runId}`);
+					const store = ensureCtx.store;
+					if (store) {
+						const key = definition.key(ensureCtx);
+						const existing = await store.get(key);
+						if (existing) await store.upsert({
+							...existing,
+							latestSnapshotId: snapshot.id,
+							updatedAt: Date.now()
+						});
+					}
+				}
+				if (lifecycle?.destroyOnComplete) {
+					await definition.destroy(ensureCtx);
+					await definition.hooks?.onDestroy?.();
+				}
+			} catch (error) {
+				primaryError = error;
+				if (definition.lifecycle?.destroyOnComplete) try {
+					await definition.destroy(ensureCtx);
+					await definition.hooks?.onDestroy?.();
+				} catch (cleanupError) {
+					state.logger?.warn("sandbox destroy after terminal failure failed", {
+						runId: ctx.runId,
+						phase: "finish",
+						error: cleanupError
 					});
 				}
 			}
-			if (lifecycle?.destroyOnComplete) {
-				await definition.destroy(ensureCtx);
-				await definition.hooks?.onDestroy?.();
+			let snapshotCleanupError;
+			try {
+				await stopSnapshotLease(state, { closePortable: true });
+			} catch (error) {
+				snapshotCleanupError = error;
 			}
+			if (primaryError !== void 0) {
+				if (snapshotCleanupError !== void 0) state.logger?.warn("sandbox snapshot writer release failed", {
+					runId: ctx.runId,
+					phase: "finish",
+					error: snapshotCleanupError
+				});
+				throw primaryError;
+			}
+			if (snapshotCleanupError !== void 0) throw snapshotCleanupError;
 		},
 		async onAbort(ctx, info) {
 			const state = runState.get(ctx);
 			if (!state) return;
 			await drainWatcher(state, "abort");
+			let releaseError;
+			try {
+				await stopSnapshotLease(state, { closePortable: true });
+			} catch (error) {
+				releaseError = error;
+			}
 			const durability = state.durability;
 			const cancelled = await cancelIntent(durability, ctx.runId, info.cancelRequested === true);
 			if (durability !== void 0 && !cancelled && durability.detachOnDisconnect) {
-				if (await recordDetach(definition, state, durability, ctx, "abort")) return;
+				if (await recordDetach(definition, state, durability, ctx, "abort")) {
+					if (releaseError) throw releaseError;
+					return;
+				}
 				await definition.destroy(state.ensureCtx);
 				await definition.hooks?.onDestroy?.();
+				if (releaseError) throw releaseError;
 				return;
 			}
 			await definition.destroy(state.ensureCtx);
 			await definition.hooks?.onDestroy?.();
+			if (releaseError) throw releaseError;
 		},
 		async onError(ctx, info) {
 			const state = runState.get(ctx);
 			if (!state) return;
 			await drainWatcher(state, "error");
+			let releaseError;
+			try {
+				await stopSnapshotLease(state);
+			} catch (error) {
+				releaseError = error;
+			}
 			await definition.hooks?.onError?.(info.error);
 			if (definition.lifecycle?.destroyOnComplete) {
 				await definition.destroy(state.ensureCtx);
 				await definition.hooks?.onDestroy?.();
 			}
+			if (releaseError) throw releaseError;
 		}
 	});
 }