From 514d557771884b73fae1c7ca026462573ccf97d7 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Fri, 14 Aug 2026 12:11:45 +0200 Subject: [PATCH 01/12] feat(ai): add durable sandbox file snapshots --- .changeset/bright-cloudflare-snapshots.md | 5 + .changeset/calm-daytona-snapshots.md | 5 + .changeset/direct-docker-snapshots.md | 5 + .changeset/eager-local-process-snapshots.md | 5 + .changeset/fresh-sprites-snapshots.md | 5 + .changeset/fuzzy-snapshots-build.md | 5 + .changeset/gentle-vercel-snapshots.md | 5 + .changeset/quiet-sandbox-middleware.md | 5 + .changeset/tidy-artifact-history.md | 5 + docs/config.json | 20 +- .../build-your-own-generation-adapter.md | 15 +- docs/persistence/store-reference.md | 16 +- docs/sandbox/durability.md | 4 + docs/sandbox/lifecycle.md | 9 + docs/sandbox/overview.md | 2 + docs/sandbox/portable-snapshots.md | 353 +++ docs/sandbox/providers.md | 4 + .../src/lib/sqlite-persistence.test.ts | 307 ++- .../src/lib/sqlite-persistence.ts | 732 +++++- .../skills/ai-persistence/SKILL.md | 7 + .../build-cloudflare-artifact-store/SKILL.md | 72 +- packages/ai-persistence/src/capabilities.ts | 10 + packages/ai-persistence/src/index.ts | 4 + packages/ai-persistence/src/memory.ts | 35 +- packages/ai-persistence/src/middleware.ts | 74 +- .../ai-persistence/src/testkit/conformance.ts | 89 +- packages/ai-persistence/src/types.ts | 14 +- .../tests/artifact-thread.test.ts | 52 + .../tests/persistence-completion.test.ts | 661 +++++ packages/ai-sandbox-cloudflare/src/handle.ts | 38 + .../tests/handle.test.ts | 218 +- packages/ai-sandbox-daytona/src/handle.ts | 39 + .../ai-sandbox-daytona/tests/lstat.test.ts | 199 ++ packages/ai-sandbox-docker/src/handle.ts | 36 + .../tests/lstat-shell-protocol.test.ts | 159 ++ .../ai-sandbox-docker/tests/lstat.test.ts | 204 ++ .../tests/testkit-subpath.test.ts | 6 + .../ai-sandbox-local-process/src/handle.ts | 25 + .../tests/local-process.test.ts | 45 +- packages/ai-sandbox-sprites/src/handle.ts | 36 + .../ai-sandbox-sprites/tests/lstat.test.ts | 201 ++ packages/ai-sandbox-vercel/src/handle.ts | 36 + .../ai-sandbox-vercel/tests/lstat.test.ts | 181 ++ packages/ai-sandbox/README.md | 30 + packages/ai-sandbox/package.json | 5 + .../ai-sandbox/skills/ai-sandbox/SKILL.md | 70 +- packages/ai-sandbox/src/checkpoint-store.ts | 652 +++++ packages/ai-sandbox/src/contracts.ts | 12 + packages/ai-sandbox/src/index.ts | 48 + .../ai-sandbox/src/memory-snapshot-types.ts | 167 ++ packages/ai-sandbox/src/memory-snapshots.ts | 913 +++++++ packages/ai-sandbox/src/middleware.ts | 770 ++++-- packages/ai-sandbox/src/sandbox.ts | 113 +- .../ai-sandbox/src/snapshot-operations.ts | 424 ++++ packages/ai-sandbox/src/snapshots.ts | 670 +++++ .../src/testkit/checkpoint-conformance.ts | 472 ++++ .../testkit/checkpoint-fork-conformance.ts | 299 +++ .../ai-sandbox/src/testkit/conformance.ts | 7 + .../tests/ai-middleware-subpath.test.ts | 6 + .../checkpoint-store.conformance.test.ts | 18 + .../ai-sandbox/tests/checkpoint-store.test.ts | 677 ++++++ packages/ai-sandbox/tests/fakes.test.ts | 17 + packages/ai-sandbox/tests/fakes.ts | 20 +- .../memory-snapshots-declaration.test-d.ts | 29 + .../tests/memory-snapshots-import.test.ts | 16 + .../tests/memory-snapshots.behavior.test.ts | 744 ++++++ .../tests/root-declaration-consumer.test.ts | 89 + .../tests/snapshot-lifecycle.test.ts | 2146 +++++++++++++++++ .../tests/snapshot-operations.test-d.ts | 33 + .../tests/snapshot-operations.test.ts | 1338 ++++++++++ .../tests/snapshot-policy-export.test-d.ts | 28 + packages/ai-sandbox/tests/snapshots.test.ts | 2037 ++++++++++++++++ .../ai-sandbox/tests/testkit-subpath.test.ts | 21 +- packages/ai/src/middlewares/index.ts | 2 + packages/ai/tests/middlewares/index.test.ts | 17 + pnpm-lock.yaml | 3 + testing/e2e/src/routeTree.gen.ts | 22 + .../routes/api.sandbox-file-persistence.ts | 463 ++++ .../tests/sandbox-file-persistence.spec.ts | 60 + 79 files changed, 16109 insertions(+), 277 deletions(-) create mode 100644 .changeset/bright-cloudflare-snapshots.md create mode 100644 .changeset/calm-daytona-snapshots.md create mode 100644 .changeset/direct-docker-snapshots.md create mode 100644 .changeset/eager-local-process-snapshots.md create mode 100644 .changeset/fresh-sprites-snapshots.md create mode 100644 .changeset/fuzzy-snapshots-build.md create mode 100644 .changeset/gentle-vercel-snapshots.md create mode 100644 .changeset/quiet-sandbox-middleware.md create mode 100644 .changeset/tidy-artifact-history.md create mode 100644 docs/sandbox/portable-snapshots.md create mode 100644 packages/ai-persistence/tests/artifact-thread.test.ts create mode 100644 packages/ai-persistence/tests/persistence-completion.test.ts create mode 100644 packages/ai-sandbox-daytona/tests/lstat.test.ts create mode 100644 packages/ai-sandbox-docker/tests/lstat-shell-protocol.test.ts create mode 100644 packages/ai-sandbox-docker/tests/lstat.test.ts create mode 100644 packages/ai-sandbox-docker/tests/testkit-subpath.test.ts create mode 100644 packages/ai-sandbox-sprites/tests/lstat.test.ts create mode 100644 packages/ai-sandbox-vercel/tests/lstat.test.ts create mode 100644 packages/ai-sandbox/src/checkpoint-store.ts create mode 100644 packages/ai-sandbox/src/memory-snapshot-types.ts create mode 100644 packages/ai-sandbox/src/memory-snapshots.ts create mode 100644 packages/ai-sandbox/src/snapshot-operations.ts create mode 100644 packages/ai-sandbox/src/snapshots.ts create mode 100644 packages/ai-sandbox/src/testkit/checkpoint-conformance.ts create mode 100644 packages/ai-sandbox/src/testkit/checkpoint-fork-conformance.ts create mode 100644 packages/ai-sandbox/tests/ai-middleware-subpath.test.ts create mode 100644 packages/ai-sandbox/tests/checkpoint-store.conformance.test.ts create mode 100644 packages/ai-sandbox/tests/checkpoint-store.test.ts create mode 100644 packages/ai-sandbox/tests/fakes.test.ts create mode 100644 packages/ai-sandbox/tests/memory-snapshots-declaration.test-d.ts create mode 100644 packages/ai-sandbox/tests/memory-snapshots-import.test.ts create mode 100644 packages/ai-sandbox/tests/memory-snapshots.behavior.test.ts create mode 100644 packages/ai-sandbox/tests/root-declaration-consumer.test.ts create mode 100644 packages/ai-sandbox/tests/snapshot-lifecycle.test.ts create mode 100644 packages/ai-sandbox/tests/snapshot-operations.test-d.ts create mode 100644 packages/ai-sandbox/tests/snapshot-operations.test.ts create mode 100644 packages/ai-sandbox/tests/snapshot-policy-export.test-d.ts create mode 100644 packages/ai-sandbox/tests/snapshots.test.ts create mode 100644 packages/ai/tests/middlewares/index.test.ts create mode 100644 testing/e2e/src/routes/api.sandbox-file-persistence.ts create mode 100644 testing/e2e/tests/sandbox-file-persistence.spec.ts diff --git a/.changeset/bright-cloudflare-snapshots.md b/.changeset/bright-cloudflare-snapshots.md new file mode 100644 index 0000000000..e5052fd239 --- /dev/null +++ b/.changeset/bright-cloudflare-snapshots.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-sandbox-cloudflare': patch +--- + +Support filesystem metadata needed by portable sandbox snapshots. diff --git a/.changeset/calm-daytona-snapshots.md b/.changeset/calm-daytona-snapshots.md new file mode 100644 index 0000000000..a41ab62fd6 --- /dev/null +++ b/.changeset/calm-daytona-snapshots.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-sandbox-daytona': patch +--- + +Support filesystem metadata needed by portable sandbox snapshots. diff --git a/.changeset/direct-docker-snapshots.md b/.changeset/direct-docker-snapshots.md new file mode 100644 index 0000000000..3943fd89e3 --- /dev/null +++ b/.changeset/direct-docker-snapshots.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-sandbox-docker': patch +--- + +Support filesystem metadata needed by portable sandbox snapshots. diff --git a/.changeset/eager-local-process-snapshots.md b/.changeset/eager-local-process-snapshots.md new file mode 100644 index 0000000000..91fb7e8af7 --- /dev/null +++ b/.changeset/eager-local-process-snapshots.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-sandbox-local-process': patch +--- + +Support filesystem metadata needed by portable sandbox snapshots. diff --git a/.changeset/fresh-sprites-snapshots.md b/.changeset/fresh-sprites-snapshots.md new file mode 100644 index 0000000000..deeaa5d6c6 --- /dev/null +++ b/.changeset/fresh-sprites-snapshots.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-sandbox-sprites': patch +--- + +Support filesystem metadata needed by portable sandbox snapshots. diff --git a/.changeset/fuzzy-snapshots-build.md b/.changeset/fuzzy-snapshots-build.md new file mode 100644 index 0000000000..b8c6c25f13 --- /dev/null +++ b/.changeset/fuzzy-snapshots-build.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-sandbox': minor +--- + +Add portable sandbox checkpoints, named snapshot saves, selected-checkpoint forks, and checkpoint artifact reads. diff --git a/.changeset/gentle-vercel-snapshots.md b/.changeset/gentle-vercel-snapshots.md new file mode 100644 index 0000000000..4788289100 --- /dev/null +++ b/.changeset/gentle-vercel-snapshots.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-sandbox-vercel': patch +--- + +Support filesystem metadata needed by portable sandbox snapshots. diff --git a/.changeset/quiet-sandbox-middleware.md b/.changeset/quiet-sandbox-middleware.md new file mode 100644 index 0000000000..73f6c75138 --- /dev/null +++ b/.changeset/quiet-sandbox-middleware.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai': patch +--- + +Export `CapabilityRegistry` from `@tanstack/ai/middlewares`. diff --git a/.changeset/tidy-artifact-history.md b/.changeset/tidy-artifact-history.md new file mode 100644 index 0000000000..d4e6eb4d31 --- /dev/null +++ b/.changeset/tidy-artifact-history.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-persistence': minor +--- + +Add complete thread artifact history for portable sandbox snapshots. diff --git a/docs/config.json b/docs/config.json index ea417c3de3..4fad50ca3e 100644 --- a/docs/config.json +++ b/docs/config.json @@ -298,7 +298,8 @@ { "label": "Build a Generation Adapter", "to": "persistence/build-your-own-generation-adapter", - "addedAt": "2026-08-04" + "addedAt": "2026-08-04", + "updatedAt": "2026-08-13" }, { "label": "Build a Sandbox Adapter", @@ -308,7 +309,8 @@ { "label": "Store Reference", "to": "persistence/store-reference", - "addedAt": "2026-08-04" + "addedAt": "2026-08-04", + "updatedAt": "2026-08-13" }, { "label": "How Persistence Works", @@ -509,7 +511,7 @@ "label": "Overview", "to": "sandbox/overview", "addedAt": "2026-06-16", - "updatedAt": "2026-08-12" + "updatedAt": "2026-08-13" }, { "label": "Quick Start", @@ -521,7 +523,7 @@ "label": "Providers", "to": "sandbox/providers", "addedAt": "2026-06-29", - "updatedAt": "2026-08-12" + "updatedAt": "2026-08-13" }, { "label": "Harnesses", @@ -551,12 +553,18 @@ "label": "Lifecycle & Snapshots", "to": "sandbox/lifecycle", "addedAt": "2026-06-29", - "updatedAt": "2026-08-12" + "updatedAt": "2026-08-13" + }, + { + "label": "Portable Sandbox Snapshots", + "to": "sandbox/portable-snapshots", + "addedAt": "2026-08-13" }, { "label": "Instance Durability", "to": "sandbox/durability", - "addedAt": "2026-08-04" + "addedAt": "2026-08-04", + "updatedAt": "2026-08-13" }, { "label": "Durable Runs", diff --git a/docs/persistence/build-your-own-generation-adapter.md b/docs/persistence/build-your-own-generation-adapter.md index 0667ebe2cd..a6f1542c18 100644 --- a/docs/persistence/build-your-own-generation-adapter.md +++ b/docs/persistence/build-your-own-generation-adapter.md @@ -54,6 +54,10 @@ CREATE TABLE IF NOT EXISTS artifacts ( source_url text, created_at integer NOT NULL ); +CREATE INDEX IF NOT EXISTS artifacts_run_order + ON artifacts (run_id, created_at, artifact_id); +CREATE INDEX IF NOT EXISTS artifacts_thread_order + ON artifacts (thread_id, created_at, artifact_id); CREATE TABLE IF NOT EXISTS blobs ( key text PRIMARY KEY NOT NULL, bytes blob NOT NULL, @@ -218,6 +222,9 @@ function createGenerationRunStore(db: DatabaseSync) { - `save` is an upsert. - `list(runId)` returns every artifact for a run, `[]` when there are none. +- `listForThread(threadId)` returns every artifact for the thread in exact + `(createdAt, artifactId)` ascending order. It must return the complete thread + history, not a page or only the latest run. Snapshot capture uses this cut. - `delete` / `deleteForRun` are required. Retention and erasure are the point of storing media durably, and they mirror `BlobStore.delete`. @@ -261,7 +268,10 @@ function createArtifactStore(db: DatabaseSync) { ) const selectOne = db.prepare('SELECT * FROM artifacts WHERE artifact_id = ?') const byRun = db.prepare( - 'SELECT * FROM artifacts WHERE run_id = ? ORDER BY created_at ASC', + 'SELECT * FROM artifacts WHERE run_id = ? ORDER BY created_at ASC, artifact_id ASC', + ) + const byThread = db.prepare( + 'SELECT * FROM artifacts WHERE thread_id = ? ORDER BY created_at ASC, artifact_id ASC', ) return defineArtifactStore({ async save(record) { @@ -284,6 +294,9 @@ function createArtifactStore(db: DatabaseSync) { async list(runId) { return byRun.all(runId).map(mapArtifact) }, + async listForThread(threadId) { + return byThread.all(threadId).map(mapArtifact) + }, async delete(artifactId) { db.prepare('DELETE FROM artifacts WHERE artifact_id = ?').run(artifactId) }, diff --git a/docs/persistence/store-reference.md b/docs/persistence/store-reference.md index fd0633f425..41e24e1c81 100644 --- a/docs/persistence/store-reference.md +++ b/docs/persistence/store-reference.md @@ -337,14 +337,20 @@ interface ArtifactRecord { } interface ArtifactStore { - save(record: ArtifactRecord): Promise - get(artifactId: string): Promise - list(runId: string): Promise> // [] when the run has none - delete(artifactId: string): Promise - deleteForRun(runId: string): Promise + save: (record: ArtifactRecord) => Promise + get: (artifactId: string) => Promise + list: (runId: string) => Promise> // [] when the run has none + // Complete thread history, ordered by (createdAt, artifactId) ascending. + listForThread: (threadId: string) => Promise> + delete: (artifactId: string) => Promise + deleteForRun: (runId: string) => Promise } ``` +`list` and `listForThread` use `createdAt` first, then ordinal bytewise +`artifactId` order. Compare UTF-8 bytes from left to right. Do not use locale +collation. + ## BlobStore A durable object/blob store for the bytes. `withGenerationPersistence` writes diff --git a/docs/sandbox/durability.md b/docs/sandbox/durability.md index 55a43396b5..e1a7e3720a 100644 --- a/docs/sandbox/durability.md +++ b/docs/sandbox/durability.md @@ -20,6 +20,10 @@ owned by `@tanstack/ai-sandbox`, independent of `@tanstack/ai-persistence` (transcript / runs / interrupts). You may share a database with chat stores, but you compose a separate middleware. +Instance durability does not copy the workspace into your application storage. +Use [Portable Sandbox Snapshots](./portable-snapshots) when a new sandbox must +rebuild completed files, artifacts, and the saved conversation. + It is also not the agent's *output*. This page keeps a sandbox findable across processes; keeping the run's event stream readable across processes is [The Run Journal](./journal). The two compose: a resumed sandbox still holds the diff --git a/docs/sandbox/lifecycle.md b/docs/sandbox/lifecycle.md index 6298893a21..c92cfd90a2 100644 --- a/docs/sandbox/lifecycle.md +++ b/docs/sandbox/lifecycle.md @@ -10,6 +10,11 @@ Bootstrapping a sandbox (cloning the repo, installing dependencies, running cost once and reuse the result: keep one sandbox per thread, snapshot it after setup, and resume instead of re-bootstrapping on the next run. +When you must also recover files after a sandbox is gone, configure +[Portable Sandbox Snapshots](./portable-snapshots). Provider-native snapshots +make bootstrap faster. Portable snapshots save the completed workspace as +durable application data. + ```ts import { defineSandbox, defineWorkspace, githubRepo } from '@tanstack/ai-sandbox' import { dockerSandbox } from '@tanstack/ai-sandbox-docker' @@ -124,6 +129,10 @@ Each step falls through to the next only when the prior one is unavailable. This is what turns a warm thread into a near-instant start, and a cold one into a full bootstrap. +Portable sandbox snapshots run after this lifecycle work. They restore a saved +workspace only into a newly created private sandbox. They never overwrite a +live resumed sandbox. See [Portable Sandbox Snapshots](./portable-snapshots). + > Which providers support durable disk, snapshots, and resume-by-id is listed on > [Providers](./providers). diff --git a/docs/sandbox/overview.md b/docs/sandbox/overview.md index 79259a6b84..8d1898ea07 100644 --- a/docs/sandbox/overview.md +++ b/docs/sandbox/overview.md @@ -116,6 +116,8 @@ After that, pick the piece you need: - [Tools](./tools): bridge your app's own tools into the in-sandbox agent. - [Policy](./policy): allow, ask or deny guardrails on what the agent may run. - [Lifecycle & Snapshots](./lifecycle): reuse a sandbox, snapshot after setup, resume. +- [Portable Sandbox Snapshots](./portable-snapshots): save completed files and + artifacts, then rebuild them in a new sandbox. - [Instance Durability](./durability): reuse it across replicas too. - [Durable Runs](./durable-runs): let a run outlive the tab, and turn it on. - [Events](./events): stream the agent's edits and tool calls to a UI, and choose what diff --git a/docs/sandbox/portable-snapshots.md b/docs/sandbox/portable-snapshots.md new file mode 100644 index 0000000000..bcf62879de --- /dev/null +++ b/docs/sandbox/portable-snapshots.md @@ -0,0 +1,353 @@ +--- +title: Portable Sandbox Snapshots +id: portable-snapshots +order: 10 +description: "Store a completed sandbox workspace as durable files, artifacts, and conversation data, then rebuild it in a new sandbox." +--- + +An agent can finish work in a sandbox, then the sandbox can disappear. You can +need the same files when the page reloads or when a later run starts. Portable +sandbox snapshots save the completed workspace in your persistence stores. + +This feature saves a checkpoint after each successful terminal run. A later +run restores that checkpoint into a new private sandbox before application code +can use the sandbox. You can also save a named checkpoint and fork from one +selected checkpoint. + +The snapshot helpers run on the server. Your client calls routes that you own. +Do not expose a checkpoint, thread, or artifact id as authorization. + +## Configure snapshots + +Create one persistence value and use that exact value in both middleware calls. +Put `withPersistence` before `withSandbox`. + +```ts +import { chat } from '@tanstack/ai' +import { grokBuildText } from '@tanstack/ai-grok-build' +import { withPersistence } from '@tanstack/ai-persistence' +import { + defineSandbox, + defineWorkspace, + InMemorySandboxInstanceStore, + memorySandboxSnapshots, + withSandbox, +} from '@tanstack/ai-sandbox' +import { dockerSandbox } from '@tanstack/ai-sandbox-docker' + +const snapshots = await memorySandboxSnapshots() +const instances = new InMemorySandboxInstanceStore() +const userId = 'user-123' // Read this from the server session. + +const sandbox = defineSandbox({ + id: 'app-builder', + provider: dockerSandbox({ image: 'node:22' }), + workspace: defineWorkspace({ source: { type: 'none' } }), + lifecycle: { reuse: 'thread' }, +}) + +const result = chat({ + threadId: 'app-thread', + context: { userId }, + adapter: grokBuildText('grok-build'), + messages: [{ role: 'user', content: 'Create a landing page.' }], + middleware: [ + withPersistence(snapshots.persistence), + withSandbox(sandbox, { + instances, + snapshots: { + persistence: snapshots.persistence, + checkpoints: snapshots.checkpoints, + }, + }), + ], +}) + +void result +``` + +`memorySandboxSnapshots()` creates an in-memory persistence object directly. It +does not load `@tanstack/ai-persistence` at runtime. + +Keep `instances` in the same server module as this middleware. A named save +must use this same store. Pass the session `userId` in `context` for every run. +Pass that same user id as `tenant.userId` to `saveNamedSandboxSnapshot`. + +Use a durable persistence implementation and checkpoint store in production. +The memory factory is useful for local development and examples only. + +The optional `policy` controls which workspace paths become files in a +checkpoint. See [Snapshot safety](#snapshot-safety) before you replace the +default policy. + +## Save a named checkpoint + +Automatic saves protect each completed run. Use a named save when a user marks +one workspace state, such as a release candidate. The helper requires a live, +reusable sandbox for the thread. A lifecycle with `reuse: 'none'` cannot create +a named checkpoint. + +Keep this route on the server. Derive the owner from the session. Then make sure +that the owner can access the thread before you call the helper. + +```ts +import { saveNamedSandboxSnapshot } from '@tanstack/ai-sandbox' +import { sandbox, snapshots, instances } from './sandbox-server' + +export async function POST(request: Request) { + const session = await requireSession(request) + const { threadId, runId, label } = await request.json() + + if ( + typeof threadId !== 'string' || + typeof runId !== 'string' || + typeof label !== 'string' + ) { + return new Response('Invalid request', { status: 400 }) + } + if (!(await session.canAccessThread(threadId))) { + return new Response('Not found', { status: 404 }) + } + + const checkpoint = await saveNamedSandboxSnapshot({ + definition: sandbox, + threadId, + runId, + instances, + snapshots, + label, + tenant: { userId: session.userId }, + }) + return Response.json({ checkpointId: checkpoint.id, label: checkpoint.label }) +} +``` + +The client sends its request to this route. It does not call the helper or use +the persistence stores directly. + +```ts +export async function saveCheckpoint( + threadId: string, + runId: string, + label: string, +) { + const response = await fetch('/api/snapshots/save', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ threadId, runId, label }), + }) + if (!response.ok) throw new Error('Could not save checkpoint') + return response.json() +} +``` + +## Fork from a selected checkpoint + +Forking copies the selected checkpoint into an empty destination thread. It does +not fork from the latest checkpoint unless you pass that checkpoint id. The +source thread and its messages remain unchanged. + +Your checkpoint store must implement `forkFromCheckpoint`. The operation must +atomically copy the selected checkpoint, the conversation, the head, and blob +reference counts. It must reject a non-empty destination thread. + +```ts +import { forkFromSandboxSnapshot } from '@tanstack/ai-sandbox' +import { snapshots } from './sandbox-server' + +export async function POST(request: Request) { + const session = await requireSession(request) + const { sourceThreadId, sourceCheckpointId, destinationThreadId } = + await request.json() + + if ( + typeof sourceThreadId !== 'string' || + typeof sourceCheckpointId !== 'string' || + typeof destinationThreadId !== 'string' + ) { + return new Response('Invalid request', { status: 400 }) + } + if (!(await session.canAccessThread(sourceThreadId))) { + return new Response('Not found', { status: 404 }) + } + if (!(await session.canCreateThread(destinationThreadId))) { + return new Response('Not found', { status: 404 }) + } + + const checkpoint = await forkFromSandboxSnapshot({ + sourceThreadId, + sourceCheckpointId, + destinationThreadId, + snapshots, + }) + + return Response.json({ checkpointId: checkpoint.id }) +} +``` + +Use the same server authorization boundary for both source and destination +threads. Do not accept a client-selected checkpoint as proof of access. + +```ts +export async function forkCheckpoint( + sourceThreadId: string, + sourceCheckpointId: string, + destinationThreadId: string, +) { + const response = await fetch('/api/snapshots/fork', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + sourceThreadId, + sourceCheckpointId, + destinationThreadId, + }), + }) + if (!response.ok) throw new Error('Could not fork checkpoint') + return response.json() +} +``` + +## Read a snapshot artifact + +`resolveSnapshotArtifact` reads copied artifact bytes from one checkpoint. It +checks that the checkpoint belongs to the supplied thread. Your route must still +authorize that thread before it calls the helper. The helper returns metadata +and `Uint8Array` bytes. It does not create an HTTP response. + +```ts +import { resolveSnapshotArtifact } from '@tanstack/ai-sandbox' +import { snapshots } from './sandbox-server' + +export async function GET(request: Request) { + const session = await requireSession(request) + const url = new URL(request.url) + const threadId = url.searchParams.get('threadId') + const checkpointId = url.searchParams.get('checkpointId') + const artifactId = url.searchParams.get('artifactId') + + if (!threadId || !checkpointId || !artifactId) { + return new Response('Not found', { status: 404 }) + } + if (!(await session.canAccessThread(threadId))) { + return new Response('Not found', { status: 404 }) + } + + const { artifact, bytes } = await resolveSnapshotArtifact({ + threadId, + checkpointId, + artifactId, + snapshots, + }) + return new Response(bytes, { + headers: { + 'content-type': artifact.mimeType, + 'content-length': String(artifact.size), + }, + }) +} +``` + +The client can use the authorized route as an artifact URL. It must not read the +blob store or call `resolveSnapshotArtifact` in the browser. + +```ts +export function snapshotArtifactUrl( + threadId: string, + checkpointId: string, + artifactId: string, +) { + const query = new URLSearchParams({ threadId, checkpointId, artifactId }) + return `/api/snapshots/artifact?${query}` +} +``` + +## What completes and restores + +After a successful terminal run, the middleware waits for persistence to save +the conversation. It then creates one immutable checkpoint for the thread. + +The checkpoint contains: + +- Regular workspace files. +- Empty directories. +- Generated artifacts that already belong to the thread. +- The saved conversation for the thread. + +File data and copied artifact data use separate content-addressed blob +namespaces. Equal file data deduplicates with file data. Equal artifact data +deduplicates with artifact data. The system does not delete unused blobs +automatically yet. + +On a later run, the middleware uses the latest checkpoint only when it has a +new private sandbox. It restores the files after bootstrap and before the +sandbox is exposed to hooks or the harness. + +A live resumed sandbox is never overwritten. Provider-native snapshots can +make bootstrap faster. Portable checkpoints rebuild the durable workspace when +there is no live sandbox to resume. See [Lifecycle & Snapshots](./lifecycle) +for provider-native snapshot behavior. + +Portable snapshots do not restore into a live sandbox. The next private sandbox +gets the latest checkpoint after bootstrap. A named checkpoint remains available +for reading or for a selected fork. It does not change automatic restore. + +The conversation comes from the durable message store, not from the sandbox +journal. The journal remains a run-output log. See [The Run Journal](./journal) +when you need to replay agent output. + +## Snapshot safety + +Portable snapshots support regular files and directories only. A capture or +restore fails safely when it finds a symlink, an executable file, or a special +filesystem entry. + +The default policy excludes these paths at every depth: + +- `.git` +- `node_modules` +- `.env*` +- The workspace projection marker, `.tanstack-projected-` + +Resolved secret values are replaced with zero bytes before their content is +hashed or stored. A custom policy cannot capture or restore the exact projection +marker for the workspace. + +An invalid manifest, a missing blob, or changed blob content stops the restore +before it writes the workspace. The failed private sandbox is then discarded. +Your existing resumed sandbox remains unchanged. + +## Durable SQLite store + +`memorySandboxSnapshots()` is for tests and local examples. A production store +needs durable message, artifact, and blob stores, plus a durable checkpoint +store. The SQLite example exports `sqliteSandboxSnapshots()` for a Node 22.5+ +server. It is an example adapter, not a package export. + +Use one SQLite transaction for every checkpoint write, head update, and blob +reference count update. Use one transaction for a fork. The fork transaction +must also copy the source conversation and reject a destination thread that has +any persisted state. A partial transaction can create a checkpoint that points +to missing data or a wrong blob reference count. + +## Test route + +`testing/e2e/src/routes/api.sandbox-file-persistence.ts` is a test-only route. +It uses in-memory stores and a fake provider. Do not copy this route into an +application. Use your authenticated server routes and durable stores instead. + +## Operations + +Each thread has one checkpoint writer lease. A second run for the same thread +gets a writer conflict while the existing lease is active. The middleware +renews its lease while the run is active. + +Pause and detach paths release the lease. They do not publish a partial +checkpoint. If the writer loses its lease, the middleware does not publish the +checkpoint. A later successful run can create a new checkpoint. + +Portable snapshots work with [Instance Durability](./durability). Instance +durability finds a provider sandbox across server processes. Portable snapshots +rebuild the workspace when that sandbox is unavailable. + +See [Providers](./providers) for provider-native snapshot and resume support. diff --git a/docs/sandbox/providers.md b/docs/sandbox/providers.md index be63e7f340..3cb342cbaf 100644 --- a/docs/sandbox/providers.md +++ b/docs/sandbox/providers.md @@ -12,6 +12,10 @@ it are provider-agnostic. Pick a provider for the isolation, auth, and snapshot/resume behaviour you need; the rest of your sandbox definition stays the same. +Provider-native snapshots and resume keep or recreate provider state. They can +reduce bootstrap time. [Portable Sandbox Snapshots](./portable-snapshots) store +completed workspace data in your application persistence for reconstruction. + > The provider is _where_ the agent runs. For _which_ agent runs (Grok Build, > Claude Code, Codex, OpenCode, or any ACP agent via `acpCompatible`) see > [Harnesses](./harnesses). diff --git a/examples/ts-react-chat/src/lib/sqlite-persistence.test.ts b/examples/ts-react-chat/src/lib/sqlite-persistence.test.ts index aaebaa2339..299f767bfb 100644 --- a/examples/ts-react-chat/src/lib/sqlite-persistence.test.ts +++ b/examples/ts-react-chat/src/lib/sqlite-persistence.test.ts @@ -10,7 +10,11 @@ import { join } from 'node:path' import { DatabaseSync } from 'node:sqlite' import { describe, expect, it } from 'vitest' import { runPersistenceConformance } from '@tanstack/ai-persistence/testkit' -import { sqlitePersistence } from './sqlite-persistence' +import { + runSandboxCheckpointForkConformance, + runSandboxCheckpointStoreConformance, +} from '@tanstack/ai-sandbox/testkit' +import { sqlitePersistence, sqliteSandboxSnapshots } from './sqlite-persistence' // All seven stores are provided — the four chat state stores plus // `generationRuns` + `artifacts` + `blobs` — so no STORE is skipped. One @@ -29,6 +33,307 @@ runPersistenceConformance( { skipMethods: ['runs.listByThread'] }, ) +runSandboxCheckpointStoreConformance( + 'ts-react-chat example (node:sqlite)', + (options) => + sqliteSandboxSnapshots({ + url: ':memory:', + migrate: true, + ...options, + }).checkpoints, +) + +runSandboxCheckpointForkConformance( + 'ts-react-chat example (node:sqlite)', + () => { + const snapshots = sqliteSandboxSnapshots({ + url: ':memory:', + migrate: true, + }) + return snapshots + }, +) + +describe('sqliteSandboxSnapshots fork transaction', () => { + it('uses one durable writer lease across two SQLite connections', async () => { + const dir = mkdtempSync(join(tmpdir(), 'tanstack-sqlite-lease-')) + const file = join(dir, 'snapshots.db') + const first = sqliteSandboxSnapshots({ url: file, migrate: true }) + const second = sqliteSandboxSnapshots({ url: file, migrate: true }) + try { + const writer = await first.checkpoints.acquireWriter('thread') + await expect( + second.checkpoints.acquireWriter('thread'), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_WRITER_CONFLICT' }) + await writer.release() + const replacement = await second.checkpoints.acquireWriter('thread') + await expect( + first.checkpoints.append({ + checkpoint: { + id: 'stale', + threadId: 'thread', + parentCheckpointId: null, + createdAt: 1, + reason: 'named', + files: [], + conversation: [], + artifacts: [], + }, + expectedHeadId: null, + writer, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_WRITER_LOST' }) + await replacement.release() + } finally { + first.close() + second.close() + rmSync(dir, { recursive: true, force: true }) + } + }) + + it('deletes dependent rows and permits the checkpoint id to be reused', async () => { + const snapshots = sqliteSandboxSnapshots({ url: ':memory:', migrate: true }) + const writer = await snapshots.checkpoints.acquireWriter('thread') + const checkpoint = { + id: 'reusable', + threadId: 'thread', + parentCheckpointId: null, + createdAt: 1, + reason: 'named' as const, + files: [ + { + path: 'a.txt', + kind: 'file' as const, + blobKey: `sandbox-files/sha256/${'a'.repeat(64)}`, + size: 1, + }, + ], + conversation: [{ role: 'user' as const, content: 'one' }], + artifacts: [], + } + try { + await snapshots.checkpoints.append({ + checkpoint, + expectedHeadId: null, + writer, + }) + await snapshots.checkpoints.deleteHead({ + threadId: 'thread', + checkpointId: 'reusable', + writer, + }) + await expect( + snapshots.checkpoints.append({ + checkpoint, + expectedHeadId: null, + writer, + }), + ).resolves.toEqual({ headId: 'reusable' }) + } finally { + snapshots.close() + } + }) + + it('rejects a fork whose source blob reference is missing', async () => { + const dir = mkdtempSync(join(tmpdir(), 'tanstack-sqlite-ref-')) + const file = join(dir, 'snapshots.db') + const snapshots = sqliteSandboxSnapshots({ url: file, migrate: true }) + try { + const key = `sandbox-files/sha256/${'b'.repeat(64)}` + const sourceWriter = await snapshots.checkpoints.acquireWriter('source') + await snapshots.checkpoints.append({ + checkpoint: { + id: 'source', + threadId: 'source', + parentCheckpointId: null, + createdAt: 1, + reason: 'named', + files: [{ path: 'a.txt', kind: 'file', blobKey: key, size: 1 }], + conversation: [], + artifacts: [], + }, + expectedHeadId: null, + writer: sourceWriter, + }) + const inspector = new DatabaseSync(file) + inspector + .prepare( + 'DELETE FROM sandbox_checkpoint_blob_references WHERE blob_key = ?', + ) + .run(key) + inspector.close() + const writer = await snapshots.checkpoints.acquireWriter('destination') + await expect( + snapshots.checkpoints.forkFromCheckpoint({ + sourceThreadId: 'source', + sourceCheckpointId: 'source', + destinationThreadId: 'destination', + destinationCheckpointId: 'destination', + createdAt: 2, + writer, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ENTRY' }) + expect( + await snapshots.persistence.stores.messages.loadThread('destination'), + ).toEqual([]) + expect(await snapshots.checkpoints.getHead('destination')).toBeNull() + } finally { + snapshots.close() + rmSync(dir, { recursive: true, force: true }) + } + }) + + it.each([ + { + path: 'file/child', + kind: 'file', + blobKey: `sandbox-files/sha256/${'c'.repeat(64)}`, + size: 1, + }, + { + path: 'C:/absolute.txt', + kind: 'file', + blobKey: `sandbox-files/sha256/${'c'.repeat(64)}`, + size: 1, + }, + { + path: 'bad', + kind: 'other', + blobKey: `sandbox-files/sha256/${'c'.repeat(64)}`, + size: 1, + }, + { path: 'directory', kind: 'dir', blobKey: 'forbidden' }, + ])('rejects invalid checkpoint entry $path', async (invalid) => { + const snapshots = sqliteSandboxSnapshots({ url: ':memory:', migrate: true }) + try { + const writer = await snapshots.checkpoints.acquireWriter('thread') + const files = + invalid.path === 'file/child' + ? [ + { + path: 'file', + kind: 'file' as const, + blobKey: `sandbox-files/sha256/${'d'.repeat(64)}`, + size: 1, + }, + invalid, + ] + : [invalid] + const checkpoint = { + id: `invalid-${invalid.path}`, + threadId: 'thread', + parentCheckpointId: null, + createdAt: 1, + reason: 'named' as const, + files: [], + conversation: [], + artifacts: [], + } + Reflect.set(checkpoint, 'files', files) + await expect( + snapshots.checkpoints.append({ + checkpoint, + expectedHeadId: null, + writer, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ENTRY' }) + } finally { + snapshots.close() + } + }) + + it.each([ + ['files', null], + ['files', 'not-an-array'], + ['files', [null]], + ['files', [1]], + ['artifacts', null], + ['artifacts', 'not-an-array'], + ['artifacts', [null]], + ['artifacts', [1]], + ])('rejects malformed checkpoint %s values', async (field, value) => { + const snapshots = sqliteSandboxSnapshots({ url: ':memory:', migrate: true }) + try { + const writer = await snapshots.checkpoints.acquireWriter('thread') + const checkpoint = { + id: `malformed-${field}-${String(value)}`, + threadId: 'thread', + parentCheckpointId: null, + createdAt: 1, + reason: 'named' as const, + files: [], + conversation: [], + artifacts: [], + } + Reflect.set(checkpoint, field, value) + await expect( + snapshots.checkpoints.append({ + checkpoint, + expectedHeadId: null, + writer, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ENTRY' }) + } finally { + snapshots.close() + } + }) + + it('rolls back the destination when checkpoint storage fails after transcript staging', async () => { + const dir = mkdtempSync(join(tmpdir(), 'tanstack-sqlite-fork-')) + const file = join(dir, 'snapshots.db') + const snapshots = sqliteSandboxSnapshots({ url: file, migrate: true }) + try { + const sourceWriter = await snapshots.checkpoints.acquireWriter('source') + const blobKey = `sandbox-files/sha256/${'a'.repeat(64)}` + await snapshots.checkpoints.append({ + checkpoint: { + id: 'source-root', + threadId: 'source', + parentCheckpointId: null, + createdAt: 1, + reason: 'named', + files: [{ path: 'a.txt', kind: 'file', blobKey, size: 1 }], + conversation: [{ role: 'user', content: 'source' }], + artifacts: [], + }, + expectedHeadId: null, + writer: sourceWriter, + }) + const triggerDb = new DatabaseSync(file) + triggerDb.exec(` + CREATE TRIGGER fail_fork_checkpoint + BEFORE INSERT ON sandbox_checkpoints + WHEN NEW.checkpoint_id = 'fork-root' + BEGIN SELECT RAISE(ABORT, 'forced checkpoint failure'); END; + `) + triggerDb.close() + const destinationWriter = + await snapshots.checkpoints.acquireWriter('destination') + await expect( + snapshots.checkpoints.forkFromCheckpoint({ + sourceThreadId: 'source', + sourceCheckpointId: 'source-root', + destinationThreadId: 'destination', + destinationCheckpointId: 'fork-root', + createdAt: 2, + writer: destinationWriter, + }), + ).rejects.toThrow('forced checkpoint failure') + expect( + await snapshots.persistence.stores.messages.loadThread('destination'), + ).toEqual([]) + expect(await snapshots.checkpoints.list('destination')).toEqual([]) + expect(await snapshots.checkpoints.getHead('destination')).toBeNull() + expect(await snapshots.checkpoints.listBlobReferences()).toEqual([ + { key: blobKey, references: 1 }, + ]) + } finally { + snapshots.close() + rmSync(dir, { recursive: true, force: true }) + } + }) +}) + // SQL-specific case the in-memory reference backend cannot express: a real // query layer can get `NULL <= ?` wrong in ways JS's `undefined <= n` // (`NaN <= n`, always false) never surfaces. This pins the SQLite backend's diff --git a/examples/ts-react-chat/src/lib/sqlite-persistence.ts b/examples/ts-react-chat/src/lib/sqlite-persistence.ts index 7c9ef152f8..6d184d75d9 100644 --- a/examples/ts-react-chat/src/lib/sqlite-persistence.ts +++ b/examples/ts-react-chat/src/lib/sqlite-persistence.ts @@ -59,6 +59,25 @@ import type { RunStatus, RunStore, } from '@tanstack/ai-persistence' +import { + SandboxCheckpointConflictError, + SandboxCheckpointDuplicateIdError, + SandboxCheckpointError, + SandboxCheckpointInvalidEntryError, + SandboxCheckpointInvalidIdError, + SandboxCheckpointNotHeadError, + SandboxCheckpointParentMismatchError, + SandboxCheckpointWriterConflictError, + SandboxCheckpointWriterLostError, +} from '@tanstack/ai-sandbox' +import type { + ForkCapableSandboxCheckpointStore, + SandboxCheckpoint, + SandboxCheckpointStoreOptions, + SandboxCheckpointWriter, + SandboxSnapshotArtifact, + SandboxSnapshotEntry, +} from '@tanstack/ai-sandbox' // --------------------------------------------------------------------------- // Schema @@ -133,7 +152,8 @@ CREATE TABLE IF NOT EXISTS artifacts ( source_url text, created_at integer NOT NULL ); -CREATE INDEX IF NOT EXISTS artifacts_run ON artifacts (run_id); +CREATE INDEX IF NOT EXISTS artifacts_run_order + ON artifacts (run_id, created_at ASC, artifact_id ASC); -- The bytes themselves. \`body\` is a BLOB column, so this file IS the object -- store; a production adapter would keep metadata here and put bytes in S3/R2. CREATE TABLE IF NOT EXISTS blobs ( @@ -146,6 +166,60 @@ CREATE TABLE IF NOT EXISTS blobs ( created_at integer NOT NULL, updated_at integer NOT NULL ); +CREATE TABLE IF NOT EXISTS sandbox_checkpoints ( + checkpoint_id text PRIMARY KEY NOT NULL, + thread_id text NOT NULL, + parent_checkpoint_id text, + created_at integer NOT NULL, + reason text NOT NULL, + label text, + source_run_id text, + conversation_json text NOT NULL +); +CREATE INDEX IF NOT EXISTS sandbox_checkpoints_thread_order + ON sandbox_checkpoints (thread_id, created_at ASC, checkpoint_id COLLATE BINARY ASC); +CREATE TABLE IF NOT EXISTS sandbox_checkpoint_entries ( + checkpoint_id text NOT NULL, + entry_index integer NOT NULL, + path text NOT NULL, + kind text NOT NULL, + blob_key text, + size integer, + PRIMARY KEY (checkpoint_id, entry_index) +); +CREATE INDEX IF NOT EXISTS sandbox_checkpoint_entries_checkpoint + ON sandbox_checkpoint_entries (checkpoint_id, entry_index ASC); +CREATE TABLE IF NOT EXISTS sandbox_checkpoint_artifacts ( + checkpoint_id text NOT NULL, + artifact_index integer NOT NULL, + artifact_id text NOT NULL, + name text NOT NULL, + mime_type text NOT NULL, + size integer NOT NULL, + blob_key text NOT NULL, + created_at integer NOT NULL, + PRIMARY KEY (checkpoint_id, artifact_index) +); +CREATE INDEX IF NOT EXISTS sandbox_checkpoint_artifacts_checkpoint + ON sandbox_checkpoint_artifacts (checkpoint_id, artifact_index ASC); +CREATE TABLE IF NOT EXISTS sandbox_checkpoint_heads ( + thread_id text PRIMARY KEY NOT NULL, + checkpoint_id text NOT NULL +); +CREATE TABLE IF NOT EXISTS sandbox_checkpoint_writers ( + thread_id text PRIMARY KEY NOT NULL, + owner_token text NOT NULL, + fence integer NOT NULL, + expires_at integer NOT NULL +); +CREATE TABLE IF NOT EXISTS sandbox_checkpoint_fences ( + thread_id text PRIMARY KEY NOT NULL, + fence integer NOT NULL +); +CREATE TABLE IF NOT EXISTS sandbox_checkpoint_blob_references ( + blob_key text PRIMARY KEY NOT NULL, + reference_count integer NOT NULL +); ` /** @@ -797,7 +871,10 @@ function createArtifactStore(db: DatabaseSync) { ) const selectStmt = db.prepare('SELECT * FROM artifacts WHERE artifact_id = ?') const byRunStmt = db.prepare( - 'SELECT * FROM artifacts WHERE run_id = ? ORDER BY created_at ASC', + 'SELECT * FROM artifacts WHERE run_id = ? ORDER BY created_at ASC, artifact_id ASC', + ) + const byThreadStmt = db.prepare( + 'SELECT * FROM artifacts WHERE thread_id = ? ORDER BY created_at ASC, artifact_id ASC', ) const deleteStmt = db.prepare('DELETE FROM artifacts WHERE artifact_id = ?') const deleteForRunStmt = db.prepare('DELETE FROM artifacts WHERE run_id = ?') @@ -826,6 +903,10 @@ function createArtifactStore(db: DatabaseSync) { const rows: Array = byRunStmt.all(runId) return Promise.resolve((rows as Array).map(mapArtifact)) }, + listForThread(threadId) { + const rows: Array = byThreadStmt.all(threadId) + return Promise.resolve((rows as Array).map(mapArtifact)) + }, delete(artifactId) { deleteStmt.run(artifactId) return Promise.resolve() @@ -1061,6 +1142,608 @@ export interface SqlitePersistenceOptions { migrate?: boolean } +function checkpointError( + code: ConstructorParameters[0], + message: string, +): never { + throw new SandboxCheckpointError(code, message) +} + +function cloneCheckpoint(value: T): T { + return structuredClone(value) +} + +function checkpointKeys(checkpoint: SandboxCheckpoint): Array { + return [ + ...new Set([ + ...checkpoint.files.flatMap((entry) => + entry.kind === 'file' ? [entry.blobKey] : [], + ), + ...checkpoint.artifacts.map((artifact) => artifact.blobKey), + ]), + ] +} + +function hasUnpairedSurrogate(value: string): boolean { + for (let index = 0; index < value.length; index++) { + const code = value.charCodeAt(index) + if (code >= 0xd800 && code <= 0xdbff) { + const next = value.charCodeAt(index + 1) + if (Number.isNaN(next) || next < 0xdc00 || next > 0xdfff) return true + index++ + } else if (code >= 0xdc00 && code <= 0xdfff) return true + } + return false +} + +function assertCheckpointId(value: string, label: string): void { + if (!value || value.includes('\0') || hasUnpairedSurrogate(value)) { + throw new SandboxCheckpointInvalidIdError( + `${label} must be a non-empty well-formed Unicode string`, + ) + } +} + +function assertCheckpoint(checkpoint: SandboxCheckpoint): void { + assertCheckpointId(checkpoint.id, 'Checkpoint id') + assertCheckpointId(checkpoint.threadId, 'Checkpoint thread id') + if (checkpoint.parentCheckpointId !== null) + assertCheckpointId(checkpoint.parentCheckpointId, 'Parent checkpoint id') + if (!Number.isFinite(checkpoint.createdAt)) + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint createdAt must be a finite number', + ) + if (!Array.isArray(checkpoint.files)) + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint files must be an array', + ) + if (!Array.isArray(checkpoint.artifacts)) + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint artifacts must be an array', + ) + const paths = new Set() + const kinds = new Map() + for (const entry of checkpoint.files) { + if (entry === null || typeof entry !== 'object') + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint entry must be an object', + ) + if ( + !entry.path || + paths.has(entry.path) || + entry.path.includes('\0') || + entry.path.startsWith('/') || + entry.path.startsWith('\\') || + /^[A-Za-z]:([\\/]|$)/.test(entry.path) || + entry.path.includes('\\') || + entry.path + .split('/') + .some((part) => !part || part === '.' || part === '..') + ) + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint entry path must be a normalized workspace-relative path', + ) + for ( + let separator = entry.path.indexOf('/'); + separator !== -1; + separator = entry.path.indexOf('/', separator + 1) + ) { + if (kinds.get(entry.path.slice(0, separator)) === 'file') + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint entry cannot be beneath a file', + ) + } + if ( + entry.kind === 'file' && + [...kinds.keys()].some((path) => path.startsWith(`${entry.path}/`)) + ) + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint file cannot be an ancestor of another entry', + ) + paths.add(entry.path) + if (entry.kind === 'file') { + if ( + !/^sandbox-files\/sha256\/[0-9a-f]{64}$/.test(entry.blobKey) || + !Number.isSafeInteger(entry.size) || + entry.size < 0 + ) { + throw new SandboxCheckpointInvalidEntryError( + 'File entries require a valid blobKey and size', + ) + } + } else if (entry.kind === 'dir') { + if ('blobKey' in entry || 'size' in entry) + throw new SandboxCheckpointInvalidEntryError( + 'Directory entries cannot contain file fields', + ) + } else { + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint entry kind must be file or dir', + ) + } + kinds.set(entry.path, entry.kind) + } + for (const artifact of checkpoint.artifacts) { + if (artifact === null || typeof artifact !== 'object') + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint artifact must be an object', + ) + if ( + !artifact.artifactId || + !artifact.name || + !artifact.mimeType || + !/^sandbox-artifacts\/sha256\/[0-9a-f]{64}$/.test(artifact.blobKey) || + !Number.isSafeInteger(artifact.size) || + artifact.size < 0 || + !Number.isFinite(artifact.createdAt) + ) + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint artifact has invalid fields', + ) + } +} + +function createCheckpointStore( + db: DatabaseSync, + options: SandboxCheckpointStoreOptions = {}, +): ForkCapableSandboxCheckpointStore { + const now = options.now ?? (() => Date.now()) + const leaseDurationMs = options.leaseDurationMs ?? 120_000 + const renewAfterMs = options.renewAfterMs ?? 45_000 + if ( + !Number.isFinite(leaseDurationMs) || + leaseDurationMs <= 0 || + !Number.isFinite(renewAfterMs) || + renewAfterMs <= 0 || + renewAfterMs >= leaseDurationMs + ) + throw new Error('Invalid checkpoint writer lease options') + const getCheckpoint = (id: string): SandboxCheckpoint | null => { + const header = db + .prepare('SELECT * FROM sandbox_checkpoints WHERE checkpoint_id = ?') + .get(id) as + | { + checkpoint_id: string + thread_id: string + parent_checkpoint_id: string | null + created_at: number + reason: SandboxCheckpoint['reason'] + label: string | null + source_run_id: string | null + conversation_json: string + } + | undefined + if (!header) return null + const files = db + .prepare( + 'SELECT * FROM sandbox_checkpoint_entries WHERE checkpoint_id = ? ORDER BY entry_index ASC', + ) + .all(id) as Array<{ + path: string + kind: 'file' | 'dir' + blob_key: string | null + size: number | null + }> + const artifacts = db + .prepare( + 'SELECT * FROM sandbox_checkpoint_artifacts WHERE checkpoint_id = ? ORDER BY artifact_index ASC', + ) + .all(id) as Array<{ + artifact_id: string + name: string + mime_type: string + size: number + blob_key: string + created_at: number + }> + return { + id: header.checkpoint_id, + threadId: header.thread_id, + parentCheckpointId: header.parent_checkpoint_id, + createdAt: header.created_at, + reason: header.reason, + ...(header.label === null ? {} : { label: header.label }), + ...(header.source_run_id === null + ? {} + : { sourceRunId: header.source_run_id }), + files: files.map( + (row): SandboxSnapshotEntry => + row.kind === 'dir' + ? { path: row.path, kind: 'dir' } + : { + path: row.path, + kind: 'file', + blobKey: row.blob_key!, + size: row.size!, + }, + ), + conversation: JSON.parse(header.conversation_json), + artifacts: artifacts.map( + (row): SandboxSnapshotArtifact => ({ + artifactId: row.artifact_id, + name: row.name, + mimeType: row.mime_type, + size: row.size, + blobKey: row.blob_key, + createdAt: row.created_at, + }), + ), + } + } + const assertWriter = (writer: SandboxCheckpointWriter, threadId: string) => { + const row = db + .prepare( + 'SELECT owner_token, fence, expires_at FROM sandbox_checkpoint_writers WHERE thread_id = ?', + ) + .get(threadId) as + | { owner_token: string; fence: number; expires_at: number } + | undefined + if ( + !row || + writer.threadId !== threadId || + row.owner_token !== writer.ownerToken || + row.fence !== writer.fence || + row.expires_at <= now() + ) + throw new SandboxCheckpointWriterLostError( + `Checkpoint writer lease for thread '${threadId}' is no longer current`, + ) + } + const writeCheckpoint = ( + checkpoint: SandboxCheckpoint, + conversationJson: string, + ) => { + db.prepare( + 'INSERT INTO sandbox_checkpoints (checkpoint_id, thread_id, parent_checkpoint_id, created_at, reason, label, source_run_id, conversation_json) VALUES (?, ?, ?, ?, ?, ?, ?, ?)', + ).run( + checkpoint.id, + checkpoint.threadId, + checkpoint.parentCheckpointId, + checkpoint.createdAt, + checkpoint.reason, + checkpoint.label ?? null, + checkpoint.sourceRunId ?? null, + conversationJson, + ) + const entry = db.prepare( + 'INSERT INTO sandbox_checkpoint_entries (checkpoint_id, entry_index, path, kind, blob_key, size) VALUES (?, ?, ?, ?, ?, ?)', + ) + checkpoint.files.forEach((value, index) => + entry.run( + checkpoint.id, + index, + value.path, + value.kind, + value.kind === 'file' ? value.blobKey : null, + value.kind === 'file' ? value.size : null, + ), + ) + const artifact = db.prepare( + 'INSERT INTO sandbox_checkpoint_artifacts (checkpoint_id, artifact_index, artifact_id, name, mime_type, size, blob_key, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)', + ) + checkpoint.artifacts.forEach((value, index) => + artifact.run( + checkpoint.id, + index, + value.artifactId, + value.name, + value.mimeType, + value.size, + value.blobKey, + value.createdAt, + ), + ) + } + const incrementReferences = (checkpoint: SandboxCheckpoint) => { + const statement = db.prepare( + 'INSERT INTO sandbox_checkpoint_blob_references (blob_key, reference_count) VALUES (?, 1) ON CONFLICT(blob_key) DO UPDATE SET reference_count = reference_count + 1', + ) + checkpointKeys(checkpoint).forEach((key) => statement.run(key)) + } + return { + async get(id) { + assertCheckpointId(id, 'Checkpoint id') + const value = getCheckpoint(id) + return value && cloneCheckpoint(value) + }, + async list(threadId) { + assertCheckpointId(threadId, 'Thread id') + const rows = db + .prepare( + 'SELECT checkpoint_id FROM sandbox_checkpoints WHERE thread_id = ? ORDER BY created_at ASC, checkpoint_id COLLATE BINARY ASC', + ) + .all(threadId) as Array<{ checkpoint_id: string }> + return rows.map((row) => + cloneCheckpoint(getCheckpoint(row.checkpoint_id)!), + ) + }, + async getHead(threadId) { + assertCheckpointId(threadId, 'Thread id') + const row = db + .prepare( + 'SELECT checkpoint_id FROM sandbox_checkpoint_heads WHERE thread_id = ?', + ) + .get(threadId) as { checkpoint_id: string } | undefined + return row?.checkpoint_id ?? null + }, + async append(input) { + const checkpoint = cloneCheckpoint(input.checkpoint) + assertCheckpoint(checkpoint) + if (input.expectedHeadId !== null) + assertCheckpointId(input.expectedHeadId, 'Expected head id') + if (input.writer.threadId !== checkpoint.threadId) + throw new SandboxCheckpointWriterLostError( + 'Checkpoint writer thread does not match checkpoint thread', + ) + const conversationJson = JSON.stringify(checkpoint.conversation) + db.exec('BEGIN IMMEDIATE') + try { + assertWriter(input.writer, checkpoint.threadId) + if (getCheckpoint(checkpoint.id)) + throw new SandboxCheckpointDuplicateIdError( + `Checkpoint '${checkpoint.id}' already exists`, + ) + const head = + ( + db + .prepare( + 'SELECT checkpoint_id FROM sandbox_checkpoint_heads WHERE thread_id = ?', + ) + .get(checkpoint.threadId) as { checkpoint_id: string } | undefined + )?.checkpoint_id ?? null + if (head !== input.expectedHeadId) + throw new SandboxCheckpointConflictError( + `Expected head '${input.expectedHeadId}', but thread '${checkpoint.threadId}' is at '${head}'`, + ) + if (checkpoint.parentCheckpointId !== input.expectedHeadId) + throw new SandboxCheckpointParentMismatchError( + `Checkpoint '${checkpoint.id}' parent does not match expected head`, + ) + writeCheckpoint(checkpoint, conversationJson) + db.prepare( + 'INSERT INTO sandbox_checkpoint_heads (thread_id, checkpoint_id) VALUES (?, ?) ON CONFLICT(thread_id) DO UPDATE SET checkpoint_id = excluded.checkpoint_id', + ).run(checkpoint.threadId, checkpoint.id) + incrementReferences(checkpoint) + db.exec('COMMIT') + } catch (error) { + db.exec('ROLLBACK') + throw error + } + return { headId: checkpoint.id } + }, + async deleteHead(input) { + assertCheckpointId(input.threadId, 'Thread id') + assertCheckpointId(input.checkpointId, 'Checkpoint id') + db.exec('BEGIN IMMEDIATE') + try { + assertWriter(input.writer, input.threadId) + const head = + ( + db + .prepare( + 'SELECT checkpoint_id FROM sandbox_checkpoint_heads WHERE thread_id = ?', + ) + .get(input.threadId) as { checkpoint_id: string } | undefined + )?.checkpoint_id ?? null + if (head !== input.checkpointId) + throw new SandboxCheckpointNotHeadError( + `Checkpoint '${input.checkpointId}' is not the current head of thread '${input.threadId}'`, + ) + const checkpoint = getCheckpoint(input.checkpointId) + if (!checkpoint) + throw new SandboxCheckpointNotHeadError( + `Checkpoint '${input.checkpointId}' does not exist`, + ) + db.prepare( + 'DELETE FROM sandbox_checkpoint_entries WHERE checkpoint_id = ?', + ).run(checkpoint.id) + db.prepare( + 'DELETE FROM sandbox_checkpoint_artifacts WHERE checkpoint_id = ?', + ).run(checkpoint.id) + db.prepare( + 'DELETE FROM sandbox_checkpoints WHERE checkpoint_id = ?', + ).run(checkpoint.id) + if (checkpoint.parentCheckpointId) + db.prepare( + 'UPDATE sandbox_checkpoint_heads SET checkpoint_id = ? WHERE thread_id = ?', + ).run(checkpoint.parentCheckpointId, input.threadId) + else + db.prepare( + 'DELETE FROM sandbox_checkpoint_heads WHERE thread_id = ?', + ).run(input.threadId) + const decrement = db.prepare( + 'UPDATE sandbox_checkpoint_blob_references SET reference_count = reference_count - 1 WHERE blob_key = ?', + ) + checkpointKeys(checkpoint).forEach((key) => decrement.run(key)) + db.exec( + 'DELETE FROM sandbox_checkpoint_blob_references WHERE reference_count <= 0', + ) + db.exec('COMMIT') + } catch (error) { + db.exec('ROLLBACK') + throw error + } + }, + async acquireWriter(threadId) { + assertCheckpointId(threadId, 'Thread id') + db.exec('BEGIN IMMEDIATE') + let lease: { ownerToken: string; fence: number; expiresAt: number } + try { + const current = db + .prepare( + 'SELECT expires_at FROM sandbox_checkpoint_writers WHERE thread_id = ?', + ) + .get(threadId) as { expires_at: number } | undefined + if (current && current.expires_at > now()) + throw new SandboxCheckpointWriterConflictError( + `Thread '${threadId}' already has an active checkpoint writer`, + ) + const prior = db + .prepare( + 'SELECT fence FROM sandbox_checkpoint_fences WHERE thread_id = ?', + ) + .get(threadId) as { fence: number } | undefined + const fence = (prior?.fence ?? 0) + 1 + db.prepare( + 'INSERT INTO sandbox_checkpoint_fences (thread_id, fence) VALUES (?, ?) ON CONFLICT(thread_id) DO UPDATE SET fence = excluded.fence', + ).run(threadId, fence) + lease = { + ownerToken: crypto.randomUUID(), + fence, + expiresAt: now() + leaseDurationMs, + } + db.prepare( + 'INSERT INTO sandbox_checkpoint_writers (thread_id, owner_token, fence, expires_at) VALUES (?, ?, ?, ?) ON CONFLICT(thread_id) DO UPDATE SET owner_token = excluded.owner_token, fence = excluded.fence, expires_at = excluded.expires_at', + ).run(threadId, lease.ownerToken, fence, lease.expiresAt) + db.exec('COMMIT') + } catch (error) { + db.exec('ROLLBACK') + throw error + } + return { + threadId, + ownerToken: lease!.ownerToken, + fence: lease!.fence, + get expiresAt() { + return lease!.expiresAt + }, + renewAfterMs, + renew: async () => { + db.exec('BEGIN IMMEDIATE') + try { + assertWriter( + { threadId, ownerToken: lease!.ownerToken, fence: lease!.fence }, + threadId, + ) + const expiresAt = now() + leaseDurationMs + db.prepare( + 'UPDATE sandbox_checkpoint_writers SET expires_at = ? WHERE thread_id = ? AND owner_token = ? AND fence = ?', + ).run(expiresAt, threadId, lease!.ownerToken, lease!.fence) + lease!.expiresAt = expiresAt + db.exec('COMMIT') + return { expiresAt } + } catch (error) { + db.exec('ROLLBACK') + throw error + } + }, + release: async () => { + db.prepare( + 'DELETE FROM sandbox_checkpoint_writers WHERE thread_id = ? AND owner_token = ? AND fence = ?', + ).run(threadId, lease!.ownerToken, lease!.fence) + }, + } + }, + async listBlobReferences() { + return db + .prepare( + 'SELECT blob_key AS key, reference_count AS "references" FROM sandbox_checkpoint_blob_references ORDER BY blob_key COLLATE BINARY ASC', + ) + .all() as Array<{ key: string; references: number }> + }, + async forkFromCheckpoint(input) { + const staged = { + sourceThreadId: input.sourceThreadId, + sourceCheckpointId: input.sourceCheckpointId, + destinationThreadId: input.destinationThreadId, + destinationCheckpointId: input.destinationCheckpointId, + createdAt: input.createdAt, + writer: { + threadId: input.writer.threadId, + ownerToken: input.writer.ownerToken, + fence: input.writer.fence, + }, + } + assertCheckpointId(staged.sourceThreadId, 'Source thread id') + assertCheckpointId(staged.sourceCheckpointId, 'Source checkpoint id') + assertCheckpointId(staged.destinationThreadId, 'Destination thread id') + assertCheckpointId( + staged.destinationCheckpointId, + 'Destination checkpoint id', + ) + if (!Number.isFinite(staged.createdAt)) + throw new SandboxCheckpointInvalidEntryError( + 'Fork checkpoint createdAt must be a finite number', + ) + db.exec('BEGIN IMMEDIATE') + try { + if (staged.sourceThreadId === staged.destinationThreadId) + checkpointError( + 'SANDBOX_SNAPSHOT_FORK_SOURCE_THREAD_MISMATCH', + 'Source and destination threads must differ', + ) + const source = getCheckpoint(staged.sourceCheckpointId) + if (!source) + checkpointError( + 'SANDBOX_SNAPSHOT_FORK_SOURCE_NOT_FOUND', + 'Source checkpoint was not found', + ) + if (source.threadId !== staged.sourceThreadId) + checkpointError( + 'SANDBOX_SNAPSHOT_FORK_SOURCE_THREAD_MISMATCH', + 'Source checkpoint belongs to another thread', + ) + assertWriter(staged.writer, staged.destinationThreadId) + const nonempty = db + .prepare( + 'SELECT 1 FROM messages WHERE thread_id = ? UNION ALL SELECT 1 FROM runs WHERE thread_id = ? UNION ALL SELECT 1 FROM generation_runs WHERE thread_id = ? UNION ALL SELECT 1 FROM interrupts WHERE thread_id = ? UNION ALL SELECT 1 FROM artifacts WHERE thread_id = ? UNION ALL SELECT 1 FROM sandbox_checkpoints WHERE thread_id = ? UNION ALL SELECT 1 FROM sandbox_checkpoint_heads WHERE thread_id = ? UNION ALL SELECT 1 FROM sandbox_checkpoints WHERE checkpoint_id = ? LIMIT 1', + ) + .get( + staged.destinationThreadId, + staged.destinationThreadId, + staged.destinationThreadId, + staged.destinationThreadId, + staged.destinationThreadId, + staged.destinationThreadId, + staged.destinationThreadId, + staged.destinationCheckpointId, + ) + if (nonempty) + checkpointError( + 'SANDBOX_SNAPSHOT_FORK_DESTINATION_NOT_EMPTY', + 'Destination thread is not empty', + ) + const checkpoint: SandboxCheckpoint = cloneCheckpoint({ + id: staged.destinationCheckpointId, + threadId: staged.destinationThreadId, + parentCheckpointId: null, + createdAt: staged.createdAt, + reason: 'fork-root', + files: source.files, + conversation: source.conversation, + artifacts: source.artifacts, + }) + assertCheckpoint(checkpoint) + const conversationJson = JSON.stringify(checkpoint.conversation) + for (const key of checkpointKeys(source)) { + const reference = db + .prepare( + 'SELECT reference_count FROM sandbox_checkpoint_blob_references WHERE blob_key = ?', + ) + .get(key) as { reference_count: number } | undefined + if (!reference || reference.reference_count <= 0) + throw new SandboxCheckpointInvalidEntryError( + `Source checkpoint blob '${key}' has no positive reference count`, + ) + } + db.prepare( + 'INSERT INTO messages (thread_id, messages_json) VALUES (?, ?)', + ).run(checkpoint.threadId, conversationJson) + writeCheckpoint(checkpoint, conversationJson) + db.prepare( + 'INSERT INTO sandbox_checkpoint_heads (thread_id, checkpoint_id) VALUES (?, ?)', + ).run(checkpoint.threadId, checkpoint.id) + incrementReferences(checkpoint) + db.exec('COMMIT') + return { checkpoint: cloneCheckpoint(checkpoint) } + } catch (error) { + db.exec('ROLLBACK') + throw error + } + }, + } +} + /** * Every store this backend provides, spelled out. * @@ -1130,6 +1813,51 @@ export function sqlitePersistence( } } +/** Build the seven persistence stores and a durable SQLite checkpoint store. */ +export function sqliteSandboxSnapshots( + options: SqlitePersistenceOptions & SandboxCheckpointStoreOptions, +): { + persistence: SqliteAIPersistence + checkpoints: ForkCapableSandboxCheckpointStore + close: () => void +} { + const filename = normalizeSqliteUrl(options.url) + ensureParentDirectory(filename) + const db = new DatabaseSync(filename) + try { + if (options.migrate) { + db.exec(SCHEMA_SQL) + addMissingColumns(db) + } + const messages = createMessageStore(db) + const persistence = defineAIPersistence({ + stores: { + messages, + runs: createRunStore(db), + interrupts: createInterruptStore(db), + metadata: createMetadataStore(db), + generationRuns: createGenerationRunStore(db), + artifacts: createArtifactStore(db), + blobs: createBlobStore(db), + }, + }) + const checkpoints = createCheckpointStore(db, options) + let closed = false + return { + persistence, + checkpoints, + close() { + if (closed) return + db.close() + closed = true + }, + } + } catch (error) { + db.close() + throw error + } +} + // --------------------------------------------------------------------------- // URL / path helpers (kept identical to the packaged Node SQLite factory so the // `{ url, migrate }` call site is a drop-in). diff --git a/packages/ai-persistence/skills/ai-persistence/SKILL.md b/packages/ai-persistence/skills/ai-persistence/SKILL.md index 6c7eb1dfe4..356af45330 100644 --- a/packages/ai-persistence/skills/ai-persistence/SKILL.md +++ b/packages/ai-persistence/skills/ai-persistence/SKILL.md @@ -91,6 +91,13 @@ authorize the caller against `ArtifactRecord.threadId` before serving (404, not 403, so valid ids aren't confirmed), and `reconstructGeneration` MUST be given `authorize` on any multi-user route. Both take ids straight from the caller. +Portable sandbox snapshots use the same `messages`, `artifacts`, and `blobs` +stores. Their artifact reader checks the checkpoint thread, but it does not +authenticate a caller. Authorize the thread before any route reads a snapshot +artifact. The snapshot checkpoint store also needs atomic append and fork +operations. A SQLite adapter must write a checkpoint, its head, and blob +reference counts in one transaction. + ## Sub-skills | Need to... | Read | diff --git a/packages/ai-persistence/skills/ai-persistence/build-cloudflare-artifact-store/SKILL.md b/packages/ai-persistence/skills/ai-persistence/build-cloudflare-artifact-store/SKILL.md index ba7090892f..0255952498 100644 --- a/packages/ai-persistence/skills/ai-persistence/build-cloudflare-artifact-store/SKILL.md +++ b/packages/ai-persistence/skills/ai-persistence/build-cloudflare-artifact-store/SKILL.md @@ -1,6 +1,6 @@ --- name: ai-persistence/build-cloudflare-artifact-store -description: Use when a Cloudflare Worker needs durable byte storage for TanStack AI generated media (images, audio, video, transcripts) — writes a BlobStore backed by R2 and an ArtifactStore backed by D1 (or KV), composes them onto the generation persistence so withGenerationPersistence persists artifact bytes, and serves them back from a Worker GET route. Includes one-line sketches for S3, GCS, Vercel Blob, Supabase, and a dev filesystem BlobStore. +description: Use when a Cloudflare Worker needs durable byte storage for TanStack AI generated media (images, audio, video, transcripts) — writes a BlobStore backed by R2 and an ArtifactStore backed by D1, composes them onto the generation persistence so withGenerationPersistence persists artifact bytes, and serves them back from a Worker GET route. Includes one-line sketches for S3, GCS, Vercel Blob, Supabase, and a dev filesystem BlobStore. --- # Cloudflare Artifact + Blob Store @@ -21,7 +21,7 @@ per-request-binding rule, `wrangler` config shape, D1 migration workflow, and th chat (generation-run/message) side. This skill covers only the two byte-storage stores and how to compose them. -## The two contracts, verbatim +## The two contracts Both come from `@tanstack/ai-persistence`. `defineBlobStore` / `defineArtifactStore` type an object literal inline (autocomplete + contract checking, no separate @@ -30,28 +30,33 @@ annotation). ```ts // BlobStore — the byte layer. R2 backs it. interface BlobStore { - put( + put: ( key: string, body: BlobBody, options?: BlobPutOptions, - ): Promise + ) => Promise // metadata + byte accessors; `options.range` reads one slice (for `206`s) - get(key: string, options?: BlobGetOptions): Promise - head(key: string): Promise // metadata only - delete(key: string): Promise // no-op if absent - list(options?: BlobListOptions): Promise + get: (key: string, options?: BlobGetOptions) => Promise + head: (key: string) => Promise // metadata only + delete: (key: string) => Promise // no-op if absent + list: (options?: BlobListOptions) => Promise } -// ArtifactStore — the metadata layer. D1 (or KV) backs it. +// ArtifactStore — the metadata layer. D1 backs it. interface ArtifactStore { - save(record: ArtifactRecord): Promise // insert or overwrite - get(artifactId: string): Promise - list(runId: string): Promise> // [] when none - delete(artifactId: string): Promise - deleteForRun(runId: string): Promise + save: (record: ArtifactRecord) => Promise // insert or overwrite + get: (artifactId: string) => Promise + list: (runId: string) => Promise> // [] when none + listForThread: (threadId: string) => Promise> + delete: (artifactId: string) => Promise + deleteForRun: (runId: string) => Promise } ``` +`list` and `listForThread` return records ordered by `createdAt`, then by the +ordinal bytewise order of `artifactId`. Compare UTF-8 bytes from left to right. +Do not use locale collation. + `BlobBody` is `ReadableStream | ArrayBuffer | ArrayBufferView | string | Blob`. The non-stream shapes flow straight into `R2Bucket.put` unchanged — but a `ReadableStream` body does **not**, in the general case: @@ -362,9 +367,10 @@ Invariants that matter (asserted by the conformance testkit): ## 2. ArtifactStore backed by D1 -`ArtifactRecord` is `{ artifactId, runId, threadId, name, mimeType, size, -sourceUrl?, createdAt }` (`createdAt` epoch ms). One flat table, keyed by -`artifact_id`, indexed by `run_id` for `list`. +`ArtifactRecord` is `{ artifactId, runId, threadId, blobKey?, name, mimeType, size, +sourceUrl?, createdAt }` (`createdAt` epoch ms). One flat table is keyed by +`artifact_id`. It has run and thread ordered indexes for `list` and +`listForThread`. ```sql CREATE TABLE IF NOT EXISTS generation_artifacts ( @@ -378,7 +384,10 @@ CREATE TABLE IF NOT EXISTS generation_artifacts ( source_url text, created_at integer NOT NULL ); -CREATE INDEX IF NOT EXISTS generation_artifacts_run ON generation_artifacts (run_id); +CREATE INDEX IF NOT EXISTS generation_artifacts_run_order + ON generation_artifacts (run_id, created_at, artifact_id); +CREATE INDEX IF NOT EXISTS generation_artifacts_thread_order + ON generation_artifacts (thread_id, created_at, artifact_id); ``` ```ts ignore @@ -450,12 +459,26 @@ export function d1ArtifactStore(db: D1Database) { async list(runId) { const { results } = await db - .prepare(`SELECT * FROM generation_artifacts WHERE run_id = ?`) + .prepare( + `SELECT * FROM generation_artifacts WHERE run_id = ? ORDER BY created_at ASC, artifact_id ASC`, + ) .bind(runId) .all() return results.map(fromRow) }, + async listForThread(threadId) { + const { results } = await db + .prepare( + `SELECT * FROM generation_artifacts + WHERE thread_id = ? + ORDER BY created_at ASC, artifact_id ASC`, + ) + .bind(threadId) + .all() + return results.map(fromRow) + }, + async delete(artifactId) { await db .prepare(`DELETE FROM generation_artifacts WHERE artifact_id = ?`) @@ -478,11 +501,9 @@ keeps records comparing cleanly against the reference in-memory store. Persist `blob_key` verbatim: a `storageKey` mapper can put the bytes anywhere, so a reader cannot recompute the path — `resolveArtifactBlobKey(record)` falls back to the default convention only for rows written before the column existed. -**KV alternative:** if -you have no D1, back `save`/`get` with `KV.put(artifactId, JSON.stringify(record))` -/ `KV.get(artifactId, 'json')`, and maintain a `run:` index key (a JSON -array of artifact ids) for `list` — KV has no query, so `list` needs that -secondary index. +Cloudflare KV is not an equivalent ArtifactStore backend. The required ordering +and indexed reads need a transactional indexed database. Use D1 or another +transactional indexed database for artifact metadata. Store the bytes in R2. ## 3. Compose and wire @@ -682,7 +703,8 @@ between runs (see **ai-persistence/build-cloudflare-adapter** for the ascending keys, pages through the `cursor` when `truncated` without gaps or repeats, and returns an empty untruncated page for `limit: 0`. - The `ArtifactStore`: `save` is insert-or-overwrite, `get` returns `null` when - absent, `list(runId)` returns `[]` for an unknown run, and `delete` / + absent, `list(runId)` returns `[]` for an unknown run, + `listForThread(threadId)` returns the complete ordered history, and `delete` / `deleteForRun` remove exactly the expected rows. - The `GenerationRunStore`: `createOrResume` idempotency, no-op `update` on an unknown id, and `findLatestForThread` returning the most recently started diff --git a/packages/ai-persistence/src/capabilities.ts b/packages/ai-persistence/src/capabilities.ts index dbf1774367..d6aa892419 100644 --- a/packages/ai-persistence/src/capabilities.ts +++ b/packages/ai-persistence/src/capabilities.ts @@ -7,6 +7,11 @@ import { createCapability } from '@tanstack/ai' import type { AIPersistence, InterruptStore } from './types' +export interface PersistenceCompletion { + /** Resolves after successful terminal persistence, or rejects with the original run error or abort reason after terminal persistence settles. */ + waitForRunCompletion: () => Promise +} + export const PersistenceCapability = createCapability()('persistence') @@ -14,5 +19,10 @@ export const InterruptsCapability = createCapability()( 'persistence.interrupts', ) +export const PersistenceCompletionCapability = + createCapability()('persistence.completion') + export const [getPersistence, providePersistence] = PersistenceCapability export const [getInterrupts, provideInterrupts] = InterruptsCapability +export const [getPersistenceCompletion, providePersistenceCompletion] = + PersistenceCompletionCapability diff --git a/packages/ai-persistence/src/index.ts b/packages/ai-persistence/src/index.ts index 9519ec6235..4365c84347 100644 --- a/packages/ai-persistence/src/index.ts +++ b/packages/ai-persistence/src/index.ts @@ -53,6 +53,7 @@ export type { // Scope.threadId; authorize multi-user access with Scope.userId/tenantId. Scope, } from './types' +export type { PersistenceCompletion } from './capabilities' // AIPersistenceStores is intentionally NOT re-exported — use a named chat // shape or AIPersistence<{ messages: MessageStore, … }>. @@ -111,4 +112,7 @@ export { providePersistence, getInterrupts, provideInterrupts, + PersistenceCompletionCapability, + getPersistenceCompletion, + providePersistenceCompletion, } from './capabilities' diff --git a/packages/ai-persistence/src/memory.ts b/packages/ai-persistence/src/memory.ts index 922f5344a7..608fda24d0 100644 --- a/packages/ai-persistence/src/memory.ts +++ b/packages/ai-persistence/src/memory.ts @@ -22,6 +22,22 @@ import type { RunStore, } from './types' +const compareUtf8Bytes = (left: string, right: string): number => { + const leftBytes = new TextEncoder().encode(left) + const rightBytes = new TextEncoder().encode(right) + const length = Math.min(leftBytes.length, rightBytes.length) + + for (let index = 0; index < length; index++) { + const leftByte = leftBytes[index] + const rightByte = rightBytes[index] + if (leftByte !== rightByte) { + return (leftByte ?? 0) - (rightByte ?? 0) + } + } + + return leftBytes.length - rightBytes.length +} + class MemoryMessageStore implements MessageStore { private readonly threads = new Map>() loadThread(threadId: string): Promise> { @@ -267,7 +283,24 @@ class MemoryArtifactStore implements ArtifactStore { } list(runId: string): Promise> { return Promise.resolve( - [...this.artifacts.values()].filter((a) => a.runId === runId), + [...this.artifacts.values()] + .filter((a) => a.runId === runId) + .sort( + (a, b) => + a.createdAt - b.createdAt || + compareUtf8Bytes(a.artifactId, b.artifactId), + ), + ) + } + listForThread(threadId: string): Promise> { + return Promise.resolve( + [...this.artifacts.values()] + .filter((a) => a.threadId === threadId) + .sort( + (a, b) => + a.createdAt - b.createdAt || + compareUtf8Bytes(a.artifactId, b.artifactId), + ), ) } delete(artifactId: string): Promise { diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index 99df9d5411..ec1c983dc1 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -8,8 +8,10 @@ import { base64ToUint8Array } from '@tanstack/ai-utils' import { InterruptsCapability, PersistenceCapability, + PersistenceCompletionCapability, provideInterrupts, providePersistence, + providePersistenceCompletion, } from './capabilities' import { validateChatPersistenceStores, @@ -265,6 +267,11 @@ interface RunStateEntry { * bubble in place. */ streamingMessageId?: string + completion?: { + promise: Promise + resolve: () => void + reject: (error: unknown) => void + } } const runState = new WeakMap() @@ -1429,6 +1436,7 @@ export function withPersistence( const provides = [ PersistenceCapability, + PersistenceCompletionCapability, ...(wantsInterrupts ? [InterruptsCapability] : []), ] @@ -1438,9 +1446,27 @@ export function withPersistence( setup(ctx: ChatMiddlewareContext) { providePersistence(ctx, persistence) + let resolveCompletion: () => void = () => undefined + let rejectCompletion: (error: unknown) => void = () => undefined + const completion = new Promise((resolve, reject) => { + resolveCompletion = resolve + rejectCompletion = reject + }) + // Consumers may not need this capability. Mark the rejection handled while + // preserving the original promise for callers that do await it. + void completion.catch(() => undefined) + runState.set(ctx, { merged: false, interrupted: false, + completion: { + promise: completion, + resolve: resolveCompletion, + reject: rejectCompletion, + }, + }) + providePersistenceCompletion(ctx, { + waitForRunCompletion: () => completion, }) if (wantsInterrupts && persistence.stores.interrupts) { @@ -1618,16 +1644,28 @@ export function withPersistence( // resumes stay pending so a retry can re-apply them. Completing the run // or consuming approvals before the durable history lands leaves a // "finished" run whose transcript is missing the terminal turn. - await messageStore.saveThread( - ctx.threadId, - finishedTranscript(ctx.messages, info, state?.streamingMessageId), - ) - await completeRun(runs, ctx.runId, info.usage) - await commitPendingResumes(state, persistence.stores.interrupts) + try { + const transcript = finishedTranscript( + ctx.messages, + info, + state?.streamingMessageId, + ) + await messageStore.saveThread(ctx.threadId, transcript) + await completeRun(runs, ctx.runId, info.usage) + await commitPendingResumes(state, persistence.stores.interrupts) + state?.completion?.resolve() + } catch (error) { + state?.completion?.reject(error) + throw error + } }, async onError(ctx: ChatMiddlewareContext, info: ErrorInfo) { - await failRun(runs, ctx.runId, info.error) + try { + await failRun(runs, ctx.runId, info.error) + } finally { + runState.get(ctx)?.completion?.reject(info.error) + } }, async onAbort(ctx: ChatMiddlewareContext, info: AbortInfo) { @@ -1637,10 +1675,6 @@ export function withPersistence( // (`info.cancelRequested`, set when the cancel aborted this host's signal) // and durable (`RunRecord.cancelRequested`, the only channel that reaches // a run being driven elsewhere). - const cancelled = - info.cancelRequested === true || - (runs !== undefined && (await wasCancelRequested(runs, ctx.runId))) - // A run paused at an interrupt boundary is waiting for a HUMAN, not for // this socket. `chat()` skips its terminal hook at an actionable-wait // boundary, so its `finally` routes the disconnect here — and @@ -1649,9 +1683,21 @@ export function withPersistence( // still threw on the next request. An explicit cancel is different: the // user gave up on the approval, so the cancel band stays authoritative. const state = runState.get(ctx) - if (cancelled || (!detachableRun(ctx) && state?.interrupted !== true)) { - await abortRun(runs, ctx.runId) - return + let terminal = false + try { + // The durable cancel read is best-effort. It must not bypass the + // terminal persistence path or prevent the completion promise from + // settling when the run store is unavailable. + const cancelled = + info.cancelRequested === true || + (runs !== undefined && (await wasCancelRequested(runs, ctx.runId))) + terminal = + cancelled || (!detachableRun(ctx) && state?.interrupted !== true) + if (terminal) { + await abortRun(runs, ctx.runId) + } + } finally { + if (terminal) state?.completion?.reject(info.reason) } // A plain disconnect on a detachable or interrupted run: write NOTHING. // Either the agent is still running and a later attach can take it over diff --git a/packages/ai-persistence/src/testkit/conformance.ts b/packages/ai-persistence/src/testkit/conformance.ts index eaa941ebb5..df573723ed 100644 --- a/packages/ai-persistence/src/testkit/conformance.ts +++ b/packages/ai-persistence/src/testkit/conformance.ts @@ -922,24 +922,24 @@ export function runPersistenceConformance( await store.save( artifact({ - artifactId: 'art-a', + artifactId: '\u{10000}', blobKey: 'artifacts/run-art/art-a', createdAt: 100, }), ) await store.save( artifact({ - artifactId: 'art-b', + artifactId: '\u{e000}', sourceUrl: 'https://provider.example/expiring.png', - createdAt: 200, + createdAt: 100, }), ) await store.save( artifact({ artifactId: 'art-c', runId: 'run-art-other' }), ) - expect(await store.get('art-a')).toMatchObject({ - artifactId: 'art-a', + expect(await store.get('\u{10000}')).toMatchObject({ + artifactId: '\u{10000}', runId: 'run-art', threadId: 'thread-art', blobKey: 'artifacts/run-art/art-a', @@ -948,28 +948,93 @@ export function runPersistenceConformance( size: 3, createdAt: 100, }) - expect(await store.get('art-b')).toMatchObject({ + expect(await store.get('\u{e000}')).toMatchObject({ sourceUrl: 'https://provider.example/expiring.png', }) expect((await store.list('run-art')).map((r) => r.artifactId)).toEqual([ - 'art-a', - 'art-b', + '\u{e000}', + '\u{10000}', ]) // save() is insert-OR-OVERWRITE: re-saving an id corrects the record. await store.save( - artifact({ artifactId: 'art-a', name: 'renamed.png', size: 9 }), + artifact({ artifactId: '\u{10000}', name: 'renamed.png', size: 9 }), ) - const updated = await store.get('art-a') + const updated = await store.get('\u{10000}') expect(updated).toMatchObject({ name: 'renamed.png', size: 9 }) expect(updated?.blobKey).toBeUndefined() expect((await store.list('run-art')).map((r) => r.artifactId)).toEqual([ - 'art-a', - 'art-b', + '\u{e000}', + '\u{10000}', ]) }) + it('lists a thread in deterministic createdAt and artifactId order', async () => { + const store = resolveStore('artifacts') + if (!store) return + + await store.save( + artifact({ + artifactId: 'thread-b', + threadId: 'thread-order', + createdAt: 2, + }), + ) + await store.save( + artifact({ + artifactId: 'thread-a', + threadId: 'thread-order', + createdAt: 2, + }), + ) + await store.save( + artifact({ + artifactId: 'thread-early', + threadId: 'thread-order', + createdAt: 1, + }), + ) + await store.save( + artifact({ + artifactId: 'other', + threadId: 'thread-other', + createdAt: 0, + }), + ) + + expect( + (await store.listForThread('thread-order')).map((r) => r.artifactId), + ).toEqual(['thread-early', 'thread-a', 'thread-b']) + }) + + it('orders artifact IDs by UTF-8 bytes after createdAt', async () => { + const store = resolveStore('artifacts') + if (!store) return + + await store.save( + artifact({ + artifactId: '\u{10000}', + threadId: 'thread-utf8', + createdAt: 1, + }), + ) + await store.save( + artifact({ + artifactId: '\u{e000}', + threadId: 'thread-utf8', + createdAt: 1, + }), + ) + await store.save( + artifact({ artifactId: 'a', threadId: 'thread-utf8', createdAt: 1 }), + ) + + expect( + (await store.listForThread('thread-utf8')).map((r) => r.artifactId), + ).toEqual(['a', '\u{e000}', '\u{10000}']) + }) + it('deletes one artifact and every artifact for a run', async () => { const store = resolveStore('artifacts') if (!store) return diff --git a/packages/ai-persistence/src/types.ts b/packages/ai-persistence/src/types.ts index 7362141b57..c6c2068156 100644 --- a/packages/ai-persistence/src/types.ts +++ b/packages/ai-persistence/src/types.ts @@ -388,8 +388,20 @@ export interface ArtifactStore { save: (record: ArtifactRecord) => Promise /** Return the artifact for `artifactId`, or `null` if none exists. */ get: (artifactId: string) => Promise - /** All artifacts for a run. Returns `[]` when the run has none. */ + /** + * All artifacts for a run in deterministic snapshot order: `createdAt` + * ascending, then `artifactId` ascending by the unsigned UTF-8 bytes of + * each string (compare bytes left-to-right; shorter equal prefixes first). + * Returns `[]` when the run has none. + */ list: (runId: string) => Promise> + /** + * All artifacts for a thread in deterministic snapshot order. + * Records are ordered by `createdAt` ascending, then by `artifactId` using + * the unsigned UTF-8 bytes of each string (compare bytes left-to-right; shorter + * equal prefixes first). + */ + listForThread: (threadId: string) => Promise> /** * Delete a single artifact by id. A no-op if absent, mirroring * {@link BlobStore.delete} — the two are written and deleted as a pair, so diff --git a/packages/ai-persistence/tests/artifact-thread.test.ts b/packages/ai-persistence/tests/artifact-thread.test.ts new file mode 100644 index 0000000000..6ebd5af1f1 --- /dev/null +++ b/packages/ai-persistence/tests/artifact-thread.test.ts @@ -0,0 +1,52 @@ +import { describe, expect, it } from 'vitest' +import { memoryPersistence } from '../src/memory' +import type { ArtifactRecord } from '../src' + +const artifact = ( + overrides: Partial & Pick, +): ArtifactRecord => ({ + runId: 'run-1', + threadId: 'thread-1', + name: 'file.txt', + mimeType: 'text/plain', + size: 1, + createdAt: 1, + ...overrides, +}) + +describe('ArtifactStore.listForThread', () => { + it('lists a thread artifacts by createdAt then artifactId', async () => { + const artifacts = memoryPersistence().stores.artifacts + if (!artifacts) + throw new Error('memory persistence should provide artifacts') + + await artifacts.save(artifact({ artifactId: 'b', createdAt: 2 })) + await artifacts.save(artifact({ artifactId: 'a', createdAt: 2 })) + await artifacts.save(artifact({ artifactId: 'earlier', createdAt: 1 })) + await artifacts.save( + artifact({ + artifactId: 'other-thread', + threadId: 'thread-2', + createdAt: 0, + }), + ) + + expect( + (await artifacts.listForThread('thread-1')).map((x) => x.artifactId), + ).toEqual(['earlier', 'a', 'b']) + }) + + it('orders mixed ASCII, accented, and astral IDs by UTF-8 bytes', async () => { + const artifacts = memoryPersistence().stores.artifacts + if (!artifacts) + throw new Error('memory persistence should provide artifacts') + + await artifacts.save(artifact({ artifactId: '😀' })) + await artifacts.save(artifact({ artifactId: 'é' })) + await artifacts.save(artifact({ artifactId: 'a' })) + + expect( + (await artifacts.listForThread('thread-1')).map((x) => x.artifactId), + ).toEqual(['a', 'é', '😀']) + }) +}) diff --git a/packages/ai-persistence/tests/persistence-completion.test.ts b/packages/ai-persistence/tests/persistence-completion.test.ts new file mode 100644 index 0000000000..4b9d1eaa3d --- /dev/null +++ b/packages/ai-persistence/tests/persistence-completion.test.ts @@ -0,0 +1,661 @@ +import { describe, expect, it, vi } from 'vitest' +import { + DetachableRunCapability, + EventType, + chat, + defineChatMiddleware, + provideDetachableRun, +} from '@tanstack/ai' +import type { AnyTextAdapter, ModelMessage, StreamChunk } from '@tanstack/ai' +import { memoryPersistence } from '../src/memory' +import { withPersistence } from '../src/middleware' +import { + PersistenceCompletionCapability, + getPersistenceCompletion, +} from '../src/capabilities' + +function mockAdapter(chunks: Array) { + return { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': { + providerOptions: undefined, + inputModalities: undefined, + messageMetadataByModality: undefined, + toolCapabilities: undefined, + toolCallMetadata: undefined, + systemPromptMetadata: undefined, + }, + chatStream: () => + (async function* () { + for (const chunk of chunks) yield chunk + })(), + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + } satisfies AnyTextAdapter +} + +function waitingAdapter(signal: AbortSignal) { + return { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': { + providerOptions: undefined, + inputModalities: undefined, + messageMetadataByModality: undefined, + toolCapabilities: undefined, + toolCallMetadata: undefined, + systemPromptMetadata: undefined, + }, + chatStream: () => + (async function* () { + yield { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: 1, + } + await new Promise((resolve) => + signal.addEventListener('abort', () => resolve(), { once: true }), + ) + })(), + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + } satisfies AnyTextAdapter +} + +async function collect(stream: AsyncIterable) { + for await (const _chunk of stream) { + // Drain terminal middleware hooks. + } +} + +async function nextEventLoopTurn() { + await new Promise((resolve) => setTimeout(resolve, 0)) +} + +function createCompletionConsumer() { + let completion: ReturnType | undefined + let resolveSetup: (() => void) | undefined + const setupReady = new Promise((resolve) => { + resolveSetup = resolve + }) + const consumer = defineChatMiddleware({ + name: 'completion-consumer', + requires: [PersistenceCompletionCapability], + setup(ctx) { + completion = getPersistenceCompletion(ctx) + resolveSetup?.() + }, + }) + return { completion: () => completion, consumer, setupReady } +} + +function createCompletionRun( + persistence: ReturnType, + input: { + adapter: AnyTextAdapter + messages: Array + runId: string + threadId: string + resume?: Parameters[0]['resume'] + abortController?: AbortController + middleware?: Parameters[0]['middleware'] + }, +) { + const completionSetup = createCompletionConsumer() + const run = collect( + chat({ + ...input, + middleware: [ + withPersistence(persistence), + ...(input.middleware ?? []), + completionSetup.consumer, + ], + }), + ) + return { ...completionSetup, run } +} + +function requireCompletion( + completion: ReturnType | undefined, +) { + if (!completion) throw new Error('Completion middleware was not initialized') + return completion +} + +function requireValue(value: T | undefined) { + if (value === undefined) + throw new Error('Expected test value to be initialized') + return value +} + +const detachableProvider = defineChatMiddleware({ + name: 'detachable-provider', + provides: [DetachableRunCapability], + setup(ctx) { + provideDetachableRun(ctx, true) + }, +}) + +describe('PersistenceCompletionCapability', () => { + function finishedChunks(): Array { + return [ + { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: 1, + }, + { + type: EventType.TEXT_MESSAGE_START, + messageId: 'assistant-1', + role: 'assistant', + runId: 'run-1', + threadId: 'thread-1', + timestamp: 2, + }, + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'assistant-1', + delta: 'done', + runId: 'run-1', + threadId: 'thread-1', + timestamp: 3, + }, + { + type: EventType.TEXT_MESSAGE_END, + messageId: 'assistant-1', + runId: 'run-1', + threadId: 'thread-1', + timestamp: 4, + }, + { + type: EventType.RUN_FINISHED, + runId: 'run-1', + threadId: 'thread-1', + finishReason: 'stop', + timestamp: 5, + }, + ] + } + + it('resolves only after final transcript persistence succeeds', async () => { + const persistence = memoryPersistence() + let resolveSave: (() => void) | undefined + const saved = new Promise((resolve) => { + resolveSave = resolve + }) + let saveCount = 0 + const originalSave = persistence.stores.messages.saveThread.bind( + persistence.stores.messages, + ) + persistence.stores.messages.saveThread = async (...args) => { + saveCount += 1 + if (saveCount > 1) { + await saved + } + await originalSave(...args) + } + + const { completion, setupReady, run } = createCompletionRun(persistence, { + adapter: mockAdapter(finishedChunks()), + messages: [{ role: 'user', content: 'hi' }], + runId: 'run-1', + threadId: 'thread-1', + }) + + await setupReady + expect(completion).toBeDefined() + let settled = false + const waiting = completion() + ?.waitForRunCompletion() + .then(() => { + settled = true + }) + await Promise.resolve() + expect(settled).toBe(false) + + resolveSave?.() + await run + await waiting + expect(settled).toBe(true) + }) + + it('rejects with the original final persistence error', async () => { + const persistence = memoryPersistence() + const failure = new Error('final transcript failed') + persistence.stores.messages.saveThread = async () => { + throw failure + } + const { completion, setupReady, run } = createCompletionRun(persistence, { + adapter: mockAdapter([ + { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: 1, + }, + { + type: EventType.RUN_FINISHED, + runId: 'run-1', + threadId: 'thread-1', + finishReason: 'stop', + timestamp: 2, + }, + ]), + messages: [{ role: 'user', content: 'hi' }], + runId: 'run-1', + threadId: 'thread-1', + }) + await setupReady + await run.catch(() => undefined) + + await expect(completion()?.waitForRunCompletion()).rejects.toBe(failure) + }) + + it('rejects when the initial save succeeds but the final save fails', async () => { + const persistence = memoryPersistence() + const failure = new Error('final save failed') + let saves = 0 + persistence.stores.messages.saveThread = async (...args) => { + saves += 1 + if (saves === 2) throw failure + await memoryPersistence().stores.messages.saveThread(...args) + } + const { completion, setupReady, run } = createCompletionRun(persistence, { + adapter: mockAdapter(finishedChunks()), + messages: [{ role: 'user', content: 'hi' }], + runId: 'run-1', + threadId: 'thread-1', + }) + await setupReady + await run.catch(() => undefined) + await expect(completion()?.waitForRunCompletion()).rejects.toBe(failure) + }) + + it('does not clone or freeze arbitrary message metadata', async () => { + const persistence = memoryPersistence() + const metadata = new Map([['kind', 'typed']]) + const providerBytes = new Uint8Array([1, 2, 3]) + const saved: Array> = [] + const originalSave = persistence.stores.messages.saveThread.bind( + persistence.stores.messages, + ) + persistence.stores.messages.saveThread = async (threadId, messages) => { + saved.push(messages) + await originalSave(threadId, messages) + } + const { completion, run } = createCompletionRun(persistence, { + adapter: mockAdapter(finishedChunks()), + messages: [ + { + role: 'user', + content: [ + { + type: 'text', + content: 'hi', + metadata: { metadata, providerBytes }, + }, + ], + }, + ], + runId: 'run-1', + threadId: 'thread-1', + }) + await run + await requireCompletion(completion()).waitForRunCompletion() + const savedMessage = saved.at(-1)?.[0] + expect(savedMessage).toBeDefined() + expect(savedMessage).toMatchObject({ + content: [{ metadata: { metadata, providerBytes } }], + }) + expect(savedMessage?.content).toBeDefined() + expect(metadata.get('kind')).toBe('typed') + providerBytes[0] = 9 + const firstContent = savedMessage?.content?.[0] + if (typeof firstContent === 'string' || firstContent === undefined) + throw new Error('expected content part') + expect(firstContent.metadata).toMatchObject({ providerBytes }) + }) + + it('handles completion rejection when the caller does not await it', async () => { + const persistence = memoryPersistence() + const failure = new Error('unhandled completion') + persistence.stores.messages.saveThread = async () => { + throw failure + } + const unhandled = vi.fn() + process.on('unhandledRejection', unhandled) + try { + await collect( + chat({ + adapter: mockAdapter(finishedChunks()), + messages: [{ role: 'user', content: 'hi' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [withPersistence(persistence)], + }), + ).catch(() => undefined) + await nextEventLoopTurn() + expect(unhandled).not.toHaveBeenCalled() + } finally { + process.off('unhandledRejection', unhandled) + } + }) + + // These cases are kept as separate tests because each deferred terminal + // write is a distinct completion barrier. The detailed race setup lives in + // the lifecycle tests; keeping the names here makes the capability contract + // explicit at its public boundary. + it('keeps success pending while completeRun is pending', async () => { + const persistence = memoryPersistence() + const runs = requireValue(persistence.stores.runs) + let releaseUpdate: (() => void) | undefined + const updateBlocked = new Promise((resolve) => { + releaseUpdate = resolve + }) + const originalUpdate = runs.update.bind(runs) + runs.update = async (...args) => { + await updateBlocked + await originalUpdate(...args) + } + const { completion, setupReady, run } = createCompletionRun(persistence, { + adapter: mockAdapter(finishedChunks()), + messages: [{ role: 'user', content: 'hi' }], + runId: 'run-1', + threadId: 'thread-1', + }) + await setupReady + let settled = false + const waiting = requireCompletion(completion()) + .waitForRunCompletion() + .then(() => { + settled = true + }) + await Promise.resolve() + expect(settled).toBe(false) + requireValue(releaseUpdate)() + await run + await waiting + expect(settled).toBe(true) + }) + + it('keeps success pending while commitPendingResumes is pending', async () => { + const persistence = memoryPersistence() + await requireValue(persistence.stores.interrupts).create({ + interruptId: 'interrupt-1', + runId: 'run-1', + threadId: 'thread-1', + requestedAt: 1, + payload: { reason: 'approval_required' }, + }) + const interrupts = requireValue(persistence.stores.interrupts) + let releaseResolve: (() => void) | undefined + const resolveBlocked = new Promise((resolve) => { + releaseResolve = resolve + }) + const originalResolve = interrupts.resolve.bind(interrupts) + interrupts.resolve = async (...args) => { + await resolveBlocked + await originalResolve(...args) + } + const { completion, setupReady, run } = createCompletionRun(persistence, { + adapter: mockAdapter(finishedChunks()), + messages: [], + runId: 'run-1', + threadId: 'thread-1', + resume: [{ interruptId: 'interrupt-1', status: 'resolved', payload: {} }], + }) + await setupReady + let settled = false + const waiting = requireCompletion(completion()) + .waitForRunCompletion() + .then(() => { + settled = true + }) + await Promise.resolve() + expect(settled).toBe(false) + requireValue(releaseResolve)() + await run + await waiting + expect(settled).toBe(true) + }) + it('keeps error pending while failRun is pending, then rejects the same error', async () => { + const persistence = memoryPersistence() + const failure = new Error('provider failed') + const runs = requireValue(persistence.stores.runs) + let releaseUpdate: (() => void) | undefined + const updateBlocked = new Promise((resolve) => { + releaseUpdate = resolve + }) + const originalUpdate = runs.update.bind(runs) + runs.update = async (...args) => { + await updateBlocked + await originalUpdate(...args) + } + const adapter = mockAdapter([]) + adapter.chatStream = () => + (async function* () { + yield* [] + throw failure + })() + const { completion, setupReady, run } = createCompletionRun(persistence, { + adapter, + messages: [{ role: 'user', content: 'hi' }], + runId: 'run-1', + threadId: 'thread-1', + }) + await setupReady + const waiting = requireCompletion(completion()).waitForRunCompletion() + let settled = false + void waiting.catch(() => { + settled = true + }) + await Promise.resolve() + expect(settled).toBe(false) + requireValue(releaseUpdate)() + await expect(run).rejects.toBe(failure) + await expect(waiting).rejects.toBe(failure) + }) + + it('rejects the same error when the terminal error write fails', async () => { + const persistence = memoryPersistence() + const failure = new Error('provider failed') + const terminalFailure = new Error('failed to persist error status') + requireValue(persistence.stores.runs).update = async () => { + throw terminalFailure + } + const adapter = mockAdapter([]) + adapter.chatStream = () => + (async function* () { + yield* [] + throw failure + })() + const { completion, setupReady, run } = createCompletionRun(persistence, { + adapter, + messages: [{ role: 'user', content: 'hi' }], + runId: 'run-1', + threadId: 'thread-1', + }) + await setupReady + await expect(run).rejects.toBe(failure) + await expect( + requireCompletion(completion()).waitForRunCompletion(), + ).rejects.toBe(failure) + }) + it('keeps abort pending while the abort status write is pending, then rejects the same reason', async () => { + const persistence = memoryPersistence() + const runs = requireValue(persistence.stores.runs) + const reason = 'caller stopped' + const controller = new AbortController() + let releaseUpdate: (() => void) | undefined + const updateBlocked = new Promise((resolve) => { + releaseUpdate = resolve + }) + const originalUpdate = runs.update.bind(runs) + runs.update = async (...args) => { + await updateBlocked + await originalUpdate(...args) + } + const { + completion, + setupReady, + run: rawRun, + } = createCompletionRun(persistence, { + adapter: waitingAdapter(controller.signal), + messages: [{ role: 'user', content: 'hi' }], + runId: 'run-1', + threadId: 'thread-1', + abortController: controller, + }) + const run = rawRun.catch(() => undefined) + await setupReady + await Promise.resolve() + controller.abort(reason) + const waiting = requireCompletion(completion()).waitForRunCompletion() + let settled = false + void waiting.catch(() => { + settled = true + }) + await Promise.resolve() + expect(settled).toBe(false) + requireValue(releaseUpdate)() + await expect(waiting).rejects.toBe(reason) + await run + }) + it('rejects the same reason when the initial cancel-status lookup fails', async () => { + const persistence = memoryPersistence() + const lookupFailure = new Error('cancel lookup failed') + requireValue(persistence.stores.runs).get = async () => { + throw lookupFailure + } + const reason = 'caller stopped' + const controller = new AbortController() + const { completion, setupReady, run } = createCompletionRun(persistence, { + adapter: waitingAdapter(controller.signal), + messages: [], + runId: 'run-1', + threadId: 'thread-1', + abortController: controller, + }) + await setupReady + await Promise.resolve() + controller.abort(reason) + await expect( + requireCompletion(completion()).waitForRunCompletion(), + ).rejects.toBe(reason) + await run.catch(() => undefined) + }) + + it('keeps completion pending for a detachable disconnect', async () => { + const persistence = memoryPersistence() + const controller = new AbortController() + const { completion, setupReady, run } = createCompletionRun(persistence, { + adapter: waitingAdapter(controller.signal), + messages: [], + runId: 'run-1', + threadId: 'thread-1', + abortController: controller, + middleware: [detachableProvider], + }) + await setupReady + controller.abort('socket closed') + await run + + let settled = false + const waiting = requireCompletion(completion()) + .waitForRunCompletion() + .then(() => { + settled = true + }) + await nextEventLoopTurn() + expect(settled).toBe(false) + void waiting + }) + + it('rejects the same reason when the terminal abort write fails', async () => { + const persistence = memoryPersistence() + const terminalFailure = new Error('abort write failed') + requireValue(persistence.stores.runs).update = async () => { + throw terminalFailure + } + const reason = 'caller stopped' + const controller = new AbortController() + const { completion, setupReady, run } = createCompletionRun(persistence, { + adapter: waitingAdapter(controller.signal), + messages: [], + runId: 'run-1', + threadId: 'thread-1', + abortController: controller, + }) + await setupReady + await Promise.resolve() + controller.abort(reason) + await expect( + requireCompletion(completion()).waitForRunCompletion(), + ).rejects.toBe(reason) + await run.catch(() => undefined) + }) + it('does not produce an unhandled rejection on the error path', async () => { + const persistence = memoryPersistence() + const failure = new Error('provider failed') + const adapter = mockAdapter([]) + adapter.chatStream = () => + (async function* () { + yield* [] + throw failure + })() + const { completion, setupReady, run } = createCompletionRun(persistence, { + adapter, + messages: [{ role: 'user', content: 'hi' }], + runId: 'run-1', + threadId: 'thread-1', + }) + const unhandled: Array = [] + const onUnhandled = (reason: unknown): void => { + unhandled.push(reason) + } + process.on('unhandledRejection', onUnhandled) + try { + await setupReady + expect(completion).toBeDefined() + await expect(run).rejects.toBe(failure) + void requireCompletion(completion()).waitForRunCompletion() + await nextEventLoopTurn() + expect(unhandled).toEqual([]) + } finally { + process.off('unhandledRejection', onUnhandled) + } + }) + + it('does not produce an unhandled rejection on the abort path', async () => { + const persistence = memoryPersistence() + const reason = 'caller stopped' + const controller = new AbortController() + const { completion, setupReady, run } = createCompletionRun(persistence, { + adapter: waitingAdapter(controller.signal), + messages: [{ role: 'user', content: 'hi' }], + runId: 'run-1', + threadId: 'thread-1', + abortController: controller, + }) + const unhandled: Array = [] + const onUnhandled = (unhandledReason: unknown): void => { + unhandled.push(unhandledReason) + } + process.on('unhandledRejection', onUnhandled) + try { + await setupReady + controller.abort(reason) + await Promise.resolve() + expect(completion).toBeDefined() + await run + void requireCompletion(completion()).waitForRunCompletion() + await nextEventLoopTurn() + expect(unhandled).toEqual([]) + } finally { + process.off('unhandledRejection', onUnhandled) + } + }) +}) diff --git a/packages/ai-sandbox-cloudflare/src/handle.ts b/packages/ai-sandbox-cloudflare/src/handle.ts index 14880222cf..9693fc4fcc 100644 --- a/packages/ai-sandbox-cloudflare/src/handle.ts +++ b/packages/ai-sandbox-cloudflare/src/handle.ts @@ -29,6 +29,7 @@ import type { SandboxChannel, SandboxHandle, SpawnHandle, + SandboxFsStat, } from '@tanstack/ai-sandbox' export const CLOUDFLARE_CAPS: SandboxCapabilities = { @@ -55,6 +56,32 @@ function q(value: string): string { return `'${value.replace(/'/g, `'\\''`)}'` } +/** Verify a missing path by listing parent entries. `test -e` also fails for inaccessible parents. */ +function lstatCommand(path: string): string { + return `tanstack_lstat_path=${q(path)}; tanstack_lstat_output=$(stat -c '%f:%s' -- "$tanstack_lstat_path" 2>&1); tanstack_lstat_status=$?; if [ "$tanstack_lstat_status" -eq 0 ]; then printf '%s\n' "$tanstack_lstat_output"; else tanstack_lstat_missing() { tanstack_missing_path=$1; case "$tanstack_missing_path" in /|.) return 1 ;; */*) tanstack_parent=${'$'}{tanstack_missing_path%/*}; tanstack_name=${'$'}{tanstack_missing_path##*/}; [ -n "$tanstack_parent" ] || tanstack_parent=/ ;; *) tanstack_parent=.; tanstack_name=$tanstack_missing_path ;; esac; tanstack_parent_mode=$(stat -L -c '%f' -- "$tanstack_parent" 2>/dev/null); tanstack_parent_status=$?; if [ "$tanstack_parent_status" -ne 0 ]; then tanstack_lstat_missing "$tanstack_parent"; else case "$tanstack_parent_mode" in 4[0-9a-fA-F][0-9a-fA-F][0-9a-fA-F]) case "$tanstack_parent" in /*) tanstack_find_parent=$tanstack_parent ;; *) tanstack_find_parent=./$tanstack_parent ;; esac; tanstack_match=$(find -H "$tanstack_find_parent" -mindepth 1 -maxdepth 1 -exec sh -c 'tanstack_target=$1; shift; for tanstack_candidate do [ "${'$'}{tanstack_candidate##*/}" = "$tanstack_target" ] && { printf 1; exit 0; }; done; exit 0' sh "$tanstack_name" '{}' + 2>/dev/null); tanstack_find_status=$?; [ "$tanstack_find_status" -eq 0 ] && [ -z "$tanstack_match" ] ;; *) return 1 ;; esac; fi; }; if tanstack_lstat_missing "$tanstack_lstat_path"; then printf '%s' '__TANSTACK_LSTAT_MISSING__'; else printf '%s\n' "$tanstack_lstat_output" >&2; exit "$tanstack_lstat_status"; fi; fi` +} + +function parseLstatOutput(output: string): SandboxFsStat { + const fields = /^(?[0-9a-fA-F]{4}):(?\d+)\n?$/.exec(output) + const mode = fields?.groups?.mode + const size = fields?.groups?.size + if (!mode || !size) throw new Error(`invalid lstat output: ${output}`) + const parsedMode = Number.parseInt(mode, 16) + const parsedSize = Number(size) + if ( + !Number.isSafeInteger(parsedMode) || + !Number.isSafeInteger(parsedSize) || + parsedSize < 0 + ) + throw new Error(`invalid lstat output: ${output}`) + const type = parsedMode & 0xf000 + if (type === 0x8000) + return { type: 'file', mode: parsedMode, size: parsedSize } + if (type === 0x4000) return { type: 'dir', mode: parsedMode } + if (type === 0xa000) return { type: 'symlink', mode: parsedMode } + return { type: 'other', mode: parsedMode } +} + /** A push-driven async string queue used to adapt CF's onOutput callback. */ class OutputQueue { private readonly buffer: Array = [] @@ -162,6 +189,7 @@ export class CloudflareHandle implements SandboxHandle { } }) }, + lstat: async (p) => this.lstat(this.abs(p)), mkdir: async (p) => { await this.exec(`mkdir -p ${q(this.abs(p))}`) }, @@ -197,6 +225,16 @@ export class CloudflareHandle implements SandboxHandle { return p } + private async lstat(path: string): Promise { + const r = await this.exec(lstatCommand(path)) + if (r.exitCode === 0 && r.stdout === '__TANSTACK_LSTAT_MISSING__') + return undefined + if (r.exitCode !== 0) { + throw new Error(`lstat failed: ${r.stderr.trim()}`) + } + return parseLstatOutput(r.stdout) + } + private async exec( command: string, opts?: ProcessOptions, diff --git a/packages/ai-sandbox-cloudflare/tests/handle.test.ts b/packages/ai-sandbox-cloudflare/tests/handle.test.ts index 17d28f68d7..4c1c37e4e9 100644 --- a/packages/ai-sandbox-cloudflare/tests/handle.test.ts +++ b/packages/ai-sandbox-cloudflare/tests/handle.test.ts @@ -9,12 +9,55 @@ import { CLOUDFLARE_CAPS, CloudflareHandle } from '../src/handle' import type { Sandbox } from '@cloudflare/sandbox' import type { ExecResult } from '@tanstack/ai-sandbox' +function lstatCommand(path: string): string { + const quoted = `'${path.replace(/'/g, `'\\''`)}'` + return `tanstack_lstat_path=${quoted}; tanstack_lstat_output=$(stat -c '%f:%s' -- "$tanstack_lstat_path" 2>&1); tanstack_lstat_status=$?; if [ "$tanstack_lstat_status" -eq 0 ]; then printf '%s\n' "$tanstack_lstat_output"; else tanstack_lstat_missing() { tanstack_missing_path=$1; case "$tanstack_missing_path" in /|.) return 1 ;; */*) tanstack_parent=${'$'}{tanstack_missing_path%/*}; tanstack_name=${'$'}{tanstack_missing_path##*/}; [ -n "$tanstack_parent" ] || tanstack_parent=/ ;; *) tanstack_parent=.; tanstack_name=$tanstack_missing_path ;; esac; tanstack_parent_mode=$(stat -L -c '%f' -- "$tanstack_parent" 2>/dev/null); tanstack_parent_status=$?; if [ "$tanstack_parent_status" -ne 0 ]; then tanstack_lstat_missing "$tanstack_parent"; else case "$tanstack_parent_mode" in 4[0-9a-fA-F][0-9a-fA-F][0-9a-fA-F]) case "$tanstack_parent" in /*) tanstack_find_parent=$tanstack_parent ;; *) tanstack_find_parent=./$tanstack_parent ;; esac; tanstack_match=$(find -H "$tanstack_find_parent" -mindepth 1 -maxdepth 1 -exec sh -c 'tanstack_target=$1; shift; for tanstack_candidate do [ "${'$'}{tanstack_candidate##*/}" = "$tanstack_target" ] && { printf 1; exit 0; }; done; exit 0' sh "$tanstack_name" '{}' + 2>/dev/null); tanstack_find_status=$?; [ "$tanstack_find_status" -eq 0 ] && [ -z "$tanstack_match" ] ;; *) return 1 ;; esac; fi; }; if tanstack_lstat_missing "$tanstack_lstat_path"; then printf '%s' '__TANSTACK_LSTAT_MISSING__'; else printf '%s\n' "$tanstack_lstat_output" >&2; exit "$tanstack_lstat_status"; fi; fi` +} +function lstatPath(command: string): string | undefined { + return /^tanstack_lstat_path='([^']*)';/.exec(command)?.[1] +} + /** Options the handle passes to `sandbox.exec` (streaming spawn + one-shot). */ interface MockExecOpts { stream?: boolean onOutput?: (stream: 'stdout' | 'stderr', data: string) => void } +interface SandboxFixtureMethods { + exec: (command: string, opts?: MockExecOpts) => Promise + setEnvVars: () => Promise + exposePort: (port: number) => Promise<{ url: string }> + destroy: () => Promise +} + +class SandboxFixturePrototype implements SandboxFixtureMethods { + exec(): Promise { + return Promise.reject(new Error('sandbox fixture exec is not configured')) + } + + setEnvVars(): Promise { + return Promise.resolve() + } + + exposePort(port: number): Promise<{ url: string }> { + return Promise.resolve({ url: `https://${port}.example.workers.dev` }) + } + + destroy(): Promise { + return Promise.resolve() + } +} + +function sandboxFixture(methods: SandboxFixtureMethods): Sandbox { + // The runtime class imports a `cloudflare:` module that Node cannot load. + // This concrete fixture supplies every method CloudflareHandle uses. + const sandbox: Sandbox = Object.assign( + Object.create(SandboxFixturePrototype.prototype), + methods, + ) + return sandbox +} + /** A minimal in-memory Sandbox stub: fs lives in a Map; exec emulates the * base64/test/mkdir commands the handle issues, plus the streaming path * `spawn()` relies on (`exec({ stream: true, onOutput })`). */ @@ -64,24 +107,165 @@ function mockSandbox(): { sandbox: Sandbox; files: Map } { if (exists) { return Promise.resolve(files.has(exists[1]!) ? ok() : fail('')) } + const statPath = lstatPath(command) + if (statPath) { + expect(command).toBe(lstatCommand(statPath)) + const values = new Map([ + ['/workspace/file', '81A4:12\n'], + ['/workspace/empty', '81a4:0\n'], + ['/workspace/dir', '41ed:4096\n'], + ['/workspace/link', 'a1ff:4\n'], + ['/workspace/other', 'c1b6:0\n'], + ['/workspace/char', '21b6:0\n'], + ['/workspace/block', '61b6:0\n'], + ['/workspace/fifo', '11b6:0\n'], + ['/workspace/unknown', '71b6:7\n'], + ]) + if ( + statPath === '/workspace/missing' || + statPath === '/workspace/missing-parent/child' || + statPath === '-H/missing' || + statPath === '-delete/missing' + ) + return Promise.resolve(ok('__TANSTACK_LSTAT_MISSING__')) + if ( + [ + '/workspace/file/child', + '/workspace/loop/child', + '/workspace/denied-link/child', + ].includes(statPath) + ) + return Promise.resolve(fail('stat: permission denied')) + return Promise.resolve(ok(values.get(statPath) ?? '')) + } if (command.startsWith('mkdir -p')) return Promise.resolve(ok()) if (command.startsWith('echo ')) return Promise.resolve(ok(command.slice(5))) return Promise.resolve(ok()) } - const sandbox = { + const sandbox = sandboxFixture({ exec, setEnvVars: () => Promise.resolve(), exposePort: (port: number) => Promise.resolve({ url: `https://${port}.example.workers.dev` }), destroy: () => Promise.resolve(), - } as unknown as Sandbox + }) return { sandbox, files } } +function createLstatHandle(): CloudflareHandle { + return new CloudflareHandle('sbx-1', mockSandbox().sandbox, '/workspace') +} + describe('CloudflareHandle', () => { + it('parses lstat command output with size only for files', async () => { + const handle = createLstatHandle() + expect(await handle.fs.lstat!('/workspace/file')).toEqual({ + type: 'file', + mode: 33188, + size: 12, + }) + expect(await handle.fs.lstat!('/workspace/dir')).toEqual({ + type: 'dir', + mode: 16877, + }) + expect(await handle.fs.lstat!('/workspace/link')).toEqual({ + type: 'symlink', + mode: 41471, + }) + expect(await handle.fs.lstat!('/workspace/other')).toEqual({ + type: 'other', + mode: 49590, + }) + }) + it.each([ + '/workspace/missing', + '/workspace/missing-parent/child', + '-H/missing', + '-delete/missing', + ])('returns undefined for a verified missing path: %s', async (path) => { + const handle = createLstatHandle() + await expect(handle.fs.lstat!(path)).resolves.toBeUndefined() + }) + it.each([ + '/workspace/file/child', + '/workspace/loop/child', + '/workspace/denied-link/child', + ])('preserves an unverified parent failure: %s', async (path) => { + const handle = createLstatHandle() + await expect(handle.fs.lstat!(path)).rejects.toThrow( + 'lstat failed: stat: permission denied', + ) + }) + it('parses a character special file as other without size', async () => { + const handle = createLstatHandle() + await expect(handle.fs.lstat!('/workspace/char')).resolves.toEqual({ + type: 'other', + mode: 0x21b6, + }) + }) + it('parses a zero-byte regular empty file with size zero', async () => { + const handle = createLstatHandle() + await expect(handle.fs.lstat!('/workspace/empty')).resolves.toEqual({ + type: 'file', + mode: 0x81a4, + size: 0, + }) + }) + it.each([ + ['not-a-number', '81a4'], + ['Infinity', '81a4'], + ['-1', '81a4'], + ['', '81a4'], + ['5', '81a4junk'], + ['5', '-81a4'], + [' 5', '81a4'], + ['5 ', '81a4'], + ['5\n6', '81a4'], + ['NaN', '81a4'], + ['5', ''], + ['5', '0x81a4'], + ['5', '81a'], + ['5', '81a45'], + ['5', '81a4 '], + ['9007199254740992', '81a4'], + ])('rejects malformed lstat fields', async (size, mode) => { + const { sandbox } = mockSandbox() + const handle = new CloudflareHandle('sbx-1', sandbox, '/workspace') + Object.defineProperty(sandbox, 'exec', { + value: async () => ({ + stdout: `${mode}:${size}\n`, + stderr: '', + exitCode: 0, + }), + }) + await expect(handle.fs.lstat!('/workspace/file')).rejects.toThrow( + 'invalid lstat output', + ) + }) + it('parses a block special file as other without size', async () => { + const handle = createLstatHandle() + await expect(handle.fs.lstat!('/workspace/block')).resolves.toEqual({ + type: 'other', + mode: 0x61b6, + }) + }) + it('parses a fifo as other without size', async () => { + const handle = createLstatHandle() + await expect(handle.fs.lstat!('/workspace/fifo')).resolves.toEqual({ + type: 'other', + mode: 0x11b6, + }) + }) + it('parses an unknown file type as other without size', async () => { + const handle = createLstatHandle() + await expect(handle.fs.lstat!('/workspace/unknown')).resolves.toEqual({ + type: 'other', + mode: 0x71b6, + }) + }) it('advertises edge capabilities (ephemeral disk, no snapshots/fork)', () => { expect(CLOUDFLARE_CAPS.snapshots).toBe(false) expect(CLOUDFLARE_CAPS.durableFilesystem).toBe(false) @@ -102,15 +286,17 @@ describe('CloudflareHandle', () => { const gate = new Promise((resolve) => { release = resolve }) - const execOpts: Array> = [] - const sandbox = { - exec: (_command: string, opts: Record) => { - execOpts.push(opts) + const execOpts: Array = [] + const sandbox = sandboxFixture({ + exec: (_command: string, opts?: MockExecOpts) => { + execOpts.push(opts ?? {}) return gate.then(() => ({ stdout: '', stderr: '', exitCode: 0 })) }, setEnvVars: () => Promise.resolve(), + exposePort: (port: number) => + Promise.resolve({ url: `https://${port}.example.workers.dev` }), destroy: () => Promise.resolve(), - } as unknown as Sandbox + }) const handle = new CloudflareHandle('sbx-1', sandbox, '/workspace') const proc = await handle.process.spawn('tail -c +1 -f /tmp/journal') @@ -140,24 +326,21 @@ describe('CloudflareHandle', () => { }) it('round-trips files over base64 exec', async () => { - const { sandbox } = mockSandbox() - const handle = new CloudflareHandle('sbx-1', sandbox, '/workspace') + const handle = createLstatHandle() await handle.fs.write('/workspace/a.txt', 'hello edge') expect(await handle.fs.exists('/workspace/a.txt')).toBe(true) expect(await handle.fs.read('/workspace/a.txt')).toBe('hello edge') }) it('exec passes stdout/exit through', async () => { - const { sandbox } = mockSandbox() - const handle = new CloudflareHandle('sbx-1', sandbox, '/workspace') + const handle = createLstatHandle() const r = await handle.process.exec('echo hi') expect(r.stdout).toContain('hi') expect(r.exitCode).toBe(0) }) it('spawn streams output via the queue and resolves wait()', async () => { - const { sandbox } = mockSandbox() - const handle = new CloudflareHandle('sbx-1', sandbox, '/workspace') + const handle = createLstatHandle() const proc = await handle.process.spawn('run something') let out = '' for await (const chunk of proc.stdout) out += chunk @@ -166,8 +349,7 @@ describe('CloudflareHandle', () => { }) it('surfaces a command failure by rejecting wait()', async () => { - const { sandbox } = mockSandbox() - const handle = new CloudflareHandle('sbx-1', sandbox, '/workspace') + const handle = createLstatHandle() const proc = await handle.process.spawn('reject-me') // The stdout reader must still terminate even though the command failed... for await (const _chunk of proc.stdout) void _chunk @@ -177,8 +359,7 @@ describe('CloudflareHandle', () => { }) it('rejects stdin writes (documented CF limitation)', async () => { - const { sandbox } = mockSandbox() - const handle = new CloudflareHandle('sbx-1', sandbox, '/workspace') + const handle = createLstatHandle() const proc = await handle.process.spawn('run something') await expect(proc.stdin.write('x')).rejects.toThrow(/do not expose stdin/i) }) @@ -196,8 +377,7 @@ describe('CloudflareHandle', () => { }) it('ports.connect throws without a previewHostname', async () => { - const { sandbox } = mockSandbox() - const handle = new CloudflareHandle('sbx-1', sandbox, '/workspace') + const handle = createLstatHandle() await expect(handle.ports.connect(3000)).rejects.toThrow(/previewHostname/i) }) }) diff --git a/packages/ai-sandbox-daytona/src/handle.ts b/packages/ai-sandbox-daytona/src/handle.ts index 8e2d476cf0..8cb03b8ab2 100644 --- a/packages/ai-sandbox-daytona/src/handle.ts +++ b/packages/ai-sandbox-daytona/src/handle.ts @@ -25,6 +25,7 @@ import type { SandboxHandle, SnapshotRef, SpawnHandle, + SandboxFsStat, } from '@tanstack/ai-sandbox' export const DAYTONA_CAPS: SandboxCapabilities = { @@ -72,6 +73,32 @@ function q(value: string): string { return `'${value.replace(/'/g, `'\\''`)}'` } +/** Verify a missing path by listing parent entries. `test -e` also fails for inaccessible parents. */ +function lstatCommand(path: string): string { + return `tanstack_lstat_path=${q(path)}; tanstack_lstat_output=$(stat -c '%f:%s' -- "$tanstack_lstat_path" 2>&1); tanstack_lstat_status=$?; if [ "$tanstack_lstat_status" -eq 0 ]; then printf '%s\n' "$tanstack_lstat_output"; else tanstack_lstat_missing() { tanstack_missing_path=$1; case "$tanstack_missing_path" in /|.) return 1 ;; */*) tanstack_parent=${'$'}{tanstack_missing_path%/*}; tanstack_name=${'$'}{tanstack_missing_path##*/}; [ -n "$tanstack_parent" ] || tanstack_parent=/ ;; *) tanstack_parent=.; tanstack_name=$tanstack_missing_path ;; esac; tanstack_parent_mode=$(stat -L -c '%f' -- "$tanstack_parent" 2>/dev/null); tanstack_parent_status=$?; if [ "$tanstack_parent_status" -ne 0 ]; then tanstack_lstat_missing "$tanstack_parent"; else case "$tanstack_parent_mode" in 4[0-9a-fA-F][0-9a-fA-F][0-9a-fA-F]) case "$tanstack_parent" in /*) tanstack_find_parent=$tanstack_parent ;; *) tanstack_find_parent=./$tanstack_parent ;; esac; tanstack_match=$(find -H "$tanstack_find_parent" -mindepth 1 -maxdepth 1 -exec sh -c 'tanstack_target=$1; shift; for tanstack_candidate do [ "${'$'}{tanstack_candidate##*/}" = "$tanstack_target" ] && { printf 1; exit 0; }; done; exit 0' sh "$tanstack_name" '{}' + 2>/dev/null); tanstack_find_status=$?; [ "$tanstack_find_status" -eq 0 ] && [ -z "$tanstack_match" ] ;; *) return 1 ;; esac; fi; }; if tanstack_lstat_missing "$tanstack_lstat_path"; then printf '%s' '__TANSTACK_LSTAT_MISSING__'; else printf '%s\n' "$tanstack_lstat_output" >&2; exit "$tanstack_lstat_status"; fi; fi` +} + +function parseLstatOutput(output: string): SandboxFsStat { + const fields = /^(?[0-9a-fA-F]{4}):(?\d+)\n?$/.exec(output) + const mode = fields?.groups?.mode + const size = fields?.groups?.size + if (!mode || !size) throw new Error(`invalid lstat output: ${output}`) + const parsedMode = Number.parseInt(mode, 16) + const parsedSize = Number(size) + if ( + !Number.isSafeInteger(parsedMode) || + !Number.isSafeInteger(parsedSize) || + parsedSize < 0 + ) + throw new Error(`invalid lstat output: ${output}`) + const type = parsedMode & 0xf000 + if (type === 0x8000) + return { type: 'file', mode: parsedMode, size: parsedSize } + if (type === 0x4000) return { type: 'dir', mode: parsedMode } + if (type === 0xa000) return { type: 'symlink', mode: parsedMode } + return { type: 'other', mode: parsedMode } +} + function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)) } @@ -177,6 +204,7 @@ export class DaytonaHandle implements SandboxHandle { type: entry.isDir ? ('dir' as const) : ('file' as const), })) }, + lstat: async (p) => this.lstat(this.abs(p)), mkdir: async (p) => { await this.sandbox.fs.createFolder(this.abs(p), '755') }, @@ -223,6 +251,17 @@ export class DaytonaHandle implements SandboxHandle { return { ...this.envVars, ...extra } } + private async lstat(path: string): Promise { + const r = await this.exec(lstatCommand(path)) + if (r.exitCode === 0 && r.stdout === '__TANSTACK_LSTAT_MISSING__') + return undefined + if (r.exitCode !== 0) { + const output = `${r.stdout}\n${r.stderr}` + throw new Error(`lstat failed: ${output.trim()}`) + } + return parseLstatOutput(r.stdout) + } + /** * Session execute has no env field. Write values to a workdir file, * then source that file so the stored command never contains secrets. diff --git a/packages/ai-sandbox-daytona/tests/lstat.test.ts b/packages/ai-sandbox-daytona/tests/lstat.test.ts new file mode 100644 index 0000000000..da6dfbfc16 --- /dev/null +++ b/packages/ai-sandbox-daytona/tests/lstat.test.ts @@ -0,0 +1,199 @@ +import { describe, expect, it } from 'vitest' +import { Sandbox } from '@daytona/sdk' +import { DaytonaHandle } from '../src/handle' + +type ExecResult = { exitCode: number; stdout: string; stderr: string } +type ExecFn = (command: string) => Promise +function lstatCommand(path: string): string { + const quoted = `'${path.replace(/'/g, `'\\''`)}'` + return `tanstack_lstat_path=${quoted}; tanstack_lstat_output=$(stat -c '%f:%s' -- "$tanstack_lstat_path" 2>&1); tanstack_lstat_status=$?; if [ "$tanstack_lstat_status" -eq 0 ]; then printf '%s\n' "$tanstack_lstat_output"; else tanstack_lstat_missing() { tanstack_missing_path=$1; case "$tanstack_missing_path" in /|.) return 1 ;; */*) tanstack_parent=${'$'}{tanstack_missing_path%/*}; tanstack_name=${'$'}{tanstack_missing_path##*/}; [ -n "$tanstack_parent" ] || tanstack_parent=/ ;; *) tanstack_parent=.; tanstack_name=$tanstack_missing_path ;; esac; tanstack_parent_mode=$(stat -L -c '%f' -- "$tanstack_parent" 2>/dev/null); tanstack_parent_status=$?; if [ "$tanstack_parent_status" -ne 0 ]; then tanstack_lstat_missing "$tanstack_parent"; else case "$tanstack_parent_mode" in 4[0-9a-fA-F][0-9a-fA-F][0-9a-fA-F]) case "$tanstack_parent" in /*) tanstack_find_parent=$tanstack_parent ;; *) tanstack_find_parent=./$tanstack_parent ;; esac; tanstack_match=$(find -H "$tanstack_find_parent" -mindepth 1 -maxdepth 1 -exec sh -c 'tanstack_target=$1; shift; for tanstack_candidate do [ "${'$'}{tanstack_candidate##*/}" = "$tanstack_target" ] && { printf 1; exit 0; }; done; exit 0' sh "$tanstack_name" '{}' + 2>/dev/null); tanstack_find_status=$?; [ "$tanstack_find_status" -eq 0 ] && [ -z "$tanstack_match" ] ;; *) return 1 ;; esac; fi; }; if tanstack_lstat_missing "$tanstack_lstat_path"; then printf '%s' '__TANSTACK_LSTAT_MISSING__'; else printf '%s\n' "$tanstack_lstat_output" >&2; exit "$tanstack_lstat_status"; fi; fi` +} +function lstatPath(command: string): string { + return /^tanstack_lstat_path='([^']*)';/.exec(command)?.[1] ?? '' +} +function execSlot(handle: object) { + return { + set exec(value: ExecFn) { + Object.defineProperty(handle, 'exec', { configurable: true, value }) + }, + } +} + +function sandboxFixture(): Sandbox { + const sandbox: Sandbox = Object.assign(Object.create(Sandbox.prototype), { + id: 'test', + }) + return sandbox +} + +function createHandle(): DaytonaHandle { + return new DaytonaHandle({ + sandbox: sandboxFixture(), + workdir: '/workspace', + }) +} + +describe('DaytonaHandle.fs.lstat', () => { + it('parses file, directory, symlink, and other metadata', async () => { + const values = new Map([ + ['file', '81A4:12\n'], + ['dir', '41ed:4096\n'], + ['link', 'a1ff:4\n'], + ['other', 'c1b6:0\n'], + ['char', '21b6:0\n'], + ['block', '61b6:0\n'], + ['fifo', '11b6:0\n'], + ['unknown', '71b6:7\n'], + ]) + const handle = createHandle() + execSlot(handle).exec = async (command: string) => { + const path = lstatPath(command) + expect(command).toBe(lstatCommand(path)) + return { + exitCode: 0, + stdout: values.get(path.split('/').pop()!) ?? '', + stderr: '', + } + } + await expect(handle.fs.lstat!('/workspace/file')).resolves.toEqual({ + type: 'file', + mode: 0x81a4, + size: 12, + }) + await expect(handle.fs.lstat!('/workspace/dir')).resolves.toEqual({ + type: 'dir', + mode: 0x41ed, + }) + await expect(handle.fs.lstat!('/workspace/link')).resolves.toEqual({ + type: 'symlink', + mode: 0xa1ff, + }) + await expect(handle.fs.lstat!('/workspace/other')).resolves.toEqual({ + type: 'other', + mode: 0xc1b6, + }) + }) + + it('parses a character special file as other without size', async () => { + const handle = createHandle() + execSlot(handle).exec = async () => ({ + exitCode: 0, + stdout: '21b6:0\n', + stderr: '', + }) + await expect(handle.fs.lstat!('/workspace/char')).resolves.toEqual({ + type: 'other', + mode: 0x21b6, + }) + }) + it('parses a zero-byte regular empty file with size zero', async () => { + const handle = createHandle() + execSlot(handle).exec = async () => ({ + exitCode: 0, + stdout: '81a4:0\n', + stderr: '', + }) + await expect(handle.fs.lstat!('/workspace/empty')).resolves.toEqual({ + type: 'file', + mode: 0x81a4, + size: 0, + }) + }) + it.each([ + ['not-a-number', '81a4'], + ['Infinity', '81a4'], + ['-1', '81a4'], + ['', '81a4'], + ['5', '81a4junk'], + ['5', '-81a4'], + [' 5', '81a4'], + ['5 ', '81a4'], + ['5\n6', '81a4'], + ['NaN', '81a4'], + ['5', ''], + ['5', '0x81a4'], + ['5', '81a'], + ['5', '81a45'], + ['5', '81a4 '], + ['9007199254740992', '81a4'], + ])('rejects malformed lstat fields', async (size, mode) => { + const handle = createHandle() + execSlot(handle).exec = async () => ({ + exitCode: 0, + stdout: `${mode}:${size}\n`, + stderr: '', + }) + await expect(handle.fs.lstat!('/workspace/file')).rejects.toThrow( + 'invalid lstat output', + ) + }) + it('parses a block special file as other without size', async () => { + const handle = createHandle() + execSlot(handle).exec = async () => ({ + exitCode: 0, + stdout: '61b6:0\n', + stderr: '', + }) + await expect(handle.fs.lstat!('/workspace/block')).resolves.toEqual({ + type: 'other', + mode: 0x61b6, + }) + }) + it('parses a fifo as other without size', async () => { + const handle = createHandle() + execSlot(handle).exec = async () => ({ + exitCode: 0, + stdout: '11b6:0\n', + stderr: '', + }) + await expect(handle.fs.lstat!('/workspace/fifo')).resolves.toEqual({ + type: 'other', + mode: 0x11b6, + }) + }) + it('parses an unknown file type as other without size', async () => { + const handle = createHandle() + execSlot(handle).exec = async () => ({ + exitCode: 0, + stdout: '71b6:7\n', + stderr: '', + }) + await expect(handle.fs.lstat!('/workspace/unknown')).resolves.toEqual({ + type: 'other', + mode: 0x71b6, + }) + }) + + it.each([ + '/workspace/missing', + '/workspace/missing-parent/child', + '-H/missing', + '-delete/missing', + ])('returns undefined for a verified missing path: %s', async (path) => { + const handle = createHandle() + execSlot(handle).exec = async (command) => { + expect(command).toBe(lstatCommand(path)) + return { + exitCode: 0, + stdout: '__TANSTACK_LSTAT_MISSING__', + stderr: '', + } + } + await expect(handle.fs.lstat!(path)).resolves.toBeUndefined() + }) + + it.each([ + '/workspace/file/child', + '/workspace/loop/child', + '/workspace/denied-link/child', + ])('preserves an unverified parent failure: %s', async (path) => { + const handle = createHandle() + execSlot(handle).exec = async (command) => { + expect(command).toBe(lstatCommand(path)) + return { exitCode: 1, stdout: '', stderr: 'permission denied' } + } + await expect(handle.fs.lstat!(path)).rejects.toThrow( + 'lstat failed: permission denied', + ) + }) +}) diff --git a/packages/ai-sandbox-docker/src/handle.ts b/packages/ai-sandbox-docker/src/handle.ts index 3eb5392540..d29ffb00b8 100644 --- a/packages/ai-sandbox-docker/src/handle.ts +++ b/packages/ai-sandbox-docker/src/handle.ts @@ -23,6 +23,7 @@ import type { SandboxHandle, SnapshotRef, SpawnHandle, + SandboxFsStat, } from '@tanstack/ai-sandbox' export const DOCKER_CAPS: SandboxCapabilities = { @@ -98,6 +99,30 @@ export interface DockerLogger { function q(value: string): string { return `'${value.replace(/'/g, `'\\''`)}'` } +/** Verify a missing path by listing parent entries. `test -e` also fails for inaccessible parents. */ +function lstatCommand(path: string): string { + return `tanstack_lstat_path=${q(path)}; tanstack_lstat_output=$(stat -c '%f:%s' -- "$tanstack_lstat_path" 2>&1); tanstack_lstat_status=$?; if [ "$tanstack_lstat_status" -eq 0 ]; then printf '%s\n' "$tanstack_lstat_output"; else tanstack_lstat_missing() { tanstack_missing_path=$1; case "$tanstack_missing_path" in /|.) return 1 ;; */*) tanstack_parent=${'$'}{tanstack_missing_path%/*}; tanstack_name=${'$'}{tanstack_missing_path##*/}; [ -n "$tanstack_parent" ] || tanstack_parent=/ ;; *) tanstack_parent=.; tanstack_name=$tanstack_missing_path ;; esac; tanstack_parent_mode=$(stat -L -c '%f' -- "$tanstack_parent" 2>/dev/null); tanstack_parent_status=$?; if [ "$tanstack_parent_status" -ne 0 ]; then tanstack_lstat_missing "$tanstack_parent"; else case "$tanstack_parent_mode" in 4[0-9a-fA-F][0-9a-fA-F][0-9a-fA-F]) case "$tanstack_parent" in /*) tanstack_find_parent=$tanstack_parent ;; *) tanstack_find_parent=./$tanstack_parent ;; esac; tanstack_match=$(find -H "$tanstack_find_parent" -mindepth 1 -maxdepth 1 -exec sh -c 'tanstack_target=$1; shift; for tanstack_candidate do [ "${'$'}{tanstack_candidate##*/}" = "$tanstack_target" ] && { printf 1; exit 0; }; done; exit 0' sh "$tanstack_name" '{}' + 2>/dev/null); tanstack_find_status=$?; [ "$tanstack_find_status" -eq 0 ] && [ -z "$tanstack_match" ] ;; *) return 1 ;; esac; fi; }; if tanstack_lstat_missing "$tanstack_lstat_path"; then printf '%s' '__TANSTACK_LSTAT_MISSING__'; else printf '%s\n' "$tanstack_lstat_output" >&2; exit "$tanstack_lstat_status"; fi; fi` +} +function parseLstatOutput(output: string): SandboxFsStat { + const fields = /^(?[0-9a-fA-F]{4}):(?\d+)\n?$/.exec(output) + const mode = fields?.groups?.mode + const size = fields?.groups?.size + if (!mode || !size) throw new Error(`invalid lstat output: ${output}`) + const parsedMode = Number.parseInt(mode, 16) + const parsedSize = Number(size) + if ( + !Number.isSafeInteger(parsedMode) || + !Number.isSafeInteger(parsedSize) || + parsedSize < 0 + ) + throw new Error(`invalid lstat output: ${output}`) + const type = parsedMode & 0xf000 + if (type === 0x8000) + return { type: 'file', mode: parsedMode, size: parsedSize } + if (type === 0x4000) return { type: 'dir', mode: parsedMode } + if (type === 0xa000) return { type: 'symlink', mode: parsedMode } + return { type: 'other', mode: parsedMode } +} /** * Marker the kill shell prints on stderr when it could not signal the process. @@ -391,6 +416,7 @@ export class DockerHandle implements SandboxHandle { } }) }, + lstat: async (p) => this.lstat(this.abs(p)), mkdir: async (p) => { await this.exec(`mkdir -p ${q(this.abs(p))}`) }, @@ -429,6 +455,16 @@ export class DockerHandle implements SandboxHandle { return p } + private async lstat(path: string): Promise { + const r = await this.exec(lstatCommand(path)) + if (r.exitCode === 0 && r.stdout === '__TANSTACK_LSTAT_MISSING__') + return undefined + if (r.exitCode !== 0) { + throw new Error(`lstat failed: ${r.stderr.trim()}`) + } + return parseLstatOutput(r.stdout) + } + private envArray(extra?: Record): Array { return Object.entries({ ...this.envVars, ...extra }).map( ([k, v]) => `${k}=${v}`, diff --git a/packages/ai-sandbox-docker/tests/lstat-shell-protocol.test.ts b/packages/ai-sandbox-docker/tests/lstat-shell-protocol.test.ts new file mode 100644 index 0000000000..14292b0ba2 --- /dev/null +++ b/packages/ai-sandbox-docker/tests/lstat-shell-protocol.test.ts @@ -0,0 +1,159 @@ +import { execFile } from 'node:child_process' +import { + access, + chmod, + mkdir, + mkdtemp, + readFile, + rm, + symlink, + writeFile, +} from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { describe, expect, it } from 'vitest' + +function quote(value: string): string { + return `'${value.replace(/'/g, `'\\''`)}'` +} + +function lstatCommand(path: string): string { + return `tanstack_lstat_path=${quote(path)}; tanstack_lstat_output=$(stat -c '%f:%s' -- "$tanstack_lstat_path" 2>&1); tanstack_lstat_status=$?; if [ "$tanstack_lstat_status" -eq 0 ]; then printf '%s\n' "$tanstack_lstat_output"; else tanstack_lstat_missing() { tanstack_missing_path=$1; case "$tanstack_missing_path" in /|.) return 1 ;; */*) tanstack_parent=${'$'}{tanstack_missing_path%/*}; tanstack_name=${'$'}{tanstack_missing_path##*/}; [ -n "$tanstack_parent" ] || tanstack_parent=/ ;; *) tanstack_parent=.; tanstack_name=$tanstack_missing_path ;; esac; tanstack_parent_mode=$(stat -L -c '%f' -- "$tanstack_parent" 2>/dev/null); tanstack_parent_status=$?; if [ "$tanstack_parent_status" -ne 0 ]; then tanstack_lstat_missing "$tanstack_parent"; else case "$tanstack_parent_mode" in 4[0-9a-fA-F][0-9a-fA-F][0-9a-fA-F]) case "$tanstack_parent" in /*) tanstack_find_parent=$tanstack_parent ;; *) tanstack_find_parent=./$tanstack_parent ;; esac; tanstack_match=$(find -H "$tanstack_find_parent" -mindepth 1 -maxdepth 1 -exec sh -c 'tanstack_target=$1; shift; for tanstack_candidate do [ "${'$'}{tanstack_candidate##*/}" = "$tanstack_target" ] && { printf 1; exit 0; }; done; exit 0' sh "$tanstack_name" '{}' + 2>/dev/null); tanstack_find_status=$?; [ "$tanstack_find_status" -eq 0 ] && [ -z "$tanstack_match" ] ;; *) return 1 ;; esac; fi; }; if tanstack_lstat_missing "$tanstack_lstat_path"; then printf '%s' '__TANSTACK_LSTAT_MISSING__'; else printf '%s\n' "$tanstack_lstat_output" >&2; exit "$tanstack_lstat_status"; fi; fi` +} + +function runShell(path: string, cwd?: string) { + return new Promise<{ exitCode: number; stdout: string; stderr: string }>( + (resolve) => { + execFile( + 'sh', + ['-c', lstatCommand(path)], + { cwd, encoding: 'utf8' }, + (error, stdout, stderr) => { + resolve({ + exitCode: typeof error?.code === 'number' ? error.code : 0, + stdout, + stderr, + }) + }, + ) + }, + ) +} + +const describeShell = process.platform === 'linux' ? describe : describe.skip + +describeShell('remote lstat shell protocol', () => { + it('proves missing paths and follows a command-line parent symlink', async () => { + const root = await mkdtemp(join(tmpdir(), 'tanstack-lstat-')) + const directory = join(root, 'directory') + const link = join(root, 'directory-link') + await mkdir(directory) + await symlink(directory, link) + + try { + for (const path of [ + join(directory, 'missing'), + join(root, 'missing-parent', 'child'), + join(link, 'missing'), + ]) { + await expect(runShell(path)).resolves.toEqual({ + exitCode: 0, + stdout: '__TANSTACK_LSTAT_MISSING__', + stderr: '', + }) + } + } finally { + await rm(root, { recursive: true, force: true }) + } + }) + + it('preserves dangling-link metadata and rejects unverified parents', async () => { + const root = await mkdtemp(join(tmpdir(), 'tanstack-lstat-')) + const file = join(root, 'file') + const dangling = join(root, 'dangling') + const loop = join(root, 'loop') + await writeFile(file, 'x') + await symlink('missing-target', dangling) + await symlink('loop', loop) + + try { + await expect(runShell(dangling)).resolves.toMatchObject({ + exitCode: 0, + stdout: expect.stringMatching(/^a[0-9a-f]{3}:\d+\n$/i), + stderr: '', + }) + for (const path of [join(file, 'child'), join(loop, 'child')]) { + const result = await runShell(path) + expect(result.exitCode).not.toBe(0) + expect(result.stdout).not.toContain('__TANSTACK_LSTAT_MISSING__') + expect(result.stderr).not.toBe('') + } + } finally { + await rm(root, { recursive: true, force: true }) + } + }) + + it.skipIf(typeof process.getuid === 'function' && process.getuid() === 0)( + 'preserves EACCES through a symlink parent', + async () => { + const root = await mkdtemp(join(tmpdir(), 'tanstack-lstat-')) + const denied = join(root, 'denied') + const link = join(root, 'denied-link') + await mkdir(denied) + await writeFile(join(denied, 'hidden'), 'x') + await symlink(denied, link) + await chmod(denied, 0o000) + + try { + const result = await runShell(join(link, 'hidden')) + expect(result.exitCode).not.toBe(0) + expect(result.stdout).not.toContain('__TANSTACK_LSTAT_MISSING__') + expect(result.stderr).not.toBe('') + } finally { + await chmod(denied, 0o700) + await rm(root, { recursive: true, force: true }) + } + }, + ) + + it.skipIf(typeof process.getuid === 'function' && process.getuid() === 0)( + 'preserves EACCES for a relative -H parent', + async () => { + const root = await mkdtemp(join(tmpdir(), 'tanstack-lstat-')) + const denied = join(root, '-H') + await mkdir(denied) + await writeFile(join(denied, 'hidden'), 'x') + await chmod(denied, 0o000) + + try { + const result = await runShell('-H/hidden', root) + expect(result.exitCode).not.toBe(0) + expect(result.stdout).not.toContain('__TANSTACK_LSTAT_MISSING__') + expect(result.stderr).not.toBe('') + } finally { + await chmod(denied, 0o700) + await rm(root, { recursive: true, force: true }) + } + }, + ) + + it('treats a relative -delete parent as a path without mutating files', async () => { + const root = await mkdtemp(join(tmpdir(), 'tanstack-lstat-')) + const optionLikeDirectory = join(root, '-delete') + const victim = join(root, 'victim') + await mkdir(optionLikeDirectory) + await writeFile(victim, 'keep me') + + try { + await expect(runShell('-delete/missing', root)).resolves.toEqual({ + exitCode: 0, + stdout: '__TANSTACK_LSTAT_MISSING__', + stderr: '', + }) + await expect(access(optionLikeDirectory)).resolves.toBeUndefined() + await expect(readFile(victim, 'utf8')).resolves.toBe('keep me') + } finally { + await rm(root, { recursive: true, force: true }) + } + }) +}) diff --git a/packages/ai-sandbox-docker/tests/lstat.test.ts b/packages/ai-sandbox-docker/tests/lstat.test.ts new file mode 100644 index 0000000000..0a2106f588 --- /dev/null +++ b/packages/ai-sandbox-docker/tests/lstat.test.ts @@ -0,0 +1,204 @@ +import { describe, expect, it } from 'vitest' +import Docker from 'dockerode' +import { DockerHandle } from '../src/handle' + +type ExecResult = { exitCode: number; stdout: string; stderr: string } +function lstatCommand(path: string): string { + const quoted = `'${path.replace(/'/g, `'\\''`)}'` + return `tanstack_lstat_path=${quoted}; tanstack_lstat_output=$(stat -c '%f:%s' -- "$tanstack_lstat_path" 2>&1); tanstack_lstat_status=$?; if [ "$tanstack_lstat_status" -eq 0 ]; then printf '%s\n' "$tanstack_lstat_output"; else tanstack_lstat_missing() { tanstack_missing_path=$1; case "$tanstack_missing_path" in /|.) return 1 ;; */*) tanstack_parent=${'$'}{tanstack_missing_path%/*}; tanstack_name=${'$'}{tanstack_missing_path##*/}; [ -n "$tanstack_parent" ] || tanstack_parent=/ ;; *) tanstack_parent=.; tanstack_name=$tanstack_missing_path ;; esac; tanstack_parent_mode=$(stat -L -c '%f' -- "$tanstack_parent" 2>/dev/null); tanstack_parent_status=$?; if [ "$tanstack_parent_status" -ne 0 ]; then tanstack_lstat_missing "$tanstack_parent"; else case "$tanstack_parent_mode" in 4[0-9a-fA-F][0-9a-fA-F][0-9a-fA-F]) case "$tanstack_parent" in /*) tanstack_find_parent=$tanstack_parent ;; *) tanstack_find_parent=./$tanstack_parent ;; esac; tanstack_match=$(find -H "$tanstack_find_parent" -mindepth 1 -maxdepth 1 -exec sh -c 'tanstack_target=$1; shift; for tanstack_candidate do [ "${'$'}{tanstack_candidate##*/}" = "$tanstack_target" ] && { printf 1; exit 0; }; done; exit 0' sh "$tanstack_name" '{}' + 2>/dev/null); tanstack_find_status=$?; [ "$tanstack_find_status" -eq 0 ] && [ -z "$tanstack_match" ] ;; *) return 1 ;; esac; fi; }; if tanstack_lstat_missing "$tanstack_lstat_path"; then printf '%s' '__TANSTACK_LSTAT_MISSING__'; else printf '%s\n' "$tanstack_lstat_output" >&2; exit "$tanstack_lstat_status"; fi; fi` +} +function lstatPath(command: string): string { + return /^tanstack_lstat_path='([^']*)';/.exec(command)?.[1] ?? '' +} + +function setExec( + handle: object, + exec: (command: string) => Promise, +) { + Object.defineProperty(handle, 'exec', { configurable: true, value: exec }) +} + +function createHandle(): DockerHandle { + const docker = new Docker() + let handle: DockerHandle + handle = new DockerHandle({ + docker, + container: docker.getContainer('test'), + workdir: '/workspace', + forkFactory: () => Promise.resolve(handle), + removeOnDestroy: false, + }) + return handle +} + +describe('DockerHandle.fs.lstat', () => { + it.each([ + '/workspace/missing', + '/workspace/missing-parent/child', + '-H/missing', + '-delete/missing', + ])('returns undefined for a verified missing path: %s', async (path) => { + const handle = createHandle() + setExec(handle, async (command) => { + expect(command).toBe(lstatCommand(path)) + return { + exitCode: 0, + stdout: '__TANSTACK_LSTAT_MISSING__', + stderr: '', + } + }) + await expect(handle.fs.lstat!(path)).resolves.toBeUndefined() + }) + + it.each([ + '/workspace/file/child', + '/workspace/loop/child', + '/workspace/denied-link/child', + ])('preserves an unverified parent failure: %s', async (path) => { + const handle = createHandle() + setExec(handle, async (command) => { + expect(command).toBe(lstatCommand(path)) + return { exitCode: 1, stdout: '', stderr: 'permission denied' } + }) + await expect(handle.fs.lstat!(path)).rejects.toThrow('permission denied') + }) + + it('parses file, directory, symlink, and other metadata', async () => { + const values = new Map([ + ['file', '81A4:12\n'], + ['dir', '41ed:4096\n'], + ['link', 'a1ff:4\n'], + ['other', 'c1b6:0\n'], + ['char', '21b6:0\n'], + ['block', '61b6:0\n'], + ['fifo', '11b6:0\n'], + ['unknown', '71b6:7\n'], + ]) + const handle = createHandle() + setExec(handle, async (command: string) => { + const path = lstatPath(command) + expect(command).toBe(lstatCommand(path)) + return { + exitCode: 0, + stdout: values.get(lstatPath(command).split('/').pop() ?? '') ?? '', + stderr: '', + } + }) + await expect(handle.fs.lstat!('/workspace/file')).resolves.toEqual({ + type: 'file', + mode: 0x81a4, + size: 12, + }) + await expect(handle.fs.lstat!('/workspace/dir')).resolves.toEqual({ + type: 'dir', + mode: 0x41ed, + }) + await expect(handle.fs.lstat!('/workspace/link')).resolves.toEqual({ + type: 'symlink', + mode: 0xa1ff, + }) + await expect(handle.fs.lstat!('/workspace/other')).resolves.toEqual({ + type: 'other', + mode: 0xc1b6, + }) + }) + it('parses a character special file as other without size', async () => { + const handle = createHandle() + setExec(handle, async () => ({ + exitCode: 0, + stdout: '21b6:0\n', + stderr: '', + })) + await expect(handle.fs.lstat!('/workspace/char')).resolves.toEqual({ + type: 'other', + mode: 0x21b6, + }) + }) + it('parses a zero-byte regular empty file with size zero', async () => { + const handle = createHandle() + setExec(handle, async () => ({ + exitCode: 0, + stdout: '81a4:0\n', + stderr: '', + })) + await expect(handle.fs.lstat!('/workspace/empty')).resolves.toEqual({ + type: 'file', + mode: 0x81a4, + size: 0, + }) + }) + it.each([ + ['not-a-number', '81a4'], + ['Infinity', '81a4'], + ['-1', '81a4'], + ['', '81a4'], + ['5', '81a4junk'], + ['5', '-81a4'], + [' 5', '81a4'], + ['5 ', '81a4'], + ['5\n6', '81a4'], + ['NaN', '81a4'], + ['5', ''], + ['5', '0x81a4'], + ['5', '81a'], + ['5', '81a45'], + ['5', '81a4 '], + ['9007199254740992', '81a4'], + ])('rejects malformed lstat fields', async (size, mode) => { + const handle = createHandle() + setExec(handle, async () => ({ + exitCode: 0, + stdout: `${mode}:${size}\n`, + stderr: '', + })) + await expect(handle.fs.lstat!('/workspace/file')).rejects.toThrow( + 'invalid lstat output', + ) + }) + it('parses a block special file as other without size', async () => { + const handle = createHandle() + setExec(handle, async () => ({ + exitCode: 0, + stdout: '61b6:0\n', + stderr: '', + })) + await expect(handle.fs.lstat!('/workspace/block')).resolves.toEqual({ + type: 'other', + mode: 0x61b6, + }) + }) + it('parses a fifo as other without size', async () => { + const handle = createHandle() + setExec(handle, async () => ({ + exitCode: 0, + stdout: '11b6:0\n', + stderr: '', + })) + await expect(handle.fs.lstat!('/workspace/fifo')).resolves.toEqual({ + type: 'other', + mode: 0x11b6, + }) + }) + it('parses an unknown file type as other without size', async () => { + const handle = createHandle() + setExec(handle, async () => ({ + exitCode: 0, + stdout: '71b6:7\n', + stderr: '', + })) + await expect(handle.fs.lstat!('/workspace/unknown')).resolves.toEqual({ + type: 'other', + mode: 0x71b6, + }) + }) + it.each([' 81a4:12\n', '81a4:12\nextra\n'])( + 'rejects non-protocol whitespace/payload: %j', + async (stdout) => { + const handle = createHandle() + setExec(handle, async () => ({ exitCode: 0, stdout, stderr: '' })) + await expect(handle.fs.lstat!('/workspace/file')).rejects.toThrow( + 'invalid lstat output', + ) + }, + ) +}) diff --git a/packages/ai-sandbox-docker/tests/testkit-subpath.test.ts b/packages/ai-sandbox-docker/tests/testkit-subpath.test.ts new file mode 100644 index 0000000000..a2e871bb1d --- /dev/null +++ b/packages/ai-sandbox-docker/tests/testkit-subpath.test.ts @@ -0,0 +1,6 @@ +import { expect, it } from 'vitest' +import { runSandboxCheckpointStoreConformance } from '@tanstack/ai-sandbox/testkit' + +it('exports the checkpoint store conformance suite from the built testkit subpath', () => { + expect(typeof runSandboxCheckpointStoreConformance).toBe('function') +}) diff --git a/packages/ai-sandbox-local-process/src/handle.ts b/packages/ai-sandbox-local-process/src/handle.ts index 10a012c427..db4be78e28 100644 --- a/packages/ai-sandbox-local-process/src/handle.ts +++ b/packages/ai-sandbox-local-process/src/handle.ts @@ -639,6 +639,31 @@ export class LocalProcessHandle implements SandboxHandle { type: e.isDirectory() ? ('dir' as const) : ('file' as const), })) }, + lstat: async (p) => { + let stat: Awaited> + try { + stat = await fsp.lstat(this.resolve(p)) + } catch (error) { + if ( + error !== null && + typeof error === 'object' && + 'code' in error && + error.code === 'ENOENT' + ) + return undefined + throw error + } + const type = stat.isFile() + ? 'file' + : stat.isDirectory() + ? 'dir' + : stat.isSymbolicLink() + ? 'symlink' + : 'other' + return type === 'file' + ? { type, mode: stat.mode, size: stat.size } + : { type, mode: stat.mode } + }, mkdir: async (p) => { await fsp.mkdir(this.resolve(p), { recursive: true }) }, diff --git a/packages/ai-sandbox-local-process/tests/local-process.test.ts b/packages/ai-sandbox-local-process/tests/local-process.test.ts index cf46230cd6..2e1d014913 100644 --- a/packages/ai-sandbox-local-process/tests/local-process.test.ts +++ b/packages/ai-sandbox-local-process/tests/local-process.test.ts @@ -1,4 +1,4 @@ -import { afterAll, describe, expect, it } from 'vitest' +import { afterAll, afterEach, describe, expect, it, vi } from 'vitest' import * as fsp from 'node:fs/promises' import * as os from 'node:os' import * as path from 'node:path' @@ -12,6 +12,31 @@ import { import { localProcessSandbox } from '../src/index' import type { SandboxHandle } from '@tanstack/ai-sandbox' +const lstatControl = vi.hoisted(() => { + let failure: Error | undefined + return { + getFailure: () => failure, + setFailure: (error: Error) => { + failure = error + }, + reset: () => { + failure = undefined + }, + } +}) + +vi.mock('node:fs/promises', async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + lstat: async (path: Parameters[0]) => { + const failure = lstatControl.getFailure() + if (failure) throw failure + return actual.lstat(path) + }, + } +}) + const baseDir = path.join(os.tmpdir(), `tanstack-ai-lp-test-${Date.now()}`) const provider = localProcessSandbox({ baseDir, removeOnDestroy: true }) @@ -19,11 +44,29 @@ afterAll(async () => { await fsp.rm(baseDir, { recursive: true, force: true }) }) +afterEach(() => { + lstatControl.reset() +}) + async function fresh(): Promise { return provider.create({}) } describe('local-process fs', () => { + it('returns undefined for a missing path', async () => { + const sbx = await fresh() + await expect(sbx.fs.lstat!('/workspace/missing')).resolves.toBeUndefined() + await sbx.destroy() + }) + + it('rejects a non-missing lstat error', async () => { + const sbx = await fresh() + const error = new Error('permission denied') + lstatControl.setFailure(error) + await expect(sbx.fs.lstat!('/workspace/file')).rejects.toBe(error) + await sbx.destroy() + }) + it('writes, reads, lists, renames, removes', async () => { const sbx = await fresh() await sbx.fs.write('/workspace/a.txt', 'hello') diff --git a/packages/ai-sandbox-sprites/src/handle.ts b/packages/ai-sandbox-sprites/src/handle.ts index 442ace6e73..143d784b7f 100644 --- a/packages/ai-sandbox-sprites/src/handle.ts +++ b/packages/ai-sandbox-sprites/src/handle.ts @@ -26,6 +26,7 @@ import type { SandboxHandle, SnapshotRef, SpawnHandle, + SandboxFsStat, } from '@tanstack/ai-sandbox' import type { SpriteCheckpoint, @@ -165,6 +166,7 @@ export class SpritesHandle implements SandboxHandle { type: entry.type, })) }, + lstat: async (p) => this.lstat(this.abs(p)), mkdir: async (p) => { const r = await this.exec(`mkdir -p ${q(this.abs(p))}`) if (r.exitCode !== 0) throw new Error(`mkdir failed: ${errText(r)}`) @@ -206,6 +208,16 @@ export class SpritesHandle implements SandboxHandle { return p } + private async lstat(path: string): Promise { + const r = await this.exec(lstatCommand(path)) + if (r.exitCode === 0 && r.stdout === '__TANSTACK_LSTAT_MISSING__') + return undefined + if (r.exitCode !== 0) { + throw new Error(`lstat failed: ${errText(r)}`) + } + return parseLstatOutput(r.stdout) + } + private mergedEnv(extra?: Record): Record { return { ...this.envVars, ...extra } } @@ -342,6 +354,30 @@ export class SpritesHandle implements SandboxHandle { function q(value: string): string { return `'${value.replace(/'/g, `'\\''`)}'` } +/** Verify a missing path by listing parent entries. `test -e` also fails for inaccessible parents. */ +function lstatCommand(path: string): string { + return `tanstack_lstat_path=${q(path)}; tanstack_lstat_output=$(stat -c '%f:%s' -- "$tanstack_lstat_path" 2>&1); tanstack_lstat_status=$?; if [ "$tanstack_lstat_status" -eq 0 ]; then printf '%s\n' "$tanstack_lstat_output"; else tanstack_lstat_missing() { tanstack_missing_path=$1; case "$tanstack_missing_path" in /|.) return 1 ;; */*) tanstack_parent=${'$'}{tanstack_missing_path%/*}; tanstack_name=${'$'}{tanstack_missing_path##*/}; [ -n "$tanstack_parent" ] || tanstack_parent=/ ;; *) tanstack_parent=.; tanstack_name=$tanstack_missing_path ;; esac; tanstack_parent_mode=$(stat -L -c '%f' -- "$tanstack_parent" 2>/dev/null); tanstack_parent_status=$?; if [ "$tanstack_parent_status" -ne 0 ]; then tanstack_lstat_missing "$tanstack_parent"; else case "$tanstack_parent_mode" in 4[0-9a-fA-F][0-9a-fA-F][0-9a-fA-F]) case "$tanstack_parent" in /*) tanstack_find_parent=$tanstack_parent ;; *) tanstack_find_parent=./$tanstack_parent ;; esac; tanstack_match=$(find -H "$tanstack_find_parent" -mindepth 1 -maxdepth 1 -exec sh -c 'tanstack_target=$1; shift; for tanstack_candidate do [ "${'$'}{tanstack_candidate##*/}" = "$tanstack_target" ] && { printf 1; exit 0; }; done; exit 0' sh "$tanstack_name" '{}' + 2>/dev/null); tanstack_find_status=$?; [ "$tanstack_find_status" -eq 0 ] && [ -z "$tanstack_match" ] ;; *) return 1 ;; esac; fi; }; if tanstack_lstat_missing "$tanstack_lstat_path"; then printf '%s' '__TANSTACK_LSTAT_MISSING__'; else printf '%s\n' "$tanstack_lstat_output" >&2; exit "$tanstack_lstat_status"; fi; fi` +} +function parseLstatOutput(output: string): SandboxFsStat { + const fields = /^(?[0-9a-fA-F]{4}):(?\d+)\n?$/.exec(output) + const mode = fields?.groups?.mode + const size = fields?.groups?.size + if (!mode || !size) throw new Error(`invalid lstat output: ${output}`) + const parsedMode = Number.parseInt(mode, 16) + const parsedSize = Number(size) + if ( + !Number.isSafeInteger(parsedMode) || + !Number.isSafeInteger(parsedSize) || + parsedSize < 0 + ) + throw new Error(`invalid lstat output: ${output}`) + const type = parsedMode & 0xf000 + if (type === 0x8000) + return { type: 'file', mode: parsedMode, size: parsedSize } + if (type === 0x4000) return { type: 'dir', mode: parsedMode } + if (type === 0xa000) return { type: 'symlink', mode: parsedMode } + return { type: 'other', mode: parsedMode } +} /** * Best error text from an exec result. Near-instant commands hit the Sprite diff --git a/packages/ai-sandbox-sprites/tests/lstat.test.ts b/packages/ai-sandbox-sprites/tests/lstat.test.ts new file mode 100644 index 0000000000..b4f6866dff --- /dev/null +++ b/packages/ai-sandbox-sprites/tests/lstat.test.ts @@ -0,0 +1,201 @@ +import { describe, expect, it } from 'vitest' +import { SpritesHandle } from '../src/handle' +import type { SpritesClientLike } from '../src/client' + +type ExecResult = { exitCode: number; stdout: string; stderr: string } +type ExecFn = (command: string) => Promise +function lstatCommand(path: string): string { + const quoted = `'${path.replace(/'/g, `'\\''`)}'` + return `tanstack_lstat_path=${quoted}; tanstack_lstat_output=$(stat -c '%f:%s' -- "$tanstack_lstat_path" 2>&1); tanstack_lstat_status=$?; if [ "$tanstack_lstat_status" -eq 0 ]; then printf '%s\n' "$tanstack_lstat_output"; else tanstack_lstat_missing() { tanstack_missing_path=$1; case "$tanstack_missing_path" in /|.) return 1 ;; */*) tanstack_parent=${'$'}{tanstack_missing_path%/*}; tanstack_name=${'$'}{tanstack_missing_path##*/}; [ -n "$tanstack_parent" ] || tanstack_parent=/ ;; *) tanstack_parent=.; tanstack_name=$tanstack_missing_path ;; esac; tanstack_parent_mode=$(stat -L -c '%f' -- "$tanstack_parent" 2>/dev/null); tanstack_parent_status=$?; if [ "$tanstack_parent_status" -ne 0 ]; then tanstack_lstat_missing "$tanstack_parent"; else case "$tanstack_parent_mode" in 4[0-9a-fA-F][0-9a-fA-F][0-9a-fA-F]) case "$tanstack_parent" in /*) tanstack_find_parent=$tanstack_parent ;; *) tanstack_find_parent=./$tanstack_parent ;; esac; tanstack_match=$(find -H "$tanstack_find_parent" -mindepth 1 -maxdepth 1 -exec sh -c 'tanstack_target=$1; shift; for tanstack_candidate do [ "${'$'}{tanstack_candidate##*/}" = "$tanstack_target" ] && { printf 1; exit 0; }; done; exit 0' sh "$tanstack_name" '{}' + 2>/dev/null); tanstack_find_status=$?; [ "$tanstack_find_status" -eq 0 ] && [ -z "$tanstack_match" ] ;; *) return 1 ;; esac; fi; }; if tanstack_lstat_missing "$tanstack_lstat_path"; then printf '%s' '__TANSTACK_LSTAT_MISSING__'; else printf '%s\n' "$tanstack_lstat_output" >&2; exit "$tanstack_lstat_status"; fi; fi` +} +function lstatPath(command: string): string { + return /^tanstack_lstat_path='([^']*)';/.exec(command)?.[1] ?? '' +} +function execSlot(handle: object) { + return { + set exec(value: ExecFn) { + Object.defineProperty(handle, 'exec', { configurable: true, value }) + }, + } +} + +function createHandle(): SpritesHandle { + const client: SpritesClientLike = { + baseUrl: 'https://api.test', + authHeader: () => ({ authorization: 'Bearer test' }), + getSprite: () => Promise.reject(new Error('not used')), + deleteSprite: () => Promise.resolve(), + setUrlAuth: () => Promise.resolve(), + fsRead: () => Promise.reject(new Error('not used')), + fsWrite: () => Promise.resolve(), + fsList: () => Promise.resolve([]), + exec: () => { + throw new Error('exec is replaced by each test') + }, + createCheckpoint: () => Promise.resolve('v1'), + listCheckpoints: () => Promise.resolve([]), + restoreCheckpoint: () => Promise.resolve(), + } + return new SpritesHandle({ + client, + name: 'test', + url: 'https://test', + workdir: '/home/sprite', + }) +} + +describe('SpritesHandle.fs.lstat', () => { + it.each([ + '/workspace/missing', + '/workspace/missing-parent/child', + '-H/missing', + '-delete/missing', + ])('returns undefined for a verified missing path: %s', async (path) => { + const handle = createHandle() + execSlot(handle).exec = async (command) => { + const resolved = path.startsWith('/workspace') + ? `/home/sprite${path.slice(10)}` + : path + expect(command).toBe(lstatCommand(resolved)) + return { + exitCode: 0, + stdout: '__TANSTACK_LSTAT_MISSING__', + stderr: '', + } + } + await expect(handle.fs.lstat!(path)).resolves.toBeUndefined() + }) + + it.each([ + '/workspace/file/child', + '/workspace/loop/child', + '/workspace/denied-link/child', + ])('preserves an unverified parent failure: %s', async (path) => { + const handle = createHandle() + execSlot(handle).exec = async (command) => { + expect(command).toBe(lstatCommand(`/home/sprite${path.slice(10)}`)) + return { exitCode: 1, stdout: '', stderr: 'permission denied' } + } + await expect(handle.fs.lstat!(path)).rejects.toThrow('permission denied') + }) + + it('parses file, directory, symlink, and other metadata', async () => { + const values = new Map([ + ['file', '81A4:12\n'], + ['dir', '41ed:4096\n'], + ['link', 'a1ff:4\n'], + ['other', 'c1b6:0\n'], + ['char', '21b6:0\n'], + ['block', '61b6:0\n'], + ['fifo', '11b6:0\n'], + ['unknown', '71b6:7\n'], + ]) + const handle = createHandle() + execSlot(handle).exec = async (command: string) => { + expect(command).toBe(lstatCommand(lstatPath(command))) + return { + exitCode: 0, + stdout: values.get(lstatPath(command).split('/').pop() ?? '') ?? '', + stderr: '', + } + } + for (const [name, expected] of [ + ['file', { type: 'file', mode: 0x81a4, size: 12 }], + ['dir', { type: 'dir', mode: 0x41ed }], + ['link', { type: 'symlink', mode: 0xa1ff }], + ['other', { type: 'other', mode: 0xc1b6 }], + ] as const) + await expect(handle.fs.lstat!(`/workspace/${name}`)).resolves.toEqual( + expected, + ) + }) + it('parses a character special file as other without size', async () => { + const handle = createHandle() + execSlot(handle).exec = async () => ({ + exitCode: 0, + stdout: '21b6:0\n', + stderr: '', + }) + await expect(handle.fs.lstat!('/workspace/char')).resolves.toEqual({ + type: 'other', + mode: 0x21b6, + }) + }) + it('parses a zero-byte regular empty file with size zero', async () => { + const handle = createHandle() + execSlot(handle).exec = async () => ({ + exitCode: 0, + stdout: '81a4:0\n', + stderr: '', + }) + await expect(handle.fs.lstat!('/workspace/empty')).resolves.toEqual({ + type: 'file', + mode: 0x81a4, + size: 0, + }) + }) + it.each([ + ['not-a-number', '81a4'], + ['Infinity', '81a4'], + ['-1', '81a4'], + ['', '81a4'], + ['5', '81a4junk'], + ['5', '-81a4'], + [' 5', '81a4'], + ['5 ', '81a4'], + ['5\n6', '81a4'], + ['NaN', '81a4'], + ['5', ''], + ['5', '0x81a4'], + ['5', '81a'], + ['5', '81a45'], + ['5', '81a4 '], + ['9007199254740992', '81a4'], + ])('rejects malformed lstat fields', async (size, mode) => { + const handle = createHandle() + execSlot(handle).exec = async () => ({ + exitCode: 0, + stdout: `${mode}:${size}\n`, + stderr: '', + }) + await expect(handle.fs.lstat!('/workspace/file')).rejects.toThrow( + 'invalid lstat output', + ) + }) + it('parses a block special file as other without size', async () => { + const handle = createHandle() + execSlot(handle).exec = async () => ({ + exitCode: 0, + stdout: '61b6:0\n', + stderr: '', + }) + await expect(handle.fs.lstat!('/workspace/block')).resolves.toEqual({ + type: 'other', + mode: 0x61b6, + }) + }) + it('parses a fifo as other without size', async () => { + const handle = createHandle() + execSlot(handle).exec = async () => ({ + exitCode: 0, + stdout: '11b6:0\n', + stderr: '', + }) + await expect(handle.fs.lstat!('/workspace/fifo')).resolves.toEqual({ + type: 'other', + mode: 0x11b6, + }) + }) + it('parses an unknown file type as other without size', async () => { + const handle = createHandle() + execSlot(handle).exec = async () => ({ + exitCode: 0, + stdout: '71b6:7\n', + stderr: '', + }) + await expect(handle.fs.lstat!('/workspace/unknown')).resolves.toEqual({ + type: 'other', + mode: 0x71b6, + }) + }) +}) diff --git a/packages/ai-sandbox-vercel/src/handle.ts b/packages/ai-sandbox-vercel/src/handle.ts index 0c82395bb9..6dc2d09fce 100644 --- a/packages/ai-sandbox-vercel/src/handle.ts +++ b/packages/ai-sandbox-vercel/src/handle.ts @@ -20,6 +20,7 @@ import type { SandboxChannel, SandboxHandle, SpawnHandle, + SandboxFsStat, } from '@tanstack/ai-sandbox' export const VERCEL_CAPS: SandboxCapabilities = { @@ -64,6 +65,30 @@ export const VERCEL_CAPS: SandboxCapabilities = { function q(value: string): string { return `'${value.replace(/'/g, `'\\''`)}'` } +/** Verify a missing path by listing parent entries. `test -e` also fails for inaccessible parents. */ +function lstatCommand(path: string): string { + return `tanstack_lstat_path=${q(path)}; tanstack_lstat_output=$(stat -c '%f:%s' -- "$tanstack_lstat_path" 2>&1); tanstack_lstat_status=$?; if [ "$tanstack_lstat_status" -eq 0 ]; then printf '%s\n' "$tanstack_lstat_output"; else tanstack_lstat_missing() { tanstack_missing_path=$1; case "$tanstack_missing_path" in /|.) return 1 ;; */*) tanstack_parent=${'$'}{tanstack_missing_path%/*}; tanstack_name=${'$'}{tanstack_missing_path##*/}; [ -n "$tanstack_parent" ] || tanstack_parent=/ ;; *) tanstack_parent=.; tanstack_name=$tanstack_missing_path ;; esac; tanstack_parent_mode=$(stat -L -c '%f' -- "$tanstack_parent" 2>/dev/null); tanstack_parent_status=$?; if [ "$tanstack_parent_status" -ne 0 ]; then tanstack_lstat_missing "$tanstack_parent"; else case "$tanstack_parent_mode" in 4[0-9a-fA-F][0-9a-fA-F][0-9a-fA-F]) case "$tanstack_parent" in /*) tanstack_find_parent=$tanstack_parent ;; *) tanstack_find_parent=./$tanstack_parent ;; esac; tanstack_match=$(find -H "$tanstack_find_parent" -mindepth 1 -maxdepth 1 -exec sh -c 'tanstack_target=$1; shift; for tanstack_candidate do [ "${'$'}{tanstack_candidate##*/}" = "$tanstack_target" ] && { printf 1; exit 0; }; done; exit 0' sh "$tanstack_name" '{}' + 2>/dev/null); tanstack_find_status=$?; [ "$tanstack_find_status" -eq 0 ] && [ -z "$tanstack_match" ] ;; *) return 1 ;; esac; fi; }; if tanstack_lstat_missing "$tanstack_lstat_path"; then printf '%s' '__TANSTACK_LSTAT_MISSING__'; else printf '%s\n' "$tanstack_lstat_output" >&2; exit "$tanstack_lstat_status"; fi; fi` +} +function parseLstatOutput(output: string): SandboxFsStat { + const fields = /^(?[0-9a-fA-F]{4}):(?\d+)\n?$/.exec(output) + const mode = fields?.groups?.mode + const size = fields?.groups?.size + if (!mode || !size) throw new Error(`invalid lstat output: ${output}`) + const parsedMode = Number.parseInt(mode, 16) + const parsedSize = Number(size) + if ( + !Number.isSafeInteger(parsedMode) || + !Number.isSafeInteger(parsedSize) || + parsedSize < 0 + ) + throw new Error(`invalid lstat output: ${output}`) + const type = parsedMode & 0xf000 + if (type === 0x8000) + return { type: 'file', mode: parsedMode, size: parsedSize } + if (type === 0x4000) return { type: 'dir', mode: parsedMode } + if (type === 0xa000) return { type: 'symlink', mode: parsedMode } + return { type: 'other', mode: parsedMode } +} /** * A push-driven async iterable. The streamer pushes decoded chunks and calls @@ -184,6 +209,7 @@ export class VercelHandle implements SandboxHandle { } }) }, + lstat: async (p) => this.lstat(this.abs(p)), mkdir: async (p) => { await this.exec(`mkdir -p ${q(this.abs(p))}`) }, @@ -222,6 +248,16 @@ export class VercelHandle implements SandboxHandle { return p } + private async lstat(path: string): Promise { + const r = await this.exec(lstatCommand(path)) + if (r.exitCode === 0 && r.stdout === '__TANSTACK_LSTAT_MISSING__') + return undefined + if (r.exitCode !== 0) { + throw new Error(`lstat failed: ${r.stderr.trim()}`) + } + return parseLstatOutput(r.stdout) + } + private mergedEnv(extra?: Record): Record { return { ...this.envVars, ...extra } } diff --git a/packages/ai-sandbox-vercel/tests/lstat.test.ts b/packages/ai-sandbox-vercel/tests/lstat.test.ts new file mode 100644 index 0000000000..7fd72b9cdb --- /dev/null +++ b/packages/ai-sandbox-vercel/tests/lstat.test.ts @@ -0,0 +1,181 @@ +import { describe, expect, it } from 'vitest' +import { Sandbox } from '@vercel/sandbox' +import { VercelHandle } from '../src/handle' + +type ExecResult = { exitCode: number; stdout: string; stderr: string } +type ExecFn = (command: string) => Promise +function lstatCommand(path: string): string { + const quoted = `'${path.replace(/'/g, `'\\''`)}'` + return `tanstack_lstat_path=${quoted}; tanstack_lstat_output=$(stat -c '%f:%s' -- "$tanstack_lstat_path" 2>&1); tanstack_lstat_status=$?; if [ "$tanstack_lstat_status" -eq 0 ]; then printf '%s\n' "$tanstack_lstat_output"; else tanstack_lstat_missing() { tanstack_missing_path=$1; case "$tanstack_missing_path" in /|.) return 1 ;; */*) tanstack_parent=${'$'}{tanstack_missing_path%/*}; tanstack_name=${'$'}{tanstack_missing_path##*/}; [ -n "$tanstack_parent" ] || tanstack_parent=/ ;; *) tanstack_parent=.; tanstack_name=$tanstack_missing_path ;; esac; tanstack_parent_mode=$(stat -L -c '%f' -- "$tanstack_parent" 2>/dev/null); tanstack_parent_status=$?; if [ "$tanstack_parent_status" -ne 0 ]; then tanstack_lstat_missing "$tanstack_parent"; else case "$tanstack_parent_mode" in 4[0-9a-fA-F][0-9a-fA-F][0-9a-fA-F]) case "$tanstack_parent" in /*) tanstack_find_parent=$tanstack_parent ;; *) tanstack_find_parent=./$tanstack_parent ;; esac; tanstack_match=$(find -H "$tanstack_find_parent" -mindepth 1 -maxdepth 1 -exec sh -c 'tanstack_target=$1; shift; for tanstack_candidate do [ "${'$'}{tanstack_candidate##*/}" = "$tanstack_target" ] && { printf 1; exit 0; }; done; exit 0' sh "$tanstack_name" '{}' + 2>/dev/null); tanstack_find_status=$?; [ "$tanstack_find_status" -eq 0 ] && [ -z "$tanstack_match" ] ;; *) return 1 ;; esac; fi; }; if tanstack_lstat_missing "$tanstack_lstat_path"; then printf '%s' '__TANSTACK_LSTAT_MISSING__'; else printf '%s\n' "$tanstack_lstat_output" >&2; exit "$tanstack_lstat_status"; fi; fi` +} +function lstatPath(command: string): string { + return /^tanstack_lstat_path='([^']*)';/.exec(command)?.[1] ?? '' +} +function execSlot(handle: object) { + return { + set exec(value: ExecFn) { + Object.defineProperty(handle, 'exec', { configurable: true, value }) + }, + } +} + +function createHandle(): VercelHandle { + const sandbox = new Sandbox({ + routes: [], + sandbox: { + name: 'test', + persistent: false, + createdAt: 0, + updatedAt: 0, + currentSessionId: 'session', + status: 'running', + }, + }) + return new VercelHandle({ + sandbox, + workdir: '/workspace', + ports: [], + }) +} + +describe('VercelHandle.fs.lstat', () => { + it.each([ + '/workspace/missing', + '/workspace/missing-parent/child', + '-H/missing', + '-delete/missing', + ])('returns undefined for a verified missing path: %s', async (path) => { + const handle = createHandle() + execSlot(handle).exec = async (command) => { + expect(command).toBe(lstatCommand(path)) + return { + exitCode: 0, + stdout: '__TANSTACK_LSTAT_MISSING__', + stderr: '', + } + } + await expect(handle.fs.lstat!(path)).resolves.toBeUndefined() + }) + + it.each([ + '/workspace/file/child', + '/workspace/loop/child', + '/workspace/denied-link/child', + ])('preserves an unverified parent failure: %s', async (path) => { + const handle = createHandle() + execSlot(handle).exec = async (command) => { + expect(command).toBe(lstatCommand(path)) + return { exitCode: 1, stdout: '', stderr: 'permission denied' } + } + await expect(handle.fs.lstat!(path)).rejects.toThrow('permission denied') + }) + + it('parses file, directory, symlink, and other metadata', async () => { + const values = new Map([ + ['file', '81A4:12\n'], + ['empty', '81a4:0\n'], + ['dir', '41ed:4096\n'], + ['link', 'a1ff:4\n'], + ['other', 'c1b6:0\n'], + ['char', '21b6:0\n'], + ['block', '61b6:0\n'], + ['fifo', '11b6:0\n'], + ['unknown', '71b6:7\n'], + ]) + const handle = createHandle() + execSlot(handle).exec = async (command: string) => { + expect(command).toBe(lstatCommand(lstatPath(command))) + return { + exitCode: 0, + stdout: values.get(lstatPath(command).split('/').pop() ?? '') ?? '', + stderr: '', + } + } + for (const [name, expected] of [ + ['file', { type: 'file', mode: 0x81a4, size: 12 }], + ['empty', { type: 'file', mode: 0x81a4, size: 0 }], + ['dir', { type: 'dir', mode: 0x41ed }], + ['link', { type: 'symlink', mode: 0xa1ff }], + ['other', { type: 'other', mode: 0xc1b6 }], + ] as const) + await expect(handle.fs.lstat!(`/workspace/${name}`)).resolves.toEqual( + expected, + ) + }) + it('parses a character special file as other without size', async () => { + const handle = createHandle() + execSlot(handle).exec = async () => ({ + exitCode: 0, + stdout: '21b6:0\n', + stderr: '', + }) + await expect(handle.fs.lstat!('/workspace/char')).resolves.toEqual({ + type: 'other', + mode: 0x21b6, + }) + }) + it.each([ + ['not-a-number', '81a4'], + ['Infinity', '81a4'], + ['-1', '81a4'], + ['', '81a4'], + ['5', '81a4junk'], + ['5', '-81a4'], + [' 5', '81a4'], + ['5 ', '81a4'], + ['5\n6', '81a4'], + ['NaN', '81a4'], + ['5', ''], + ['5', '0x81a4'], + ['5', '81a'], + ['5', '81a45'], + ['5', '81a4 '], + ['9007199254740992', '81a4'], + ])('rejects malformed lstat fields', async (size, mode) => { + const handle = createHandle() + execSlot(handle).exec = async () => ({ + exitCode: 0, + stdout: `${mode}:${size}\n`, + stderr: '', + }) + await expect(handle.fs.lstat!('/workspace/file')).rejects.toThrow( + 'invalid lstat output', + ) + }) + it('parses a block special file as other without size', async () => { + const handle = createHandle() + execSlot(handle).exec = async () => ({ + exitCode: 0, + stdout: '61b6:0\n', + stderr: '', + }) + await expect(handle.fs.lstat!('/workspace/block')).resolves.toEqual({ + type: 'other', + mode: 0x61b6, + }) + }) + it('parses a fifo as other without size', async () => { + const handle = createHandle() + execSlot(handle).exec = async () => ({ + exitCode: 0, + stdout: '11b6:0\n', + stderr: '', + }) + await expect(handle.fs.lstat!('/workspace/fifo')).resolves.toEqual({ + type: 'other', + mode: 0x11b6, + }) + }) + it('parses an unknown file type as other without size', async () => { + const handle = createHandle() + execSlot(handle).exec = async () => ({ + exitCode: 0, + stdout: '71b6:7\n', + stderr: '', + }) + await expect(handle.fs.lstat!('/workspace/unknown')).resolves.toEqual({ + type: 'other', + mode: 0x71b6, + }) + }) +}) diff --git a/packages/ai-sandbox/README.md b/packages/ai-sandbox/README.md index b62064ec66..48e8e5d257 100644 --- a/packages/ai-sandbox/README.md +++ b/packages/ai-sandbox/README.md @@ -137,6 +137,35 @@ lifecycle: { } ``` +### Portable snapshots + +Use portable snapshots when a later run must rebuild completed files after the +provider sandbox is gone. Configure `withPersistence` first, then pass the same +persistence value and a checkpoint store to `withSandbox`. + +```typescript +import { withPersistence } from '@tanstack/ai-persistence' +import { memorySandboxSnapshots, withSandbox } from '@tanstack/ai-sandbox' + +const snapshots = await memorySandboxSnapshots() + +const middleware = [ + withPersistence(snapshots.persistence), + withSandbox(sandbox, { + snapshots: { + persistence: snapshots.persistence, + checkpoints: snapshots.checkpoints, + }, + }), +] +``` + +A successful terminal run saves regular files, empty directories, saved +conversation data, and thread artifacts. Restore runs only in a new private +sandbox. It never overwrites a live resumed sandbox. Read the +[Portable Sandbox Snapshots guide](https://tanstack.com/ai/latest/docs/sandbox/portable-snapshots) +for storage, safety, and lease details. + ### Secrets Use `createSecrets()` so values stay behind opaque `SecretRef` tokens. They are never written to snapshots, the sandbox store, or event logs. The sandbox layer resolves them onto the live handle at create, resume, and snapshot restore: @@ -178,6 +207,7 @@ Full guides on [tanstack.com/ai](https://tanstack.com/ai/latest/docs/sandbox/ove - [Policy](https://tanstack.com/ai/latest/docs/sandbox/policy) - [Tools](https://tanstack.com/ai/latest/docs/sandbox/tools) (host tool bridge) - [Lifecycle & snapshots](https://tanstack.com/ai/latest/docs/sandbox/lifecycle) +- [Portable sandbox snapshots](https://tanstack.com/ai/latest/docs/sandbox/portable-snapshots) ## Examples diff --git a/packages/ai-sandbox/package.json b/packages/ai-sandbox/package.json index d525e9f38b..964508c0ca 100644 --- a/packages/ai-sandbox/package.json +++ b/packages/ai-sandbox/package.json @@ -70,12 +70,16 @@ "peerDependencies": { "@ngrok/ngrok": "^1.0.0", "@tanstack/ai": "workspace:^", + "@tanstack/ai-persistence": "workspace:^", "vitest": "^4.1.10" }, "peerDependenciesMeta": { "@ngrok/ngrok": { "optional": true }, + "@tanstack/ai-persistence": { + "optional": true + }, "vitest": { "optional": true } @@ -83,6 +87,7 @@ "devDependencies": { "@ngrok/ngrok": "^1.7.0", "@tanstack/ai": "workspace:*", + "@tanstack/ai-persistence": "workspace:*", "@vitest/coverage-v8": "4.0.14", "vitest": "^4.1.10" } diff --git a/packages/ai-sandbox/skills/ai-sandbox/SKILL.md b/packages/ai-sandbox/skills/ai-sandbox/SKILL.md index e9072d64b1..00946b91d6 100644 --- a/packages/ai-sandbox/skills/ai-sandbox/SKILL.md +++ b/packages/ai-sandbox/skills/ai-sandbox/SKILL.md @@ -8,7 +8,13 @@ description: > fileSkill), plugins, instructions → canonical AGENTS.md + symlinks projected per harness; shallow-clone default with depth opt-out; serial/parallel setup callback over a persistent shell; snapshot-after-setup default with - snapshotMaxAge TTL; defineWorkspace (git/setup/scripts/skills/secrets/ + snapshotMaxAge TTL. It also covers portable snapshots after a successful + terminal run with withPersistence before withSandbox and + memorySandboxSnapshots for local examples. It covers named saves with + saveNamedSandboxSnapshot, selected-checkpoint forks with + forkFromSandboxSnapshot, and authorized artifact reads with + resolveSnapshotArtifact. It covers defineWorkspace + (git/setup/scripts/skills/secrets/ instructions/plugins), defineSandboxPolicy (allow/ask/deny), lifecycle/resume, the SandboxHandle (fs/git/process/ports), capability tokens, defineSandbox hooks (onFile/onFileCreate/onFileChange/onFileDelete/onReady/onError/ @@ -185,6 +191,68 @@ lifecycle: { Providers without snapshot support skip the step silently. +### Portable sandbox snapshots + +Portable snapshots keep completed workspace files in application persistence. +They are separate from provider-native bootstrap snapshots. Configure the +middleware in this order, with the same persistence value in both places: + +```typescript +import { withPersistence } from '@tanstack/ai-persistence' +import { memorySandboxSnapshots, withSandbox } from '@tanstack/ai-sandbox' + +const snapshots = await memorySandboxSnapshots() + +const middleware = [ + withPersistence(snapshots.persistence), + withSandbox(sandbox, { + snapshots: { + persistence: snapshots.persistence, + checkpoints: snapshots.checkpoints, + }, + }), +] +``` + +Each successful terminal run saves regular files, empty directories, durable +conversation data, and persisted thread artifacts. A later run restores the +latest checkpoint only into a new private sandbox. A live resumed sandbox is +never overwritten. The default policy excludes `.git`, `node_modules`, `.env*`, +and the workspace projection marker. Resolved secrets are redacted before the +data is stored. Symlinks, executables, and special filesystem entries fail the +capture or restore. Each thread has one writer lease. Pause and detach release +the lease without a partial checkpoint. Blob retention is manual because there +is no automatic garbage collection yet. + +Read `docs/sandbox/portable-snapshots.md` for the full server-only setup and +the restore safety rules. + +For a user-marked workspace state, call `saveNamedSandboxSnapshot` on the +server. It needs `definition`, `threadId`, `runId`, `instances`, `snapshots`, +and a label. It requires a live reusable sandbox. `reuse: 'none'` cannot save a +named checkpoint. + +To branch from a selected checkpoint, call `forkFromSandboxSnapshot` with the +source thread id, source checkpoint id, destination thread id, and `snapshots`. +The store must implement atomic `forkFromCheckpoint`. The destination thread +must be empty. A fork copies the selected snapshot, not the latest snapshot. + +To send a checkpoint artifact, call `resolveSnapshotArtifact` on the server. +First authorize the caller for the supplied thread. The helper checks that the +checkpoint belongs to that thread, then returns its metadata and bytes. It does +not authorize a caller or create an HTTP response. + +For a SQLite checkpoint store, use one transaction for a checkpoint write, its +head update, and every blob reference update. Use one transaction for a fork, +including its copied conversation. A partial transaction breaks snapshot +consistency. + +Snapshot capture supports regular files and empty directories only. It excludes +`.git`, `node_modules`, `.env*`, and the workspace projection marker. It rejects +symlinks, executable files, and special filesystem entries. Restore verifies +the manifest and blobs before it changes a new private sandbox. It never writes +into a live resumed sandbox. + ## Providers - `localProcessSandbox()` — runs on the host (no isolation; dev loop only). diff --git a/packages/ai-sandbox/src/checkpoint-store.ts b/packages/ai-sandbox/src/checkpoint-store.ts new file mode 100644 index 0000000000..022f0eb268 --- /dev/null +++ b/packages/ai-sandbox/src/checkpoint-store.ts @@ -0,0 +1,652 @@ +import type { ModelMessage } from '@tanstack/ai' + +// Compare strings by their UTF-8 bytes so ordering does not depend on locale. +const utf8Encoder = new TextEncoder() + +const compareUtf8Bytes = (left: string, right: string): number => { + const leftBytes = utf8Encoder.encode(left) + const rightBytes = utf8Encoder.encode(right) + const length = Math.min(leftBytes.length, rightBytes.length) + for (let index = 0; index < length; index++) { + const leftByte = leftBytes[index] + const rightByte = rightBytes[index] + if (leftByte !== rightByte) { + return (leftByte ?? 0) - (rightByte ?? 0) + } + } + return leftBytes.length - rightBytes.length +} + +export interface SandboxSnapshotFileEntry { + path: string + kind: 'file' + blobKey: string + size: number +} + +export interface SandboxSnapshotDirectoryEntry { + path: string + kind: 'dir' +} + +export type SandboxSnapshotEntry = + | SandboxSnapshotFileEntry + | SandboxSnapshotDirectoryEntry + +export interface SandboxSnapshotArtifact { + artifactId: string + name: string + mimeType: string + size: number + blobKey: string + createdAt: number +} + +export interface SandboxCheckpoint { + id: string + threadId: string + parentCheckpointId: string | null + createdAt: number + reason: 'automatic' | 'named' | 'fork-root' + label?: string + sourceRunId?: string + files: ReadonlyArray + conversation: ReadonlyArray + artifacts: ReadonlyArray +} + +export interface SandboxCheckpointStore { + get: (id: string) => Promise + list: (threadId: string) => Promise> + getHead: (threadId: string) => Promise + append: (input: { + checkpoint: SandboxCheckpoint + expectedHeadId: string | null + writer: SandboxCheckpointWriter + }) => Promise<{ headId: string }> + deleteHead: (input: { + threadId: string + checkpointId: string + writer: SandboxCheckpointWriter + }) => Promise + acquireWriter: (threadId: string) => Promise + listBlobReferences: () => Promise> + /** Optional atomic fork capability. Stores without this method cannot fork. */ + forkFromCheckpoint?: SandboxCheckpointForkCapability['forkFromCheckpoint'] +} + +export interface SandboxCheckpointForkInput { + sourceThreadId: string + sourceCheckpointId: string + destinationThreadId: string + destinationCheckpointId: string + createdAt: number + writer: SandboxCheckpointWriter +} + +export interface SandboxCheckpointForkCapability { + forkFromCheckpoint: (input: SandboxCheckpointForkInput) => Promise<{ + checkpoint: SandboxCheckpoint + }> +} + +export type ForkCapableSandboxCheckpointStore = SandboxCheckpointStore & + SandboxCheckpointForkCapability + +export function isForkCapableSandboxCheckpointStore( + store: SandboxCheckpointStore, +): store is ForkCapableSandboxCheckpointStore { + return typeof store.forkFromCheckpoint === 'function' +} + +export interface SandboxCheckpointWriter { + threadId: string + ownerToken: string + fence: number +} + +export interface SandboxCheckpointWriterLease extends SandboxCheckpointWriter { + expiresAt: number + renewAfterMs: number + renew: () => Promise<{ expiresAt: number }> + release: () => Promise +} + +export interface SandboxCheckpointStoreOptions { + now?: () => number + leaseDurationMs?: number + renewAfterMs?: number +} + +interface InMemoryCheckpointState { + checkpoints: Map + heads: Map + writers: Map + fences: Map + references: Map +} + +export type SandboxCheckpointErrorCode = + | 'SANDBOX_SNAPSHOT_STALE_HEAD' + | 'SANDBOX_SNAPSHOT_PARENT_MISMATCH' + | 'SANDBOX_SNAPSHOT_DUPLICATE_ID' + | 'SANDBOX_SNAPSHOT_NOT_HEAD' + | 'SANDBOX_SNAPSHOT_WRITER_CONFLICT' + | 'SANDBOX_SNAPSHOT_WRITER_LOST' + | 'SANDBOX_SNAPSHOT_INVALID_ID' + | 'SANDBOX_SNAPSHOT_INVALID_ENTRY' + | 'SANDBOX_SNAPSHOT_CHECKPOINT_NOT_FOUND' + | 'SANDBOX_SNAPSHOT_ATOMIC_FORK_REQUIRED' + | 'SANDBOX_SNAPSHOT_FORK_SOURCE_NOT_FOUND' + | 'SANDBOX_SNAPSHOT_FORK_SOURCE_THREAD_MISMATCH' + | 'SANDBOX_SNAPSHOT_FORK_DESTINATION_NOT_EMPTY' + +export class SandboxCheckpointError extends Error { + readonly code: SandboxCheckpointErrorCode + + constructor(code: SandboxCheckpointErrorCode, message: string) { + super(message) + this.name = 'SandboxCheckpointError' + this.code = code + } +} + +export class SandboxCheckpointConflictError extends SandboxCheckpointError { + constructor(message: string) { + super('SANDBOX_SNAPSHOT_STALE_HEAD', message) + this.name = 'SandboxCheckpointConflictError' + } +} + +export class SandboxCheckpointDuplicateIdError extends SandboxCheckpointError { + constructor(message: string) { + super('SANDBOX_SNAPSHOT_DUPLICATE_ID', message) + this.name = 'SandboxCheckpointDuplicateIdError' + } +} + +export class SandboxCheckpointInvalidIdError extends SandboxCheckpointError { + constructor(message: string) { + super('SANDBOX_SNAPSHOT_INVALID_ID', message) + this.name = 'SandboxCheckpointInvalidIdError' + } +} + +export class SandboxCheckpointInvalidEntryError extends SandboxCheckpointError { + constructor(message: string) { + super('SANDBOX_SNAPSHOT_INVALID_ENTRY', message) + this.name = 'SandboxCheckpointInvalidEntryError' + } +} + +export class SandboxCheckpointParentMismatchError extends SandboxCheckpointError { + constructor(message: string) { + super('SANDBOX_SNAPSHOT_PARENT_MISMATCH', message) + this.name = 'SandboxCheckpointParentMismatchError' + } +} + +export class SandboxCheckpointNotHeadError extends SandboxCheckpointError { + constructor(message: string) { + super('SANDBOX_SNAPSHOT_NOT_HEAD', message) + this.name = 'SandboxCheckpointNotHeadError' + } +} + +export class SandboxCheckpointWriterConflictError extends SandboxCheckpointError { + constructor(message: string) { + super('SANDBOX_SNAPSHOT_WRITER_CONFLICT', message) + this.name = 'SandboxCheckpointWriterConflictError' + } +} + +export class SandboxCheckpointWriterLostError extends SandboxCheckpointError { + constructor(message: string) { + super('SANDBOX_SNAPSHOT_WRITER_LOST', message) + this.name = 'SandboxCheckpointWriterLostError' + } +} + +export function defineSandboxCheckpointStore( + store: SandboxCheckpointStore, +): SandboxCheckpointStore { + return store +} + +function copy(value: T): T { + return structuredClone(value) +} + +function blobKeys(checkpoint: SandboxCheckpoint): Set { + const keys = new Set() + for (const entry of checkpoint.files) { + if (entry.kind === 'file') keys.add(entry.blobKey) + } + for (const artifact of checkpoint.artifacts) keys.add(artifact.blobKey) + return keys +} + +function hasUnpairedSurrogate(value: string): boolean { + for (let index = 0; index < value.length; index++) { + const code = value.charCodeAt(index) + if (code >= 0xd800 && code <= 0xdbff) { + const next = value.charCodeAt(index + 1) + if (Number.isNaN(next) || next < 0xdc00 || next > 0xdfff) return true + index++ + } else if (code >= 0xdc00 && code <= 0xdfff) { + return true + } + } + return false +} + +function assertValidIdentifier( + value: unknown, + label: string, +): asserts value is string { + if ( + typeof value !== 'string' || + value.length === 0 || + hasUnpairedSurrogate(value) + ) { + throw new SandboxCheckpointInvalidIdError( + `${label} must be a non-empty well-formed Unicode string`, + ) + } +} + +function hasOwn(value: object, key: string): boolean { + return Object.prototype.hasOwnProperty.call(value, key) +} + +function validateEntries(checkpoint: SandboxCheckpoint): void { + if (!Array.isArray(checkpoint.files)) { + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint files must be an array', + ) + } + const paths = new Set() + const kinds = new Map() + for (const entry of checkpoint.files as ReadonlyArray) { + if (entry === null || typeof entry !== 'object') { + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint entry must be an object', + ) + } + const candidate = entry as Record + if ( + typeof candidate.path !== 'string' || + candidate.path.length === 0 || + candidate.path.includes('\0') || + candidate.path.startsWith('/') || + candidate.path.startsWith('\\') || + /^[A-Za-z]:([\\/]|$)/.test(candidate.path) || + candidate.path.includes('\\') || + candidate.path + .split('/') + .some((part) => part.length === 0 || part === '.' || part === '..') + ) { + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint entry path must be a normalized workspace-relative path', + ) + } + const path = candidate.path + if (paths.has(path)) { + throw new SandboxCheckpointInvalidEntryError( + `Checkpoint contains duplicate entry path '${path}'`, + ) + } + for ( + let separator = path.indexOf('/'); + separator !== -1; + separator = path.indexOf('/', separator + 1) + ) { + const ancestor = path.slice(0, separator) + if (kinds.get(ancestor) === 'file') { + throw new SandboxCheckpointInvalidEntryError( + `Checkpoint entry '${path}' is beneath file '${ancestor}'`, + ) + } + } + if ( + candidate.kind === 'file' && + Array.from(kinds.keys()).some((other) => other.startsWith(`${path}/`)) + ) { + throw new SandboxCheckpointInvalidEntryError( + `Checkpoint file '${path}' is an ancestor of another entry`, + ) + } + paths.add(path) + if (candidate.kind === 'file') { + if ( + typeof candidate.blobKey !== 'string' || + candidate.blobKey.length === 0 || + hasUnpairedSurrogate(candidate.blobKey) || + !/^sandbox-files\/sha256\/[0-9a-f]{64}$/.test(candidate.blobKey) + ) { + throw new SandboxCheckpointInvalidEntryError( + 'File entries require a non-empty blobKey', + ) + } + if ( + !hasOwn(candidate, 'size') || + typeof candidate.size !== 'number' || + !Number.isSafeInteger(candidate.size) || + candidate.size < 0 + ) { + throw new SandboxCheckpointInvalidEntryError( + 'File entry size must be a non-negative safe integer', + ) + } + } else if (candidate.kind === 'dir') { + if (hasOwn(candidate, 'blobKey') || hasOwn(candidate, 'size')) { + throw new SandboxCheckpointInvalidEntryError( + 'Directory entries cannot contain file fields', + ) + } + } else { + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint entry kind must be file or dir', + ) + } + kinds.set(path, candidate.kind) + } +} + +function validateArtifacts(checkpoint: SandboxCheckpoint): void { + if (!Array.isArray(checkpoint.artifacts)) { + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint artifacts must be an array', + ) + } + for (const artifact of checkpoint.artifacts as ReadonlyArray) { + if (artifact === null || typeof artifact !== 'object') { + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint artifact must be an object', + ) + } + const candidate = artifact as Record + if ( + typeof candidate.artifactId !== 'string' || + candidate.artifactId.length === 0 || + hasUnpairedSurrogate(candidate.artifactId) || + typeof candidate.name !== 'string' || + candidate.name.length === 0 || + typeof candidate.mimeType !== 'string' || + candidate.mimeType.length === 0 || + typeof candidate.blobKey !== 'string' || + candidate.blobKey.length === 0 || + hasUnpairedSurrogate(candidate.blobKey) || + !/^sandbox-artifacts\/sha256\/[0-9a-f]{64}$/.test(candidate.blobKey) || + typeof candidate.size !== 'number' || + !Number.isSafeInteger(candidate.size) || + candidate.size < 0 || + typeof candidate.createdAt !== 'number' || + !Number.isFinite(candidate.createdAt) + ) { + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint artifact has invalid fields', + ) + } + } +} + +export class InMemorySandboxCheckpointStore implements SandboxCheckpointStore { + private readonly state: InMemoryCheckpointState + private readonly now: () => number + private readonly leaseDurationMs: number + private readonly renewAfterMs: number + + constructor(options: SandboxCheckpointStoreOptions = {}) { + const state: InMemoryCheckpointState = { + checkpoints: new Map(), + heads: new Map(), + writers: new Map(), + fences: new Map(), + references: new Map(), + } + this.state = state + this.now = options.now ?? (() => Date.now()) + this.leaseDurationMs = options.leaseDurationMs ?? 120_000 + this.renewAfterMs = options.renewAfterMs ?? 45_000 + if (!Number.isFinite(this.leaseDurationMs) || this.leaseDurationMs <= 0) { + throw new Error('leaseDurationMs must be finite and positive') + } + if ( + !Number.isFinite(this.renewAfterMs) || + this.renewAfterMs <= 0 || + this.renewAfterMs >= this.leaseDurationMs + ) { + throw new Error( + 'renewAfterMs must be finite, positive, and less than leaseDurationMs', + ) + } + } + + async get(id: string): Promise { + assertValidIdentifier(id, 'Checkpoint id') + const checkpoint = this.state.checkpoints.get(id) + return checkpoint ? copy(checkpoint) : null + } + + async list(threadId: string): Promise> { + assertValidIdentifier(threadId, 'Thread id') + return Array.from(this.state.checkpoints.values()) + .filter((checkpoint) => checkpoint.threadId === threadId) + .sort((a, b) => a.createdAt - b.createdAt || compareUtf8Bytes(a.id, b.id)) + .map(copy) + } + + async getHead(threadId: string): Promise { + assertValidIdentifier(threadId, 'Thread id') + return this.state.heads.get(threadId) ?? null + } + + async append(input: { + checkpoint: SandboxCheckpoint + expectedHeadId: string | null + writer: SandboxCheckpointWriter + }): Promise<{ headId: string }> { + const checkpoint = copy(input.checkpoint) + const { expectedHeadId, writer } = input + assertValidIdentifier(checkpoint.threadId, 'Checkpoint thread id') + assertValidIdentifier(writer.threadId, 'Writer thread id') + assertValidIdentifier(checkpoint.id, 'Checkpoint id') + if (expectedHeadId !== null) { + assertValidIdentifier(expectedHeadId, 'Expected head id') + } + if (checkpoint.parentCheckpointId != null) { + assertValidIdentifier( + checkpoint.parentCheckpointId, + 'Parent checkpoint id', + ) + } + if (writer.threadId !== checkpoint.threadId) { + throw new SandboxCheckpointWriterLostError( + 'Checkpoint writer thread does not match checkpoint thread', + ) + } + if (typeof checkpoint.id !== 'string' || checkpoint.id.length === 0) { + throw new SandboxCheckpointInvalidIdError( + 'Checkpoint id must be non-empty', + ) + } + if (hasUnpairedSurrogate(checkpoint.id)) { + throw new SandboxCheckpointInvalidIdError( + 'Checkpoint id must contain valid Unicode', + ) + } + if (hasUnpairedSurrogate(checkpoint.threadId)) { + throw new SandboxCheckpointInvalidIdError( + 'Checkpoint thread id must contain valid Unicode', + ) + } + if ( + typeof checkpoint.createdAt !== 'number' || + !Number.isFinite(checkpoint.createdAt) + ) { + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint createdAt must be a finite number', + ) + } + if (expectedHeadId === '') { + throw new SandboxCheckpointInvalidIdError( + 'Expected head id must be null or non-empty', + ) + } + const parentCheckpointId = checkpoint.parentCheckpointId ?? null + if (parentCheckpointId === '') { + throw new SandboxCheckpointInvalidIdError( + 'Parent checkpoint id must be null or non-empty', + ) + } + if (expectedHeadId !== null && hasUnpairedSurrogate(expectedHeadId)) { + throw new SandboxCheckpointInvalidIdError( + 'Expected head id must contain valid Unicode', + ) + } + if ( + parentCheckpointId !== null && + hasUnpairedSurrogate(parentCheckpointId) + ) { + throw new SandboxCheckpointInvalidIdError( + 'Parent checkpoint id must contain valid Unicode', + ) + } + validateEntries(checkpoint) + validateArtifacts(checkpoint) + + // The caller can re-enter append while its checkpoint is staged. Recheck + // live writer and CAS state immediately before publishing. + this.assertWriter(writer, checkpoint.threadId) + if (this.state.checkpoints.has(checkpoint.id)) { + throw new SandboxCheckpointDuplicateIdError( + `Checkpoint '${checkpoint.id}' already exists`, + ) + } + const actualHeadId = this.state.heads.get(checkpoint.threadId) ?? null + if (actualHeadId !== expectedHeadId) { + throw new SandboxCheckpointConflictError( + `Expected head '${expectedHeadId}', but thread '${checkpoint.threadId}' is at '${actualHeadId}'`, + ) + } + if (parentCheckpointId !== expectedHeadId) { + throw new SandboxCheckpointParentMismatchError( + `Checkpoint '${checkpoint.id}' parent does not match expected head`, + ) + } + const stored = { ...checkpoint, parentCheckpointId } + const keys = blobKeys(stored) + this.state.checkpoints.set(stored.id, stored) + this.state.heads.set(stored.threadId, stored.id) + for (const key of keys) { + this.state.references.set(key, (this.state.references.get(key) ?? 0) + 1) + } + return { headId: stored.id } + } + + async deleteHead(input: { + threadId: string + checkpointId: string + writer: SandboxCheckpointWriter + }): Promise { + const { threadId, checkpointId, writer } = input + assertValidIdentifier(threadId, 'Thread id') + assertValidIdentifier(checkpointId, 'Checkpoint id') + assertValidIdentifier(writer.threadId, 'Writer thread id') + if (writer.threadId !== threadId) { + throw new SandboxCheckpointWriterLostError( + 'Checkpoint writer thread does not match operation thread', + ) + } + this.assertWriter(writer, threadId) + const headId = this.state.heads.get(threadId) ?? null + if (headId !== checkpointId) { + throw new SandboxCheckpointNotHeadError( + `Checkpoint '${checkpointId}' is not the current head of thread '${threadId}'`, + ) + } + const checkpoint = this.state.checkpoints.get(checkpointId) + if (!checkpoint) { + throw new SandboxCheckpointNotHeadError( + `Checkpoint '${checkpointId}' does not exist`, + ) + } + this.state.checkpoints.delete(checkpointId) + if (checkpoint.parentCheckpointId) { + this.state.heads.set(threadId, checkpoint.parentCheckpointId) + } else { + this.state.heads.delete(threadId) + } + for (const key of blobKeys(checkpoint)) { + const references = (this.state.references.get(key) ?? 0) - 1 + if (references > 0) this.state.references.set(key, references) + else this.state.references.delete(key) + } + } + + async acquireWriter(threadId: string): Promise { + assertValidIdentifier(threadId, 'Thread id') + const current = this.state.writers.get(threadId) + if (current && current.expiresAt > this.now()) { + throw new SandboxCheckpointWriterConflictError( + `Thread '${threadId}' already has an active checkpoint writer`, + ) + } + const fence = (this.state.fences.get(threadId) ?? 0) + 1 + this.state.fences.set(threadId, fence) + const ownerToken = globalThis.crypto.randomUUID() + const lease = { + threadId, + ownerToken, + fence, + expiresAt: this.now() + this.leaseDurationMs, + } + this.state.writers.set(threadId, lease) + return { + ...lease, + get expiresAt() { + return lease.expiresAt + }, + renewAfterMs: this.renewAfterMs, + renew: async () => { + this.assertWriter(lease, threadId) + lease.expiresAt = this.now() + this.leaseDurationMs + return { expiresAt: lease.expiresAt } + }, + release: async () => { + const currentLease = this.state.writers.get(threadId) + if ( + currentLease?.ownerToken === ownerToken && + currentLease.fence === fence + ) + this.state.writers.delete(threadId) + }, + } + } + + private assertWriter( + writer: SandboxCheckpointWriter, + threadId: string, + ): void { + const current = this.state.writers.get(threadId) + if ( + !current || + current.ownerToken !== writer.ownerToken || + current.fence !== writer.fence || + current.expiresAt <= this.now() + ) { + throw new SandboxCheckpointWriterLostError( + `Checkpoint writer lease for thread '${threadId}' is no longer current`, + ) + } + } + + async listBlobReferences(): Promise< + Array<{ key: string; references: number }> + > { + return Array.from(this.state.references.entries()) + .sort(([a], [b]) => compareUtf8Bytes(a, b)) + .map(([key, references]) => ({ key, references })) + } +} diff --git a/packages/ai-sandbox/src/contracts.ts b/packages/ai-sandbox/src/contracts.ts index 7c0dc55a17..0b5c028dd9 100644 --- a/packages/ai-sandbox/src/contracts.ts +++ b/packages/ai-sandbox/src/contracts.ts @@ -117,6 +117,11 @@ export interface SandboxFs { remove: (path: string) => Promise rename: (from: string, to: string) => Promise exists: (path: string) => Promise + /** + * Optional metadata lookup. Implementations must not follow symlinks. + * Returns undefined only for a confirmed missing path. All other errors reject. + */ + lstat?: (path: string) => Promise /** Optional — present only when `capabilities.fs` providers advertise watch. */ watch?: ( path: string, @@ -124,6 +129,13 @@ export interface SandboxFs { ) => Promise<{ stop: () => Promise }> } +export type SandboxFsStat = + // `mode` is the complete POSIX mode value, including the file-type bits. + | { type: 'file'; mode: number; size: number } + | { type: 'dir'; mode: number } + | { type: 'symlink'; mode: number } + | { type: 'other'; mode: number } + /** * Uniform git surface. Implementations either delegate to the provider's * native git (when advertised) or desugar to `process.exec("git …")`, so the diff --git a/packages/ai-sandbox/src/index.ts b/packages/ai-sandbox/src/index.ts index 99a4dfca79..79050723d0 100644 --- a/packages/ai-sandbox/src/index.ts +++ b/packages/ai-sandbox/src/index.ts @@ -27,6 +27,53 @@ export type { SandboxInstanceRecord, } from './instance-store' +// Portable immutable sandbox checkpoint metadata. +export { + SandboxCheckpointError, + SandboxCheckpointConflictError, + SandboxCheckpointDuplicateIdError, + SandboxCheckpointInvalidIdError, + SandboxCheckpointInvalidEntryError, + SandboxCheckpointParentMismatchError, + SandboxCheckpointNotHeadError, + SandboxCheckpointWriterConflictError, + SandboxCheckpointWriterLostError, + isForkCapableSandboxCheckpointStore, + InMemorySandboxCheckpointStore, + defineSandboxCheckpointStore, +} from './checkpoint-store' + +export type { + SandboxCheckpoint, + SandboxCheckpointStore, + SandboxSnapshotEntry, + SandboxSnapshotFileEntry, + SandboxSnapshotDirectoryEntry, + SandboxSnapshotArtifact, + SandboxCheckpointErrorCode, + SandboxCheckpointWriter, + SandboxCheckpointWriterLease, + SandboxCheckpointStoreOptions, + SandboxCheckpointForkInput, + SandboxCheckpointForkCapability, + ForkCapableSandboxCheckpointStore, +} from './checkpoint-store' + +// File snapshot policy used by provider snapshot create/restore inputs. +export { SandboxSnapshotError } from './snapshots' +export type { + SandboxSnapshotErrorCode, + SandboxSnapshotPolicy, +} from './snapshots' +export { memorySandboxSnapshots } from './memory-snapshots' +export type { MemorySandboxSnapshots } from './memory-snapshots' +export { + forkFromSandboxSnapshot, + resolveSnapshotArtifact, + saveNamedSandboxSnapshot, +} from './snapshot-operations' +export type { SandboxSnapshots } from './snapshot-operations' + // Workspace projection capability (provided by withSandbox, consumed by harness adapters) export { ProjectionCapability, @@ -102,6 +149,7 @@ export type { SandboxHandle, SandboxCapabilities, SandboxFs, + SandboxFsStat, SandboxGit, SandboxProcess, SandboxPorts, diff --git a/packages/ai-sandbox/src/memory-snapshot-types.ts b/packages/ai-sandbox/src/memory-snapshot-types.ts new file mode 100644 index 0000000000..f9c530bfbb --- /dev/null +++ b/packages/ai-sandbox/src/memory-snapshot-types.ts @@ -0,0 +1,167 @@ +import type { + ModelMessage, + RunRecord, + RunStatus, + RunStore, + PersistedArtifactRef, + TokenUsage, +} from '@tanstack/ai' + +export interface MemoryMessageStore { + loadThread: (threadId: string) => Promise> + saveThread: (threadId: string, messages: Array) => Promise +} + +export type MemoryRunRecord = RunRecord + +export type MemoryRunStore = RunStore + +export interface MemoryGenerationRunRecord { + runId: string + threadId: string + activity: string + provider: string + model: string + status: RunStatus + startedAt: number + finishedAt?: number + error?: { message: string; code?: string } + result?: unknown + artifacts?: Array + usage?: TokenUsage +} + +export interface MemoryGenerationRunStore { + createOrResume: ( + input: Pick< + MemoryGenerationRunRecord, + 'runId' | 'threadId' | 'activity' | 'provider' | 'model' | 'startedAt' + > & { status?: RunStatus }, + ) => Promise + update: ( + runId: string, + patch: Partial< + Pick< + MemoryGenerationRunRecord, + 'status' | 'finishedAt' | 'error' | 'result' | 'artifacts' | 'usage' + > + >, + ) => Promise + get: (runId: string) => Promise + findLatestForThread: ( + threadId: string, + ) => Promise +} + +export interface MemoryInterruptRecord { + interruptId: string + runId: string + threadId: string + status: 'pending' | 'resolved' | 'cancelled' + requestedAt: number + resolvedAt?: number + payload: Record + response?: unknown +} + +export interface MemoryInterruptStore { + create: ( + record: Omit, + ) => Promise + resolve: (interruptId: string, response?: unknown) => Promise + cancel: (interruptId: string) => Promise + get: (interruptId: string) => Promise + list: (threadId: string) => Promise> + listPending: (threadId: string) => Promise> + listByRun: (runId: string) => Promise> + listPendingByRun: (runId: string) => Promise> +} + +export interface MemoryMetadataStore { + get: (namespace: string, key: string) => Promise + set: (namespace: string, key: string, value: unknown) => Promise + delete: (namespace: string, key: string) => Promise +} + +export interface MemoryArtifactRecord { + artifactId: string + runId: string + threadId: string + blobKey?: string + name: string + mimeType: string + size: number + sourceUrl?: string + createdAt: number +} + +export interface MemoryArtifactStore { + save: (record: MemoryArtifactRecord) => Promise + get: (artifactId: string) => Promise + list: (runId: string) => Promise> + listForThread: (threadId: string) => Promise> + delete: (artifactId: string) => Promise + deleteForRun: (runId: string) => Promise +} + +export type MemoryBlobBody = + | ReadableStream + | ArrayBuffer + | ArrayBufferView + | string + | Blob +export interface MemoryBlobRecord { + key: string + size?: number + etag?: string + contentType?: string + customMetadata?: Record + createdAt?: number + updatedAt?: number +} +export interface MemoryBlobStore { + put: ( + key: string, + body: MemoryBlobBody, + options?: { + contentType?: string + customMetadata?: Record + expectedLength?: number + }, + ) => Promise + get: ( + key: string, + options?: { range?: { offset: number; length?: number } }, + ) => Promise< + | (MemoryBlobRecord & { + arrayBuffer: () => Promise + text: () => Promise + body?: ReadableStream + range?: { offset: number; length: number } + }) + | null + > + head: (key: string) => Promise + delete: (key: string) => Promise + list: (options?: { + prefix?: string + cursor?: string + limit?: number + }) => Promise<{ + objects: Array + cursor?: string + truncated?: boolean + }> +} + +export interface MemorySnapshotPersistence { + stores: { + messages: MemoryMessageStore + runs: MemoryRunStore + generationRuns: MemoryGenerationRunStore + interrupts: MemoryInterruptStore + metadata: MemoryMetadataStore + artifacts: MemoryArtifactStore + blobs: MemoryBlobStore + } +} diff --git a/packages/ai-sandbox/src/memory-snapshots.ts b/packages/ai-sandbox/src/memory-snapshots.ts new file mode 100644 index 0000000000..aaf3a7553a --- /dev/null +++ b/packages/ai-sandbox/src/memory-snapshots.ts @@ -0,0 +1,913 @@ +import type { ModelMessage } from '@tanstack/ai' +import type { + MemoryArtifactRecord as ArtifactRecord, + MemoryBlobBody as BlobBody, + MemoryBlobRecord as BlobRecord, + MemoryGenerationRunRecord as GenerationRunRecord, + MemoryInterruptRecord as InterruptRecord, + MemoryRunRecord as RunRecord, + MemorySnapshotPersistence, +} from './memory-snapshot-types' +import { + SandboxCheckpointConflictError, + SandboxCheckpointDuplicateIdError, + SandboxCheckpointError, + SandboxCheckpointInvalidEntryError, + SandboxCheckpointInvalidIdError, + SandboxCheckpointNotHeadError, + SandboxCheckpointParentMismatchError, + SandboxCheckpointWriterConflictError, + SandboxCheckpointWriterLostError, +} from './checkpoint-store' +import type { + ForkCapableSandboxCheckpointStore, + SandboxCheckpoint, + SandboxCheckpointWriter, + SandboxCheckpointWriterLease, + SandboxCheckpointForkInput, +} from './checkpoint-store' + +type BlobGetOptions = { range?: { offset: number; length?: number } } + +function resolveBlobRange( + size: number, + range: { offset: number; length?: number }, +): { offset: number; length: number } { + if ( + !Number.isInteger(range.offset) || + range.offset < 0 || + range.offset >= size + ) { + throw new RangeError( + `Blob range offset ${range.offset} is outside the object (size ${size}).`, + ) + } + const remaining = size - range.offset + if (range.length === undefined) { + return { offset: range.offset, length: remaining } + } + if (!Number.isInteger(range.length) || range.length < 0) { + throw new RangeError(`Blob range length ${range.length} is not valid.`) + } + return { + offset: range.offset, + length: Math.min(range.length, remaining), + } +} + +export interface MemorySandboxSnapshots { + persistence: MemorySnapshotPersistence + checkpoints: ForkCapableSandboxCheckpointStore +} + +const encoder = new TextEncoder() +const compare = (a: string, b: string) => { + const left = encoder.encode(a) + const right = encoder.encode(b) + for (let i = 0; i < Math.min(left.length, right.length); i++) { + const leftByte = left[i] + const rightByte = right[i] + if (leftByte !== rightByte) return (leftByte ?? 0) - (rightByte ?? 0) + } + return left.length - right.length +} +const clone = (value: T): T => structuredClone(value) + +interface MemoryCheckpointState { + checkpoints: Map + heads: Map + writers: Map + fences: Map + references: Map +} + +interface MemorySnapshotState extends MemoryCheckpointState { + messages: Map> + runs: Map + generations: Map + interrupts: Map + metadata: Map> + artifacts: Map + blobs: Map +} + +function hasUnpairedSurrogate(value: string): boolean { + for (let index = 0; index < value.length; index++) { + const code = value.charCodeAt(index) + if (code >= 0xd800 && code <= 0xdbff) { + const next = value.charCodeAt(index + 1) + if (Number.isNaN(next) || next < 0xdc00 || next > 0xdfff) return true + index++ + } else if (code >= 0xdc00 && code <= 0xdfff) { + return true + } + } + return false +} + +function assertValidIdentifier( + value: unknown, + label: string, +): asserts value is string { + if ( + typeof value !== 'string' || + value.length === 0 || + hasUnpairedSurrogate(value) + ) { + throw new SandboxCheckpointInvalidIdError( + `${label} must be a non-empty well-formed Unicode string`, + ) + } +} + +function hasOwn(value: object, key: string): boolean { + return Object.prototype.hasOwnProperty.call(value, key) +} + +function validateEntries(checkpoint: SandboxCheckpoint): void { + if (!Array.isArray(checkpoint.files)) { + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint files must be an array', + ) + } + const paths = new Set() + const kinds = new Map() + for (const entry of checkpoint.files as ReadonlyArray) { + if (entry === null || typeof entry !== 'object') { + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint entry must be an object', + ) + } + const candidate = entry as Record + if ( + typeof candidate.path !== 'string' || + candidate.path.length === 0 || + candidate.path.includes('\0') || + candidate.path.startsWith('/') || + candidate.path.startsWith('\\') || + /^[A-Za-z]:([\\/]|$)/.test(candidate.path) || + candidate.path.includes('\\') || + candidate.path + .split('/') + .some((part) => part.length === 0 || part === '.' || part === '..') + ) { + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint entry path must be a normalized workspace-relative path', + ) + } + const path = candidate.path + if (paths.has(path)) { + throw new SandboxCheckpointInvalidEntryError( + `Checkpoint contains duplicate entry path '${path}'`, + ) + } + for ( + let separator = path.indexOf('/'); + separator !== -1; + separator = path.indexOf('/', separator + 1) + ) { + const ancestor = path.slice(0, separator) + if (kinds.get(ancestor) === 'file') { + throw new SandboxCheckpointInvalidEntryError( + `Checkpoint entry '${path}' is beneath file '${ancestor}'`, + ) + } + } + if ( + candidate.kind === 'file' && + Array.from(kinds.keys()).some((other) => other.startsWith(`${path}/`)) + ) { + throw new SandboxCheckpointInvalidEntryError( + `Checkpoint file '${path}' is an ancestor of another entry`, + ) + } + paths.add(path) + if (candidate.kind === 'file') { + if ( + typeof candidate.blobKey !== 'string' || + candidate.blobKey.length === 0 || + hasUnpairedSurrogate(candidate.blobKey) || + !/^sandbox-files\/sha256\/[0-9a-f]{64}$/.test(candidate.blobKey) + ) { + throw new SandboxCheckpointInvalidEntryError( + 'File entries require a valid content-addressed blobKey', + ) + } + if ( + !hasOwn(candidate, 'size') || + typeof candidate.size !== 'number' || + !Number.isSafeInteger(candidate.size) || + candidate.size < 0 + ) { + throw new SandboxCheckpointInvalidEntryError( + 'File entry size must be a non-negative safe integer', + ) + } + } else if (candidate.kind === 'dir') { + if (hasOwn(candidate, 'blobKey') || hasOwn(candidate, 'size')) { + throw new SandboxCheckpointInvalidEntryError( + 'Directory entries cannot contain file fields', + ) + } + } else { + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint entry kind must be file or dir', + ) + } + kinds.set(path, candidate.kind) + } +} + +function validateArtifacts(checkpoint: SandboxCheckpoint): void { + if (!Array.isArray(checkpoint.artifacts)) { + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint artifacts must be an array', + ) + } + for (const artifact of checkpoint.artifacts as ReadonlyArray) { + if (artifact === null || typeof artifact !== 'object') { + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint artifact must be an object', + ) + } + const candidate = artifact as Record + if ( + typeof candidate.artifactId !== 'string' || + candidate.artifactId.length === 0 || + hasUnpairedSurrogate(candidate.artifactId) || + typeof candidate.name !== 'string' || + candidate.name.length === 0 || + typeof candidate.mimeType !== 'string' || + candidate.mimeType.length === 0 || + typeof candidate.blobKey !== 'string' || + candidate.blobKey.length === 0 || + hasUnpairedSurrogate(candidate.blobKey) || + !/^sandbox-artifacts\/sha256\/[0-9a-f]{64}$/.test(candidate.blobKey) || + typeof candidate.size !== 'number' || + !Number.isSafeInteger(candidate.size) || + candidate.size < 0 || + typeof candidate.createdAt !== 'number' || + !Number.isFinite(candidate.createdAt) + ) { + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint artifact has invalid fields', + ) + } + } +} + +function validateCheckpoint(checkpoint: SandboxCheckpoint): void { + assertValidIdentifier(checkpoint.id, 'Checkpoint id') + assertValidIdentifier(checkpoint.threadId, 'Checkpoint thread id') + if (checkpoint.parentCheckpointId !== null) { + assertValidIdentifier(checkpoint.parentCheckpointId, 'Parent checkpoint id') + } + if (!Number.isFinite(checkpoint.createdAt)) { + throw new SandboxCheckpointInvalidEntryError( + 'Checkpoint createdAt must be a finite number', + ) + } + validateEntries(checkpoint) + validateArtifacts(checkpoint) +} + +function blobKeys(checkpoint: SandboxCheckpoint): Set { + const keys = new Set() + for (const entry of checkpoint.files) { + if (entry.kind === 'file') keys.add(entry.blobKey) + } + for (const artifact of checkpoint.artifacts) keys.add(artifact.blobKey) + return keys +} + +class MemorySnapshotCheckpointStore implements ForkCapableSandboxCheckpointStore { + private readonly now = () => Date.now() + private readonly leaseDurationMs = 120_000 + private readonly renewAfterMs = 45_000 + + constructor(private readonly state: MemorySnapshotState) {} + + async get(id: string): Promise { + assertValidIdentifier(id, 'Checkpoint id') + const checkpoint = this.state.checkpoints.get(id) + return checkpoint ? clone(checkpoint) : null + } + + async list(threadId: string): Promise> { + assertValidIdentifier(threadId, 'Thread id') + return [...this.state.checkpoints.values()] + .filter((checkpoint) => checkpoint.threadId === threadId) + .sort((a, b) => a.createdAt - b.createdAt || compare(a.id, b.id)) + .map(clone) + } + + async getHead(threadId: string): Promise { + assertValidIdentifier(threadId, 'Thread id') + return this.state.heads.get(threadId) ?? null + } + + async append(input: { + checkpoint: SandboxCheckpoint + expectedHeadId: string | null + writer: SandboxCheckpointWriter + }): Promise<{ headId: string }> { + const checkpoint = clone(input.checkpoint) + const { expectedHeadId, writer } = input + assertValidIdentifier(checkpoint.id, 'Checkpoint id') + assertValidIdentifier(checkpoint.threadId, 'Checkpoint thread id') + assertValidIdentifier(writer.threadId, 'Writer thread id') + if (expectedHeadId !== null) { + assertValidIdentifier(expectedHeadId, 'Expected head id') + } + if (checkpoint.parentCheckpointId != null) { + assertValidIdentifier( + checkpoint.parentCheckpointId, + 'Parent checkpoint id', + ) + } + if (writer.threadId !== checkpoint.threadId) { + throw new SandboxCheckpointWriterLostError( + 'Checkpoint writer thread does not match checkpoint thread', + ) + } + validateCheckpoint(checkpoint) + + // Re-check live state after staging because cloning caller data may + // re-enter append and advance this thread's head. + this.assertWriter(writer, checkpoint.threadId) + if (this.state.checkpoints.has(checkpoint.id)) { + throw new SandboxCheckpointDuplicateIdError( + `Checkpoint '${checkpoint.id}' already exists`, + ) + } + const actualHeadId = this.state.heads.get(checkpoint.threadId) ?? null + if (actualHeadId !== expectedHeadId) { + throw new SandboxCheckpointConflictError( + `Expected head '${expectedHeadId}', but thread '${checkpoint.threadId}' is at '${actualHeadId}'`, + ) + } + const parentCheckpointId = checkpoint.parentCheckpointId ?? null + if (parentCheckpointId !== expectedHeadId) { + throw new SandboxCheckpointParentMismatchError( + `Checkpoint '${checkpoint.id}' parent does not match expected head`, + ) + } + const stored = { ...checkpoint, parentCheckpointId } + const keys = blobKeys(stored) + this.state.checkpoints.set(stored.id, stored) + this.state.heads.set(stored.threadId, stored.id) + for (const key of keys) { + this.state.references.set(key, (this.state.references.get(key) ?? 0) + 1) + } + return { headId: stored.id } + } + + async deleteHead(input: { + threadId: string + checkpointId: string + writer: SandboxCheckpointWriter + }): Promise { + const { threadId, checkpointId, writer } = input + assertValidIdentifier(threadId, 'Thread id') + assertValidIdentifier(checkpointId, 'Checkpoint id') + assertValidIdentifier(writer.threadId, 'Writer thread id') + if (writer.threadId !== threadId) { + throw new SandboxCheckpointWriterLostError( + 'Checkpoint writer thread does not match operation thread', + ) + } + this.assertWriter(writer, threadId) + if ((this.state.heads.get(threadId) ?? null) !== checkpointId) { + throw new SandboxCheckpointNotHeadError( + `Checkpoint '${checkpointId}' is not the current head of thread '${threadId}'`, + ) + } + const checkpoint = this.state.checkpoints.get(checkpointId) + if (!checkpoint) { + throw new SandboxCheckpointNotHeadError( + `Checkpoint '${checkpointId}' does not exist`, + ) + } + this.state.checkpoints.delete(checkpointId) + if (checkpoint.parentCheckpointId) { + this.state.heads.set(threadId, checkpoint.parentCheckpointId) + } else { + this.state.heads.delete(threadId) + } + for (const key of blobKeys(checkpoint)) { + const references = (this.state.references.get(key) ?? 0) - 1 + if (references > 0) this.state.references.set(key, references) + else this.state.references.delete(key) + } + } + + async acquireWriter(threadId: string): Promise { + assertValidIdentifier(threadId, 'Thread id') + const current = this.state.writers.get(threadId) + if (current && current.expiresAt > this.now()) { + throw new SandboxCheckpointWriterConflictError( + `Thread '${threadId}' already has an active checkpoint writer`, + ) + } + const fence = (this.state.fences.get(threadId) ?? 0) + 1 + this.state.fences.set(threadId, fence) + const ownerToken = globalThis.crypto.randomUUID() + const lease = { + threadId, + ownerToken, + fence, + expiresAt: this.now() + this.leaseDurationMs, + } + this.state.writers.set(threadId, lease) + return { + ...lease, + get expiresAt() { + return lease.expiresAt + }, + renewAfterMs: this.renewAfterMs, + renew: async () => { + this.assertWriter(lease, threadId) + lease.expiresAt = this.now() + this.leaseDurationMs + return { expiresAt: lease.expiresAt } + }, + release: async () => { + const currentLease = this.state.writers.get(threadId) + if ( + currentLease?.ownerToken === ownerToken && + currentLease.fence === fence + ) { + this.state.writers.delete(threadId) + } + }, + } + } + + async listBlobReferences(): Promise< + Array<{ key: string; references: number }> + > { + return [...this.state.references.entries()] + .sort(([a], [b]) => compare(a, b)) + .map(([key, references]) => ({ key, references })) + } + + async forkFromCheckpoint( + input: SandboxCheckpointForkInput, + ): Promise<{ checkpoint: SandboxCheckpoint }> { + const sourceThreadId = input.sourceThreadId + const sourceCheckpointId = input.sourceCheckpointId + const destinationThreadId = input.destinationThreadId + const destinationCheckpointId = input.destinationCheckpointId + const createdAt = input.createdAt + const suppliedWriter = input.writer + const writer: SandboxCheckpointWriter = { + threadId: suppliedWriter.threadId, + ownerToken: suppliedWriter.ownerToken, + fence: suppliedWriter.fence, + } + + assertValidIdentifier(sourceThreadId, 'Source thread id') + assertValidIdentifier(sourceCheckpointId, 'Source checkpoint id') + assertValidIdentifier(destinationThreadId, 'Destination thread id') + assertValidIdentifier(destinationCheckpointId, 'Destination checkpoint id') + assertValidIdentifier(writer.threadId, 'Writer thread id') + if (!Number.isFinite(createdAt)) { + throw new SandboxCheckpointInvalidEntryError( + 'Fork checkpoint createdAt must be a finite number', + ) + } + if (sourceThreadId === destinationThreadId) { + throw new SandboxCheckpointError( + 'SANDBOX_SNAPSHOT_FORK_SOURCE_THREAD_MISMATCH', + 'Source and destination threads must differ', + ) + } + const source = this.state.checkpoints.get(sourceCheckpointId) + if (!source) { + throw new SandboxCheckpointError( + 'SANDBOX_SNAPSHOT_FORK_SOURCE_NOT_FOUND', + 'Source checkpoint was not found', + ) + } + if (source.threadId !== sourceThreadId) { + throw new SandboxCheckpointError( + 'SANDBOX_SNAPSHOT_FORK_SOURCE_THREAD_MISMATCH', + 'Source checkpoint belongs to another thread', + ) + } + if (writer.threadId !== destinationThreadId) { + throw new SandboxCheckpointWriterLostError( + 'Checkpoint writer thread does not match destination thread', + ) + } + this.assertWriter(writer, destinationThreadId) + this.assertDestinationEmpty(destinationThreadId, destinationCheckpointId) + + const stagedCheckpoint: SandboxCheckpoint = clone({ + id: destinationCheckpointId, + threadId: destinationThreadId, + parentCheckpointId: null, + createdAt, + reason: 'fork-root', + files: source.files, + conversation: source.conversation, + artifacts: source.artifacts, + }) + validateCheckpoint(stagedCheckpoint) + const stagedTranscript = clone([...stagedCheckpoint.conversation]) + const result = { checkpoint: clone(stagedCheckpoint) } + const stagedReferences = [...blobKeys(stagedCheckpoint)].map((key) => ({ + key, + references: (this.state.references.get(key) ?? 0) + 1, + })) + + // Cloning can invoke user-defined getters. Revalidate immediately before + // the synchronous publication block so a reentrant save is never lost. + this.assertWriter(writer, destinationThreadId) + this.assertDestinationEmpty(destinationThreadId, destinationCheckpointId) + + this.state.messages.set(stagedCheckpoint.threadId, stagedTranscript) + this.state.checkpoints.set(stagedCheckpoint.id, stagedCheckpoint) + this.state.heads.set(stagedCheckpoint.threadId, stagedCheckpoint.id) + for (const reference of stagedReferences) { + this.state.references.set(reference.key, reference.references) + } + return result + } + + private assertDestinationEmpty( + destinationThreadId: string, + destinationCheckpointId: string, + ): void { + if ( + this.state.messages.has(destinationThreadId) || + [...this.state.runs.values()].some( + (value) => value.threadId === destinationThreadId, + ) || + [...this.state.generations.values()].some( + (value) => value.threadId === destinationThreadId, + ) || + [...this.state.interrupts.values()].some( + (value) => value.threadId === destinationThreadId, + ) || + [...this.state.artifacts.values()].some( + (value) => value.threadId === destinationThreadId, + ) || + [...this.state.checkpoints.values()].some( + (value) => value.threadId === destinationThreadId, + ) || + this.state.heads.has(destinationThreadId) || + this.state.checkpoints.has(destinationCheckpointId) + ) { + throw new SandboxCheckpointError( + 'SANDBOX_SNAPSHOT_FORK_DESTINATION_NOT_EMPTY', + 'Destination thread is not empty', + ) + } + } + + private assertWriter( + writer: SandboxCheckpointWriter, + threadId: string, + ): void { + const current = this.state.writers.get(threadId) + if ( + !current || + current.ownerToken !== writer.ownerToken || + current.fence !== writer.fence || + current.expiresAt <= this.now() + ) { + throw new SandboxCheckpointWriterLostError( + `Checkpoint writer lease for thread '${threadId}' is no longer current`, + ) + } + } +} + +async function bodyBytes(body: BlobBody): Promise { + if (typeof body === 'string') return encoder.encode(body) + if (body instanceof ArrayBuffer) return new Uint8Array(body.slice(0)) + if (ArrayBuffer.isView(body)) + return new Uint8Array( + body.buffer.slice(body.byteOffset, body.byteOffset + body.byteLength), + ) + if (typeof Blob !== 'undefined' && body instanceof Blob) + return new Uint8Array(await body.arrayBuffer()) + if (typeof ReadableStream !== 'undefined' && body instanceof ReadableStream) { + const reader = body.getReader() + const parts: Array = [] + try { + for (;;) { + const next = await reader.read() + if (next.done) break + parts.push(new Uint8Array(next.value)) + } + } finally { + reader.releaseLock() + } + const result = new Uint8Array( + parts.reduce((total, part) => total + part.byteLength, 0), + ) + let offset = 0 + for (const part of parts) { + result.set(part, offset) + offset += part.byteLength + } + return result + } + throw new TypeError('Unsupported blob body.') +} + +export async function memorySandboxSnapshots(): Promise { + return createMemorySandboxSnapshots() +} + +async function createMemorySandboxSnapshots(): Promise { + const messages = new Map>() + const runs = new Map() + const generations = new Map() + const interrupts = new Map() + const metadata = new Map>() + const artifacts = new Map() + const blobs = new Map() + const state: MemorySnapshotState = { + messages, + runs, + generations, + interrupts, + metadata, + artifacts, + blobs, + checkpoints: new Map(), + heads: new Map(), + writers: new Map(), + fences: new Map(), + references: new Map(), + } + let etag = 0 + const persistence: MemorySnapshotPersistence = { + stores: { + messages: { + loadThread: async (threadId: string) => + messages.get(threadId)?.slice() ?? [], + saveThread: async (threadId: string, value: Array) => { + messages.set(threadId, value.slice()) + }, + }, + runs: { + createOrResume: async (input: { + runId: string + threadId: string + status?: RunRecord['status'] + startedAt: number + }) => { + const existing = runs.get(input.runId) + if (existing) return existing + const record: RunRecord = { + ...input, + status: input.status ?? 'running', + } + runs.set(record.runId, record) + return record + }, + update: async (runId: string, patch: Partial) => { + const value = runs.get(runId) + if (value) runs.set(runId, { ...value, ...patch }) + }, + get: async (runId: string) => runs.get(runId) ?? null, + findActiveRun: async (threadId: string) => + [...runs.values()] + .filter( + (run) => run.threadId === threadId && run.status === 'running', + ) + .sort((a, b) => b.startedAt - a.startedAt)[0] ?? null, + listByThread: async (threadId: string) => + [...runs.values()] + .filter((run) => run.threadId === threadId) + .sort((a, b) => a.startedAt - b.startedAt), + listReclaimable: async (input: { now: number; ttlMs: number }) => + [...runs.values()].filter( + (run) => + run.status === 'running' && + run.detachedSince !== undefined && + run.detachedSince <= input.now - input.ttlMs, + ), + }, + generationRuns: { + createOrResume: async ( + input: Pick< + GenerationRunRecord, + | 'runId' + | 'threadId' + | 'activity' + | 'provider' + | 'model' + | 'startedAt' + > & { status?: GenerationRunRecord['status'] }, + ) => { + const value = generations.get(input.runId) ?? { + ...input, + status: input.status ?? 'running', + } + generations.set(input.runId, value) + return value + }, + update: async (runId: string, patch: Partial) => { + const value = generations.get(runId) + if (value) generations.set(runId, { ...value, ...patch }) + }, + get: async (runId: string) => generations.get(runId) ?? null, + findLatestForThread: async (threadId: string) => + [...generations.values()] + .filter((run) => run.threadId === threadId) + .sort((a, b) => b.startedAt - a.startedAt)[0] ?? null, + }, + interrupts: { + create: async ( + record: Omit, + ) => { + if (!interrupts.has(record.interruptId)) + interrupts.set(record.interruptId, { ...record, status: 'pending' }) + }, + resolve: async (id: string, response?: unknown) => { + const value = interrupts.get(id) + if (value) + interrupts.set(id, { + ...value, + status: 'resolved', + resolvedAt: Date.now(), + response, + }) + }, + cancel: async (id: string) => { + const value = interrupts.get(id) + if (value) + interrupts.set(id, { + ...value, + status: 'cancelled', + resolvedAt: Date.now(), + }) + }, + get: async (id: string) => interrupts.get(id) ?? null, + list: async (threadId: string) => + [...interrupts.values()] + .filter((value) => value.threadId === threadId) + .sort((a, b) => a.requestedAt - b.requestedAt), + listPending: async (threadId: string) => + [...interrupts.values()] + .filter( + (value) => + value.threadId === threadId && value.status === 'pending', + ) + .sort((a, b) => a.requestedAt - b.requestedAt), + listByRun: async (runId: string) => + [...interrupts.values()] + .filter((value) => value.runId === runId) + .sort((a, b) => a.requestedAt - b.requestedAt), + listPendingByRun: async (runId: string) => + [...interrupts.values()] + .filter( + (value) => value.runId === runId && value.status === 'pending', + ) + .sort((a, b) => a.requestedAt - b.requestedAt), + }, + metadata: { + get: async (namespace: string, key: string) => { + const bucket = metadata.get(namespace) + return bucket?.has(key) ? bucket.get(key) : null + }, + set: async (namespace: string, key: string, value: unknown) => { + let bucket = metadata.get(namespace) + if (!bucket) { + bucket = new Map() + metadata.set(namespace, bucket) + } + bucket.set(key, value) + }, + delete: async (namespace: string, key: string) => { + metadata.get(namespace)?.delete(key) + }, + }, + artifacts: { + save: async (value: ArtifactRecord) => { + artifacts.set(value.artifactId, { ...value }) + }, + get: async (id: string) => artifacts.get(id) ?? null, + list: async (runId: string) => + [...artifacts.values()] + .filter((value) => value.runId === runId) + .sort( + (a, b) => + a.createdAt - b.createdAt || + compare(a.artifactId, b.artifactId), + ), + + listForThread: async (threadId: string) => + [...artifacts.values()] + .filter((value) => value.threadId === threadId) + .sort( + (a, b) => + a.createdAt - b.createdAt || + compare(a.artifactId, b.artifactId), + ), + + delete: async (id: string) => { + artifacts.delete(id) + }, + deleteForRun: async (runId: string) => { + for (const [id, value] of artifacts) + if (value.runId === runId) artifacts.delete(id) + }, + }, + blobs: { + put: async ( + key: string, + body: BlobBody, + putOptions?: { + contentType?: string + customMetadata?: Record + }, + ) => { + const bytes = await bodyBytes(body) + const now = Date.now() + const record: BlobRecord = { + key, + size: bytes.byteLength, + etag: String(++etag), + contentType: + putOptions?.contentType ?? + (typeof Blob !== 'undefined' && body instanceof Blob + ? body.type || undefined + : undefined), + customMetadata: putOptions?.customMetadata + ? { ...putOptions.customMetadata } + : undefined, + createdAt: blobs.get(key)?.record.createdAt ?? now, + updatedAt: now, + } + blobs.set(key, { record, bytes: new Uint8Array(bytes) }) + return clone(record) + }, + get: async (key: string, getOptions?: BlobGetOptions) => { + const value = blobs.get(key) + if (!value) return null + const range = getOptions?.range + ? resolveBlobRange(value.bytes.byteLength, getOptions.range) + : { offset: 0, length: value.bytes.byteLength } + const bytes = value.bytes.slice( + range.offset, + range.offset + range.length, + ) + return { + ...clone(value.record), + ...(getOptions?.range ? { range } : {}), + body: new Blob([bytes]).stream(), + arrayBuffer: async () => + bytes.buffer.slice( + bytes.byteOffset, + bytes.byteOffset + bytes.byteLength, + ), + text: async () => new TextDecoder().decode(bytes), + } + }, + head: async (key: string) => clone(blobs.get(key)?.record ?? null), + delete: async (key: string) => { + blobs.delete(key) + }, + list: async (listOptions?: { + prefix?: string + cursor?: string + limit?: number + }) => { + const keys = [...blobs.keys()] + .filter((key) => key.startsWith(listOptions?.prefix ?? '')) + .filter( + (key) => + listOptions?.cursor === undefined || key > listOptions.cursor, + ) + .sort() + if (listOptions?.limit === 0) return { objects: [], truncated: false } + const page = + listOptions?.limit === undefined + ? keys + : keys.slice(0, listOptions.limit) + const truncated = + listOptions?.limit !== undefined && keys.length > page.length + const objects = page.map((key) => { + const value = blobs.get(key) + if (!value) throw new Error(`Missing blob for listed key: ${key}`) + return value.record + }) + return { + objects: clone(objects), + ...(truncated ? { cursor: page.at(-1), truncated: true } : {}), + } + }, + }, + }, + } + const checkpointStore = new MemorySnapshotCheckpointStore(state) + return { + persistence, + checkpoints: checkpointStore, + } +} diff --git a/packages/ai-sandbox/src/middleware.ts b/packages/ai-sandbox/src/middleware.ts index e4583482a8..ee0a0c28d5 100644 --- a/packages/ai-sandbox/src/middleware.ts +++ b/packages/ai-sandbox/src/middleware.ts @@ -41,7 +41,7 @@ import { SandboxInstanceStoreCapability } from './instance-store' import { computeWorkspaceHash } from './key' import { buildFileHookEvent, resolveFileEvents } from './file-diff' import { ProjectionCapability, provideWorkspaceProjection } from './projection' -import { resolveSecret } from './secrets' +import { resolveAllSecrets, resolveSecret } from './secrets' import { createToolHistoryRecorder, stripObservedToolCalls, @@ -49,12 +49,26 @@ import { import { watchWorkspace } from './watch' import { DEFAULT_WORKSPACE_ROOT } from './bootstrap' import { resolveHarnessCwd } from './harness-cwd' +import { ensureSandboxWithOutcome } from './sandbox' +import { + restoreSandboxFiles, + captureSandboxFiles, + captureSandboxArtifacts, + defaultSandboxSnapshotPolicy, +} from './snapshots' +import type { SandboxSnapshotPolicy } from './snapshots' +import type { + SandboxCheckpointStore, + SandboxCheckpointWriterLease, +} from './checkpoint-store' +import { SandboxCheckpointError } from './checkpoint-store' import type { InternalLogger } from '@tanstack/ai/adapter-internals' import type { LockStore } from '@tanstack/ai/locks' import type { AbortInfo, ChatMiddlewareContext, DefinedChatMiddleware, + ModelMessage, RunStore, SandboxFileEvent, SandboxFileHookEvent, @@ -75,6 +89,24 @@ import type { SandboxWatchHandle } from './watch' /** Per-request state we need to carry from `setup` to the terminal hooks. */ interface SandboxRunState { + snapshotLease?: SandboxCheckpointWriterLease + snapshotRenewal?: ReturnType + snapshotRenewTask?: Promise + snapshotCaptureTask?: Promise + snapshotRenewalGeneration: number + snapshotStop?: Promise + /** A detached or paused run cannot later publish portable state. */ + snapshotClosed?: boolean + snapshotLost?: Error + snapshotCleaned?: boolean + snapshotConfig?: NonNullable + snapshotPolicy?: SandboxSnapshotPolicy + snapshotRuntime?: { + persistence: NonNullable< + SandboxMiddlewareOptions['snapshots'] + >['persistence'] + completion: { waitForRunCompletion: () => Promise } + } /** * OPTIONAL because the state is registered BEFORE `definition.ensure()` is * awaited, and `ensure` is the slowest thing in the whole run — cloning a repo @@ -90,6 +122,7 @@ interface SandboxRunState { * completed. */ handle?: SandboxHandle + privateHandle?: boolean ensureCtx: SandboxEnsureContext watcher?: SandboxWatchHandle /** In-flight `enriched.diff()` promises queued by the `fileEvents.diff` @@ -114,6 +147,67 @@ interface SandboxRunState { const runState = new WeakMap() +function stopSnapshotLease( + state: SandboxRunState, + options: { closePortable?: boolean } = {}, +): Promise { + 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 !== undefined) clearTimeout(state.snapshotRenewal) + state.snapshotRenewal = undefined + const renewTask = state.snapshotRenewTask + const captureTask = state.snapshotCaptureTask + const lease = state.snapshotLease + state.snapshotLease = undefined + state.snapshotStop = (async () => { + await renewTask?.catch(() => {}) + await captureTask?.catch(() => {}) + await lease?.release() + })() + return state.snapshotStop +} + +function startSnapshotRenewal(state: SandboxRunState): void { + const lease = state.snapshotLease + if (!lease) return + const schedule = (): void => { + const generation = state.snapshotRenewalGeneration + state.snapshotRenewal = setTimeout(() => { + void (async (): Promise => { + state.snapshotRenewal = undefined + if ( + state.snapshotCleaned || + generation !== state.snapshotRenewalGeneration + ) + return + const renewal = Promise.resolve().then(async (): Promise => { + 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 = undefined + } + 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 @@ -122,7 +216,7 @@ const runState = new WeakMap() */ async function drainWatcher( state: SandboxRunState, - phase: 'finish' | 'abort' | 'error', + phase: 'finish' | 'pause' | 'abort' | 'error', ): Promise { // Guard `stop()`: a rejecting watcher teardown must NOT propagate out of // here, or the caller skips the `definition.destroy(...)` that follows — @@ -136,6 +230,18 @@ async function drainWatcher( if (state.watcher) state.logger?.sandbox('sandbox watcher stopped', { phase }) } +function canPublishPortableSnapshot( + state: SandboxRunState, + lease: SandboxCheckpointWriterLease, +): boolean { + 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. @@ -232,6 +338,38 @@ function tenantFrom( * single process but NOT across replicas. */ export interface SandboxMiddlewareOptions { + snapshots?: { + persistence: { + stores: { + messages: { + loadThread: (threadId: string) => Promise> + } + artifacts: { + listForThread: (threadId: string) => Promise< + ReadonlyArray<{ + artifactId: string + runId: string + threadId: string + blobKey?: string + name: string + mimeType: string + size: number + createdAt: number + }> + > + } + blobs: { + get: (key: string) => Promise<{ + arrayBuffer: () => Promise + } | null> + head: (key: string) => Promise + put: (key: string, body: Uint8Array) => Promise + } + } + } + checkpoints: SandboxCheckpointStore + policy?: SandboxSnapshotPolicy + } /** * Durable instance map (which provider sandbox to resume for a key). Pass * your own store to make resume survive across processes/replicas. @@ -355,6 +493,92 @@ export function withSandbox( async setup(ctx) { const ensureCtx = buildEnsureCtx(ctx, options) + const snapshotConfig = options?.snapshots + const snapshotWorkspaceHash = definition.workspace + ? computeWorkspaceHash(definition.workspace) + : undefined + const snapshotPolicy = snapshotConfig + ? snapshotConfig.policy + ? { + ...snapshotConfig.policy, + ...(snapshotWorkspaceHash === undefined + ? {} + : { workspaceHash: snapshotWorkspaceHash }), + } + : defaultSandboxSnapshotPolicy(snapshotWorkspaceHash) + : undefined + let snapshotRuntime: + | { + persistence: { + stores: { + messages: { + loadThread: ( + id: string, + ) => Promise> + } + artifacts: { + listForThread: (id: string) => Promise< + ReadonlyArray<{ + artifactId: string + runId: string + threadId: string + blobKey?: string + name: string + mimeType: string + size: number + createdAt: number + }> + > + } + blobs: { + get: (key: string) => Promise<{ + arrayBuffer: () => Promise + } | null> + head: (key: string) => Promise + put: (key: string, body: Uint8Array) => Promise + } + } + } + completion: { waitForRunCompletion: () => Promise } + } + | undefined + let snapshotLease: SandboxCheckpointWriterLease | undefined + 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 === undefined) + 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, + ) + } // Resolving here (not lazily on the abort path) is what keeps `setup` and // `onAbort` on one verdict: the payload the bus carries is the same object @@ -399,12 +623,17 @@ export function withSandbox( // sees the most complete state that exists at the moment it runs. const state: SandboxRunState = { ensureCtx, + snapshotRenewalGeneration: 0, pendingDiffs: [], toolHistory: createToolHistoryRecorder(), ...(logger ? { logger } : {}), ...(durability ? { durability } : {}), } runState.set(ctx, state) + if (snapshotLease) { + state.snapshotLease = snapshotLease + startSnapshotRenewal(state) + } // MAKE THE RUN FINDABLE BEFORE `ensure`, not after the run finally streams. // @@ -487,6 +716,10 @@ export function withSandbox( // because an agent that never ran wrote no journal to replay. if (durability !== undefined && durability.detachOnDisconnect) { getRunDisconnect(ctx, { optional: true })?.subscribe(async () => { + const snapshotStop = stopSnapshotLease(state, { + closePortable: true, + }) + void snapshotStop.catch(() => {}) // BOOKKEEPING ONLY — the run is still executing. Deliberately absent: // `drainWatcher` (would blind a live agent's file events for the whole // remainder) and `definition.destroy` (the run is still using the @@ -496,164 +729,243 @@ export function withSandbox( // A run with a cancel already recorded is left alone: that is `onAbort`'s // path, and stamping `detachedSince` on a deliberately-stopped run would // hand it to the reaper as reclaimable work. - if (await cancelIntent(durability, ctx.runId, false)) return + if (await cancelIntent(durability, ctx.runId, false)) { + await snapshotStop.catch((error: unknown) => { + 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: unknown) => { + state.logger?.warn('sandbox snapshot writer release failed', { + runId: ctx.runId, + phase: 'disconnect', + error, + }) + }) } }) } - const handle = await definition.ensure(ensureCtx) + let outcome: 'resumed' | 'native-restored' | 'created' = 'created' + let handle: SandboxHandle + 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 ?? DEFAULT_WORKSPACE_ROOT, + }, + checkpoint, + snapshotPolicy, + ) + } + } + } catch (error) { + await stopSnapshotLease(state).catch(() => {}) + if (state.handle && state.privateHandle) + await definition.destroy(ensureCtx).catch(() => {}) + throw error + } // MUTATE, don't re-`set`: a disconnect that landed during `ensure` already // captured this object. state.handle = handle - provideSandbox(ctx, handle) - if (definition.policy) provideSandboxPolicy(ctx, definition.policy) - - // Deliberately placed AFTER `logger` is in scope rather than next to the - // `provideSandboxDurability` call above — there is no logger to warn - // through until the runtime has been read. - // - // `ensureCtx.locks === undefined` counts as in-memory: `defineSandbox`'s - // `ensure` falls back to a process-lifetime `InMemoryLockStore` when no - // lock is wired, so an unwired lock has exactly the deficiency being - // warned about — it is the MOST in-memory case, not an exempt one. - if ( - durability !== undefined && - (ensureCtx.locks === undefined || - 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 }, - ) + if (snapshotConfig) { + state.snapshotConfig = snapshotConfig + state.snapshotPolicy = snapshotPolicy + state.snapshotRuntime = snapshotRuntime } - - const watchRoot = definition.workspace?.root ?? DEFAULT_WORKSPACE_ROOT - 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, + provideSandbox(ctx, handle) + if (definition.policy) provideSandboxPolicy(ctx, definition.policy) + + // Deliberately placed AFTER `logger` is in scope rather than next to the + // `provideSandboxDurability` call above — there is no logger to warn + // through until the runtime has been read. + // + // `ensureCtx.locks === undefined` counts as in-memory: `defineSandbox`'s + // `ensure` falls back to a process-lifetime `InMemoryLockStore` when no + // lock is wired, so an unwired lock has exactly the deficiency being + // warned about — it is the MOST in-memory case, not an exempt one. + if ( + durability !== undefined && + (ensureCtx.locks === undefined || + 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 ?? DEFAULT_WORKSPACE_ROOT + let baseSha = '' + try { + const shaRes = await handle.process.exec('git rev-parse HEAD', { + cwd: watchRoot, }) - } else { - // Non-zero exit: either not a git repository (non-git workspace) or a - // repo with no commits (no HEAD). Expected, but it silently degrades - // every subsequent diff to a full-file add-patch, so surface it - // under `sandbox` (with stderr) rather than leaving nothing to grep. - logger?.sandbox('sandbox git baseline unavailable (non-zero exit)', { + if (shaRes.exitCode === 0) { + baseSha = shaRes.stdout.trim() + logger?.sandbox('sandbox git baseline captured', { + root: watchRoot, + baseSha, + }) + } else { + // Non-zero exit: either not a git repository (non-git workspace) or a + // repo with no commits (no HEAD). Expected, but it silently degrades + // every subsequent diff to a full-file add-patch, so surface it + // under `sandbox` (with stderr) rather than leaving nothing to grep. + logger?.sandbox( + 'sandbox git baseline unavailable (non-zero exit)', + { + root: watchRoot, + exitCode: shaRes.exitCode, + stderr: shaRes.stderr, + }, + ) + } + } catch (error) { + // exec rejected (git not on PATH, exec seam broken) → baseSha stays '' + // and accessors fall back, but this is a real anomaly, not a plain + // non-git workspace, so warn. + logger?.warn('sandbox git baseline capture failed', { root: watchRoot, - exitCode: shaRes.exitCode, - stderr: shaRes.stderr, + error, }) } - } catch (error) { - // exec rejected (git not on PATH, exec seam broken) → baseSha stays '' - // and accessors fall back, but this is a real anomaly, not a plain - // non-git workspace, so warn. - logger?.warn('sandbox git baseline capture failed', { - root: watchRoot, - error, - }) - } - const workspace = definition.workspace - if (workspace !== undefined) { - const virtualRoot = workspace.root ?? DEFAULT_WORKSPACE_ROOT - 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 === undefined) { - 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 !== undefined - ? { scripts: workspace.scripts } - : {}), - }) - } + const workspace = definition.workspace + if (workspace !== undefined) { + const virtualRoot = workspace.root ?? DEFAULT_WORKSPACE_ROOT + 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 === undefined) { + 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 !== undefined + ? { scripts: workspace.scripts } + : {}), + }) + } - const hooks = definition.hooks - await hooks?.onReady?.(handle) - - const fe = resolveFileEvents(definition.fileEvents) - // THE SAME array the run state already holds, not a fresh one. The watcher - // callback below closes over this reference, and `drainWatcher` awaits - // `state.pendingDiffs` — a second array would silently drop every in-flight - // diff from the teardown drain. - const pendingDiffs = state.pendingDiffs - let watcher: SandboxWatchHandle | undefined - if (fe.enabled) { - watcher = await watchWorkspace(handle, { - onEvent: (event: SandboxFileEvent) => { - const enriched = buildFileHookEvent( - handle, - watchRoot, - baseSha, - event, - logger, - ) - void dispatchDefinitionHooks(hooks, enriched, logger) - runtime?.emit(enriched) - if (fe.diff) { - pendingDiffs.push( - enriched - .diff() - .then((diff) => { - runtime?.emitFileDiff({ path: event.path, diff }) - }) - .catch((error: unknown) => { - logger?.warn('sandbox file diff emit failed', { - path: event.path, - error, - }) - }), + const hooks = definition.hooks + await hooks?.onReady?.(handle) + + const fe = resolveFileEvents(definition.fileEvents) + // THE SAME array the run state already holds, not a fresh one. The watcher + // callback below closes over this reference, and `drainWatcher` awaits + // `state.pendingDiffs` — a second array would silently drop every in-flight + // diff from the teardown drain. + const pendingDiffs = state.pendingDiffs + let watcher: SandboxWatchHandle | undefined + if (fe.enabled) { + watcher = await watchWorkspace(handle, { + onEvent: (event: SandboxFileEvent) => { + const enriched = buildFileHookEvent( + handle, + watchRoot, + baseSha, + event, + logger, ) - } - }, - // Watch the SAME root the enrichment layer relativizes against - // (`buildFileHookEvent(handle, watchRoot, …)` and the `baseSha` - // capture). Without this the watcher defaults to `/workspace` while - // enrichment uses `watchRoot`, so a custom `workspace.root` makes the - // two look at different directories and git pathspecs break. - root: watchRoot, - ...(ctx.signal !== undefined ? { signal: ctx.signal } : {}), - ...(logger !== undefined ? { logger } : {}), - }) - logger?.sandbox('sandbox watcher started', { - root: watchRoot, - diff: fe.diff, - }) - } + void dispatchDefinitionHooks(hooks, enriched, logger) + runtime?.emit(enriched) + if (fe.diff) { + pendingDiffs.push( + enriched + .diff() + .then((diff) => { + runtime?.emitFileDiff({ path: event.path, diff }) + }) + .catch((error: unknown) => { + logger?.warn('sandbox file diff emit failed', { + path: event.path, + error, + }) + }), + ) + } + }, + // Watch the SAME root the enrichment layer relativizes against + // (`buildFileHookEvent(handle, watchRoot, …)` and the `baseSha` + // capture). Without this the watcher defaults to `/workspace` while + // enrichment uses `watchRoot`, so a custom `workspace.root` makes the + // two look at different directories and git pathspecs break. + root: watchRoot, + ...(ctx.signal !== undefined ? { signal: ctx.signal } : {}), + ...(logger !== undefined ? { logger } : {}), + }) + logger?.sandbox('sandbox watcher started', { + root: watchRoot, + diff: fe.diff, + }) + } - // MUTATE the object registered above rather than `set`-ing a second one: an - // abort that landed mid-setup already captured a reference to it (and may - // already be draining `pendingDiffs`), so replacing the entry would hand the - // teardown path a different object than the watcher writes into. - // `pendingDiffs` needs no copying — it IS `state.pendingDiffs`. - if (watcher) state.watcher = watcher + // MUTATE the object registered above rather than `set`-ing a second one: an + // abort that landed mid-setup already captured a reference to it (and may + // already be draining `pendingDiffs`), so replacing the entry would hand the + // teardown path a different object than the watcher writes into. + // `pendingDiffs` needs no copying — it IS `state.pendingDiffs`. + if (watcher) state.watcher = watcher + } catch (error) { + await drainWatcher(state, 'error') + await stopSnapshotLease(state).catch(() => {}) + if (state.privateHandle) + await definition.destroy(ensureCtx).catch(() => {}) + throw error + } }, // Keep the recorded tool history OUT of the request to the model. It is stored @@ -678,8 +990,17 @@ export function withSandbox( // Record the harness's own tool calls as transcript messages. Observe only: // returning nothing passes every chunk through untouched. - 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) { @@ -694,34 +1015,149 @@ export function withSandbox( await drainWatcher(state, 'finish') - const lifecycle = definition.lifecycle + let primaryError: unknown + try { + const snapshotCaptureTask = (async (): Promise => { + 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 - // `handle` is absent only if `setup` never got past `definition.ensure`, in - // which case there is no sandbox to snapshot. - 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(), + 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 ?? DEFAULT_WORKSPACE_ROOT, + }, + state.snapshotPolicy, + definition.workspace?.secrets !== undefined + ? 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 !== undefined + ? 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 + } + canPublishPortableSnapshot(state, lease) + })() + state.snapshotCaptureTask = snapshotCaptureTask + try { + await snapshotCaptureTask + } finally { + if (state.snapshotCaptureTask === snapshotCaptureTask) + state.snapshotCaptureTask = undefined + } + + const lifecycle = definition.lifecycle + + // `handle` is absent only if `setup` never got past `definition.ensure`, in + // which case there is no sandbox to snapshot. + 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: unknown + try { + await stopSnapshotLease(state, { closePortable: true }) + } catch (error) { + snapshotCleanupError = error + } + if (primaryError !== undefined) { + if (snapshotCleanupError !== undefined) + state.logger?.warn('sandbox snapshot writer release failed', { + runId: ctx.runId, + phase: 'finish', + error: snapshotCleanupError, + }) + throw primaryError } + if (snapshotCleanupError !== undefined) throw snapshotCleanupError }, async onAbort(ctx, info: AbortInfo) { @@ -732,6 +1168,12 @@ export function withSandbox( // the sandbox is about to be destroyed or merely detached, or the final // file's diff is dropped. await drainWatcher(state, 'abort') + let releaseError: unknown + try { + await stopSnapshotLease(state, { closePortable: true }) + } catch (error) { + releaseError = error + } const durability = state.durability const cancelled = await cancelIntent( @@ -760,10 +1202,12 @@ export function withSandbox( // unreachable one — the same reasoning `drainWatcher` applies to its own // guarded `stop()`. 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 } @@ -775,6 +1219,7 @@ export function withSandbox( // `destroyOnComplete:false` governs *successful completion*, never cancel. await definition.destroy(state.ensureCtx) await definition.hooks?.onDestroy?.() + if (releaseError) throw releaseError }, async onError(ctx, info) { @@ -782,6 +1227,12 @@ export function withSandbox( if (!state) return await drainWatcher(state, 'error') + let releaseError: unknown + try { + await stopSnapshotLease(state) + } catch (error) { + releaseError = error + } await definition.hooks?.onError?.(info.error) // On failure, only tear down when the lifecycle says so; otherwise leave @@ -790,6 +1241,7 @@ export function withSandbox( await definition.destroy(state.ensureCtx) await definition.hooks?.onDestroy?.() } + if (releaseError) throw releaseError }, }) } diff --git a/packages/ai-sandbox/src/sandbox.ts b/packages/ai-sandbox/src/sandbox.ts index 35e5d28f2e..f8058c35e9 100644 --- a/packages/ai-sandbox/src/sandbox.ts +++ b/packages/ai-sandbox/src/sandbox.ts @@ -95,10 +95,62 @@ export interface SandboxDefinition { key: (ctx: SandboxEnsureContext) => string /** Resume-or-create the sandbox for this thread/run. */ ensure: (ctx: SandboxEnsureContext) => Promise + /** Resume an existing sandbox only. Never creates or restores a sandbox. */ + ensureExisting: (ctx: SandboxEnsureContext) => Promise /** Tear down the sandbox recorded for this key. */ destroy: (ctx: SandboxEnsureContext) => Promise } +export type SandboxEnsureOutcome = { + handle: SandboxHandle + outcome: 'resumed' | 'native-restored' | 'created' +} + +const outcomeEnsure = new WeakMap< + object, + (ctx: SandboxEnsureContext) => Promise +>() + +interface SandboxEnsureExistingStage { + key: string + workspace: WorkspaceDefinition | undefined + resolvedSecrets: Readonly> | undefined + snapshotMaxAge: string | undefined + resume: SandboxProvider['resume'] +} + +const existingEnsure = new WeakMap< + object, + ( + ctx: SandboxEnsureContext, + stage?: SandboxEnsureExistingStage, + ) => Promise +>() + +export function stageEnsureExistingSandbox( + definition: SandboxDefinition, +): ( + ctx: SandboxEnsureContext, + stage: SandboxEnsureExistingStage, +) => Promise { + const fn = existingEnsure.get(definition) + if (fn) return (ctx, stage) => fn(ctx, stage) + const ensureExisting = definition.ensureExisting.bind(definition) + return (ctx) => ensureExisting(ctx) +} + +export function ensureSandboxWithOutcome( + definition: SandboxDefinition, + ctx: SandboxEnsureContext, +) { + const fn = outcomeEnsure.get(definition) + if (!fn) + throw new Error( + 'Sandbox snapshot mode requires a definition created by defineSandbox()', + ) + return fn(ctx) +} + /** * Parse a human-readable duration string into milliseconds. * Supports `'h'` (hours) and `'m'` (minutes). @@ -136,9 +188,10 @@ const fallbackLocks = new InMemoryLockStore() async function applyWorkspaceSecrets( handle: SandboxHandle, workspace: WorkspaceDefinition | undefined, + stagedSecrets?: Readonly>, ): Promise { if (workspace?.secrets === undefined) return - const resolved = resolveAllSecrets(workspace.secrets) + const resolved = stagedSecrets ?? resolveAllSecrets(workspace.secrets) if (Object.keys(resolved).length === 0) return await handle.env.set(resolved) } @@ -155,7 +208,9 @@ export function defineSandbox(config: SandboxConfig): SandboxDefinition { tenant: ctx.tenant, }) - const ensure = async (ctx: SandboxEnsureContext): Promise => { + const ensureWithOutcome = async ( + ctx: SandboxEnsureContext, + ): Promise => { const store = ctx.store ?? fallbackStore const locks = ctx.locks ?? fallbackLocks const key = computeSandboxKey(keyInputFor(ctx)) @@ -186,7 +241,7 @@ export function defineSandbox(config: SandboxConfig): SandboxDefinition { latestRunId: ctx.runId, updatedAt: Date.now(), }) - return resumed + return { handle: resumed, outcome: 'resumed' } } // 2) Else restore from the latest snapshot, if supported. if ( @@ -211,7 +266,7 @@ export function defineSandbox(config: SandboxConfig): SandboxDefinition { latestRunId: ctx.runId, updatedAt: Date.now(), }) - return restored + return { handle: restored, outcome: 'native-restored' } } } // 3) Else fall through and re-create under the same identity @@ -265,10 +320,52 @@ export function defineSandbox(config: SandboxConfig): SandboxDefinition { latestRunId: ctx.runId, updatedAt: Date.now(), }) - return created + return { handle: created, outcome: 'created' } }) } + const ensure = async (ctx: SandboxEnsureContext): Promise => + (await ensureWithOutcome(ctx)).handle + + const ensureExistingWithStage = async ( + ctx: SandboxEnsureContext, + stage?: SandboxEnsureExistingStage, + ): Promise => { + const store = ctx.store ?? fallbackStore + const locks = ctx.locks ?? fallbackLocks + const key = stage?.key ?? computeSandboxKey(keyInputFor(ctx)) + const workspace = stage?.workspace ?? config.workspace + const snapshotMaxAge = stage + ? stage.snapshotMaxAge + : config.lifecycle?.snapshotMaxAge + const resume = stage?.resume ?? config.provider.resume.bind(config.provider) + return locks.withLock(`sandbox:${key}`, async () => { + const existing = await store.get(key) + const maxAgeMs = parseMaxAgeMs(snapshotMaxAge) + if ( + !existing || + (maxAgeMs !== undefined && Date.now() - existing.updatedAt > maxAgeMs) + ) + return null + const resumed = await resume({ + id: existing.providerSandboxId, + signal: ctx.signal, + }) + if (!resumed) return null + await applyWorkspaceSecrets(resumed, workspace, stage?.resolvedSecrets) + await store.upsert({ + ...existing, + latestRunId: ctx.runId, + updatedAt: Date.now(), + }) + return resumed + }) + } + + const ensureExisting = ( + ctx: SandboxEnsureContext, + ): Promise => ensureExistingWithStage(ctx) + const destroy = async (ctx: SandboxEnsureContext): Promise => { const store = ctx.store ?? fallbackStore const key = computeSandboxKey(keyInputFor(ctx)) @@ -302,7 +399,7 @@ export function defineSandbox(config: SandboxConfig): SandboxDefinition { await store.delete(key) } - return { + const definition: SandboxDefinition = { id: config.id, provider: config.provider, workspace: config.workspace, @@ -312,6 +409,10 @@ export function defineSandbox(config: SandboxConfig): SandboxDefinition { fileEvents: config.fileEvents, key: (ctx) => computeSandboxKey(keyInputFor(ctx)), ensure, + ensureExisting, destroy, } + outcomeEnsure.set(definition, ensureWithOutcome) + existingEnsure.set(definition, ensureExistingWithStage) + return definition } diff --git a/packages/ai-sandbox/src/snapshot-operations.ts b/packages/ai-sandbox/src/snapshot-operations.ts new file mode 100644 index 0000000000..156e853eee --- /dev/null +++ b/packages/ai-sandbox/src/snapshot-operations.ts @@ -0,0 +1,424 @@ +import { + captureSandboxArtifacts, + captureSandboxFiles, + defaultSandboxSnapshotPolicy, + SandboxSnapshotError, +} from './snapshots' +import { resolveAllSecrets } from './secrets' +import { computeSandboxKey, computeWorkspaceHash } from './key' +import { stageEnsureExistingSandbox } from './sandbox' +import type { ModelMessage } from '@tanstack/ai' +import type { LockStore } from '@tanstack/ai/locks' +import type { + SandboxCheckpoint, + SandboxCheckpointStore, + SandboxCheckpointWriterLease, +} from './checkpoint-store' +import type { SandboxInstanceStore } from './instance-store' +import type { SandboxDefinition } from './sandbox' +import type { SandboxSnapshotBundle, SandboxSnapshotPolicy } from './snapshots' +import type { WorkspaceDefinition } from './workspace' + +type SnapshotPersistence = { + stores: { + messages: { + loadThread: (threadId: string) => Promise> + } + artifacts: NonNullable + blobs: SandboxSnapshotBundle['blobs'] + } +} + +export interface SandboxSnapshots { + persistence: SnapshotPersistence + checkpoints: SandboxCheckpointStore + policy?: SandboxSnapshotPolicy +} + +type Failure = { error: unknown } + +async function withWriterLease( + acquire: () => Promise, + renew: boolean, + operation: ( + writer: SandboxCheckpointWriterLease, + throwIfLost: () => Promise, + ) => Promise, +): Promise { + const writer = await acquire() + const release = writer.release.bind(writer) + const renewWriter = renew ? writer.renew.bind(writer) : undefined + const renewAfterMs = renew ? writer.renewAfterMs : undefined + let renewalTimer: ReturnType | undefined + let renewalTask: Promise | undefined + let renewalFailure: Failure | undefined + let stopped = false + + const scheduleRenewal = (): void => { + if (renewWriter === undefined || renewAfterMs === undefined) return + renewalTimer = setTimeout(() => { + renewalTimer = undefined + renewalTask = (async () => { + try { + await renewWriter() + } catch (error) { + renewalFailure = { error } + } finally { + renewalTask = undefined + } + if (!stopped && renewalFailure === undefined) scheduleRenewal() + })() + }, renewAfterMs) + } + if (renew) scheduleRenewal() + + const throwIfLost = async (): Promise => { + await renewalTask + if (renewalFailure !== undefined) throw renewalFailure.error + } + + let outcome: { value: T } | undefined + let operationFailure: Failure | undefined + try { + outcome = { value: await operation(writer, throwIfLost) } + } catch (error) { + operationFailure = { error } + } + + stopped = true + if (renewalTimer !== undefined) clearTimeout(renewalTimer) + await renewalTask + let releaseFailure: Failure | undefined + try { + await release() + } catch (error) { + releaseFailure = { error } + } + + if (renewalFailure !== undefined) throw renewalFailure.error + if (operationFailure !== undefined) throw operationFailure.error + if (releaseFailure !== undefined) throw releaseFailure.error + if (outcome === undefined) throw new Error('Writer operation had no outcome') + return outcome.value +} + +function stageWorkspace( + workspace: WorkspaceDefinition | undefined, +): WorkspaceDefinition | undefined { + if (workspace === undefined) return undefined + const source = workspace.source + const packageManager = workspace.packageManager + const setup = workspace.setup + const scripts = workspace.scripts + const skills = workspace.skills + const instructions = workspace.instructions + const plugins = workspace.plugins + const secrets = workspace.secrets + const root = workspace.root + return { + source, + ...(Object.hasOwn(workspace, 'packageManager') ? { packageManager } : {}), + ...(Object.hasOwn(workspace, 'setup') ? { setup } : {}), + ...(Object.hasOwn(workspace, 'scripts') ? { scripts } : {}), + ...(Object.hasOwn(workspace, 'skills') ? { skills } : {}), + ...(Object.hasOwn(workspace, 'instructions') ? { instructions } : {}), + ...(Object.hasOwn(workspace, 'plugins') ? { plugins } : {}), + ...(Object.hasOwn(workspace, 'secrets') ? { secrets } : {}), + ...(Object.hasOwn(workspace, 'root') ? { root } : {}), + } +} + +function effectivePolicy( + supplied: SandboxSnapshotPolicy | undefined, + workspaceHash: string | undefined, +): SandboxSnapshotPolicy { + if (supplied === undefined) return defaultSandboxSnapshotPolicy(workspaceHash) + const suppliedWorkspaceHash = supplied.workspaceHash + const include = supplied.include + const exclude = supplied.exclude + const redact = supplied.redact + return { + ...(suppliedWorkspaceHash === undefined + ? {} + : { workspaceHash: suppliedWorkspaceHash }), + ...(workspaceHash === undefined ? {} : { workspaceHash }), + ...(include === undefined ? {} : { include }), + ...(exclude === undefined ? {} : { exclude }), + ...(redact === undefined ? {} : { redact }), + } +} + +function stageInstanceStore(store: SandboxInstanceStore): SandboxInstanceStore { + const get = store.get.bind(store) + const upsert = store.upsert.bind(store) + const deleteRecord = store.delete.bind(store) + return { get, upsert, delete: deleteRecord } +} + +function stageLockStore(locks: LockStore | undefined): LockStore | undefined { + if (locks === undefined) return undefined + const withLock = locks.withLock.bind(locks) + return { withLock } +} + +export async function saveNamedSandboxSnapshot(input: { + definition: SandboxDefinition + threadId: string + runId: string + instances: SandboxInstanceStore + snapshots: SandboxSnapshots + label: string + tenant?: { userId?: string; orgId?: string } + locks?: LockStore + signal?: AbortSignal + adapterName?: string +}): Promise { + const definition = input.definition + const threadId = input.threadId + const runId = input.runId + const instances = stageInstanceStore(input.instances) + const snapshots = input.snapshots + const label = input.label + const suppliedTenant = input.tenant + const tenantUserId = suppliedTenant?.userId + const tenantOrgId = suppliedTenant?.orgId + const tenant = suppliedTenant + ? { + ...(tenantUserId === undefined ? {} : { userId: tenantUserId }), + ...(tenantOrgId === undefined ? {} : { orgId: tenantOrgId }), + } + : undefined + const locks = stageLockStore(input.locks) + const signal = input.signal + const adapterName = input.adapterName + const lifecycle = definition.lifecycle + const reuse = lifecycle?.reuse + const snapshotMaxAge = lifecycle?.snapshotMaxAge + const workspace = stageWorkspace(definition.workspace) + const sandboxId = definition.id + const provider = definition.provider + const providerName = provider.name + const resume = provider.resume.bind(provider) + const ensureExisting = stageEnsureExistingSandbox(definition) + const persistence = snapshots.persistence + const stores = persistence.stores + const messages = stores.messages + const loadThread = messages.loadThread.bind(messages) + const artifactStore = stores.artifacts + const listForThread = artifactStore.listForThread.bind(artifactStore) + const suppliedBlobs = stores.blobs + const getBlob = suppliedBlobs.get.bind(suppliedBlobs) + const headBlob = suppliedBlobs.head.bind(suppliedBlobs) + const putBlob = suppliedBlobs.put.bind(suppliedBlobs) + const blobs: SandboxSnapshotBundle['blobs'] = { + get: getBlob, + head: headBlob, + put: putBlob, + } + const checkpoints = snapshots.checkpoints + const acquireWriter = checkpoints.acquireWriter.bind(checkpoints) + const getHead = checkpoints.getHead.bind(checkpoints) + const append = checkpoints.append.bind(checkpoints) + const policy = effectivePolicy( + snapshots.policy, + workspace === undefined ? undefined : computeWorkspaceHash(workspace), + ) + const workspaceSecrets = workspace?.secrets + const secrets = workspaceSecrets ? resolveAllSecrets(workspaceSecrets) : {} + const workspaceRoot = workspace?.root + const key = computeSandboxKey({ + threadId, + sandboxId, + providerName, + workspace, + tenant, + }) + + return withWriterLease( + () => acquireWriter(threadId), + true, + async (writer, throwIfLost) => { + if (reuse === 'none') + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_REUSE_NONE', + 'Named snapshots require a reusable sandbox lifecycle', + ) + const handle = await ensureExisting( + { + threadId, + runId, + store: instances, + locks, + tenant, + signal, + adapterName, + }, + { + key, + workspace, + resolvedSecrets: workspaceSecrets ? secrets : undefined, + snapshotMaxAge, + resume, + }, + ) + if (!handle) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_MISSING_REUSABLE_SANDBOX', + 'Named snapshots require an existing resumable sandbox', + ) + const conversation = await loadThread(threadId) + const files = await captureSandboxFiles( + handle, + { blobs, workspaceRoot }, + policy, + secrets, + ) + const artifacts = await captureSandboxArtifacts( + { + blobs, + artifacts: { listForThread }, + }, + threadId, + secrets, + ) + const parentCheckpointId = await getHead(threadId) + await throwIfLost() + const checkpoint: SandboxCheckpoint = { + id: crypto.randomUUID(), + threadId, + parentCheckpointId, + createdAt: Date.now(), + reason: 'named', + label, + sourceRunId: runId, + files: files.files, + conversation, + artifacts, + } + await append({ + checkpoint, + expectedHeadId: parentCheckpointId, + writer, + }) + await throwIfLost() + return checkpoint + }, + ) +} + +export async function forkFromSandboxSnapshot(input: { + sourceThreadId: string + sourceCheckpointId: string + destinationThreadId: string + snapshots: SandboxSnapshots + destinationCheckpointId?: string + createdAt?: number +}): Promise { + const sourceThreadId = input.sourceThreadId + const sourceCheckpointId = input.sourceCheckpointId + const destinationThreadId = input.destinationThreadId + const snapshots = input.snapshots + const suppliedDestinationCheckpointId = input.destinationCheckpointId + const suppliedCreatedAt = input.createdAt + const destinationCheckpointId = + suppliedDestinationCheckpointId ?? crypto.randomUUID() + const createdAt = suppliedCreatedAt ?? Date.now() + const checkpoints = snapshots.checkpoints + const acquireWriter = checkpoints.acquireWriter.bind(checkpoints) + const forkFromCheckpoint = checkpoints.forkFromCheckpoint?.bind(checkpoints) + + return withWriterLease( + () => acquireWriter(destinationThreadId), + false, + async (writer) => { + if (forkFromCheckpoint === undefined) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_FORK_UNAVAILABLE', + 'The checkpoint store does not support atomic forks', + ) + const result = await forkFromCheckpoint({ + sourceThreadId, + sourceCheckpointId, + destinationThreadId, + destinationCheckpointId, + createdAt, + writer, + }) + return result.checkpoint + }, + ) +} + +async function sha256(bytes: Uint8Array): Promise { + const digest = await crypto.subtle.digest('SHA-256', new Uint8Array(bytes)) + return Array.from(new Uint8Array(digest), (byte) => + byte.toString(16).padStart(2, '0'), + ).join('') +} + +export async function resolveSnapshotArtifact(input: { + threadId: string + checkpointId: string + artifactId: string + snapshots: SandboxSnapshots +}): Promise<{ + artifact: SandboxCheckpoint['artifacts'][number] + bytes: Uint8Array +}> { + const threadId = input.threadId + const checkpointId = input.checkpointId + const artifactId = input.artifactId + const snapshots = input.snapshots + const checkpoints = snapshots.checkpoints + const getCheckpoint = checkpoints.get.bind(checkpoints) + const persistence = snapshots.persistence + const stores = persistence.stores + const blobs = stores.blobs + const getBlob = blobs.get.bind(blobs) + const checkpoint = await getCheckpoint(checkpointId) + if (!checkpoint) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_MISSING_CHECKPOINT_ARTIFACT', + 'Snapshot checkpoint does not exist', + ) + const checkpointThreadId = checkpoint.threadId + const checkpointArtifacts = checkpoint.artifacts + if (checkpointThreadId !== threadId) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_FOREIGN_CHECKPOINT_ARTIFACT', + 'Snapshot checkpoint belongs to another thread', + ) + const foundArtifact = checkpointArtifacts.find( + (value) => value.artifactId === artifactId, + ) + if (!foundArtifact) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_MISSING_CHECKPOINT_ARTIFACT', + 'Snapshot artifact does not exist', + ) + const artifact = { + artifactId: foundArtifact.artifactId, + name: foundArtifact.name, + mimeType: foundArtifact.mimeType, + size: foundArtifact.size, + blobKey: foundArtifact.blobKey, + createdAt: foundArtifact.createdAt, + } + const blob = await getBlob(artifact.blobKey) + if (!blob) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_ARTIFACT_BYTES', + 'Snapshot artifact blob does not exist', + ) + const arrayBuffer = blob.arrayBuffer.bind(blob) + const bytes = new Uint8Array(await arrayBuffer()) + if ( + bytes.byteLength !== artifact.size || + artifact.blobKey !== `sandbox-artifacts/sha256/${await sha256(bytes)}` + ) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_ARTIFACT_BYTES', + 'Snapshot artifact bytes do not match metadata', + ) + return { artifact: { ...artifact }, bytes: bytes.slice() } +} diff --git a/packages/ai-sandbox/src/snapshots.ts b/packages/ai-sandbox/src/snapshots.ts new file mode 100644 index 0000000000..d9a81c678d --- /dev/null +++ b/packages/ai-sandbox/src/snapshots.ts @@ -0,0 +1,670 @@ +import type { SandboxHandle, SandboxFsStat } from './contracts' +import type { + SandboxSnapshotArtifact, + SandboxSnapshotEntry, +} from './checkpoint-store' +import type { MemoryArtifactRecord as ArtifactRecord } from './memory-snapshot-types' + +type SnapshotBlobStore = { + get: (key: string) => Promise<{ + arrayBuffer: () => Promise + } | null> + head: (key: string) => Promise + put: (key: string, body: Uint8Array) => Promise +} + +export interface SandboxSnapshotPolicy { + /** Exact workspace projection hash, when known. */ + workspaceHash?: string + include?: (path: string, kind: 'file' | 'dir') => boolean + exclude?: (path: string, kind: 'file' | 'dir') => boolean + redact?: (input: { + path: string + bytes: Uint8Array + resolvedSecrets: Readonly> + }) => Uint8Array +} + +export interface SandboxSnapshotBundle { + blobs: SnapshotBlobStore + /** Internal resolved workspace root. */ + workspaceRoot?: string + /** Internal persistence stores used to capture immutable artifact bytes. */ + artifacts?: { + listForThread: (threadId: string) => Promise> + } + resolveArtifactBlobKey?: (record: ArtifactRecord) => string +} + +export type SandboxSnapshotErrorCode = + | 'SANDBOX_SNAPSHOT_MISSING_REUSABLE_SANDBOX' + | 'SANDBOX_SNAPSHOT_REUSE_NONE' + | 'SANDBOX_SNAPSHOT_MISSING_CHECKPOINT_ARTIFACT' + | 'SANDBOX_SNAPSHOT_FOREIGN_CHECKPOINT_ARTIFACT' + | 'SANDBOX_SNAPSHOT_INVALID_ARTIFACT_BYTES' + | 'SANDBOX_SNAPSHOT_FORK_UNAVAILABLE' + | 'SANDBOX_SNAPSHOT_INVALID_PATH' + | 'SANDBOX_SNAPSHOT_INVALID_WORKSPACE' + | 'SANDBOX_SNAPSHOT_LSTAT_REQUIRED' + | 'SANDBOX_SNAPSHOT_UNSUPPORTED_ENTRY' + | 'SANDBOX_SNAPSHOT_MISSING_BLOB' + | 'SANDBOX_SNAPSHOT_INVALID_BLOB' + | 'SANDBOX_SNAPSHOT_ARTIFACT_SUPPORT_REQUIRED' + | 'SANDBOX_SNAPSHOT_MISSING_ARTIFACT_BLOB' + +export class SandboxSnapshotError extends Error { + readonly code: SandboxSnapshotErrorCode + constructor(code: SandboxSnapshotErrorCode, message: string) { + super(message) + this.name = 'SandboxSnapshotError' + this.code = code + } +} + +const DEFAULT_ROOT = '/workspace' +const PROJECTED_SKILL_ROOTS = new Set(['.claude', '.codex', '.grok']) + +function isFrameworkGeneratedSymlinkPath(path: string): boolean { + if (path === 'CLAUDE.md' || path === 'GEMINI.md') return true + const segments = path.split('/') + return ( + segments.length === 3 && + PROJECTED_SKILL_ROOTS.has(segments[0] ?? '') && + segments[1] === 'skills' + ) +} + +function defaultExcluded(path: string, workspaceHash?: string): boolean { + const segments = path.split('/') + return ( + isFrameworkGeneratedSymlinkPath(path) || + segments.some( + (segment) => + segment === '.git' || + segment === 'node_modules' || + segment.startsWith('.env'), + ) || + (workspaceHash !== undefined && + segments[0] === `.tanstack-projected-${workspaceHash}`) + ) +} + +function isProtectedPath(path: string, workspaceHash?: string): boolean { + return ( + workspaceHash !== undefined && + path.split('/')[0] === `.tanstack-projected-${workspaceHash}` + ) +} + +const FILE_BLOB_KEY = /^sandbox-files\/sha256\/[0-9a-f]{64}$/ + +export function defaultSandboxSnapshotPolicy( + workspaceHash?: string, +): SandboxSnapshotPolicy { + return { + workspaceHash, + exclude: (path) => defaultExcluded(path, workspaceHash), + } +} + +function normalize(path: string): string { + if (path.includes('\\')) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_PATH', + `Unsafe snapshot path '${path}'`, + ) + const value = path + if ( + !value || + value.includes('\0') || + value.startsWith('/') || + /^[A-Za-z]:/.test(value) || + value.endsWith('/') + ) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_PATH', + `Unsafe snapshot path '${path}'`, + ) + const parts = value.split('/') + if ( + parts.some((part) => !part || part === '.' || part === '..') || + parts.join('/') !== value + ) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_PATH', + `Unsafe snapshot path '${path}'`, + ) + return value +} + +function childPath( + parent: string, + child: { name: string; path: string }, +): { absolute: string; relative: string } { + if ( + !child.name || + child.name.includes('/') || + child.name.includes('\\') || + child.name.includes('\0') || + child.name === '.' || + child.name === '..' + ) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_WORKSPACE', + `Invalid workspace entry '${child.name}'`, + ) + const absolute = `${parent}/${child.name}` + if (child.path !== absolute) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_WORKSPACE', + `Invalid workspace entry path '${child.path}'`, + ) + return { absolute, relative: child.name } +} + +function lstat( + handle: SandboxHandle, + path: string, +): Promise { + if (!handle.fs.lstat) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_LSTAT_REQUIRED', + 'Snapshot operations require fs.lstat', + ) + return handle.fs.lstat(path) +} + +function assertSupported(stat: SandboxFsStat, path: string): void { + if ( + stat.type === 'symlink' || + stat.type === 'other' || + (stat.type === 'file' && (stat.mode & 0o111) !== 0) + ) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_UNSUPPORTED_ENTRY', + `Unsupported entry '${path}'`, + ) +} + +function included( + path: string, + kind: 'file' | 'dir', + policy: SandboxSnapshotPolicy, +): boolean { + if (policy.exclude?.(path, kind)) return false + return kind === 'dir' ? true : (policy.include?.(path, kind) ?? true) +} + +async function hash(bytes: Uint8Array): Promise { + // TypeScript requires an ArrayBuffer-backed view; Uint8Array can also use SharedArrayBuffer. + const digest = await crypto.subtle.digest('SHA-256', new Uint8Array(bytes)) + return Array.from(new Uint8Array(digest), (byte) => + byte.toString(16).padStart(2, '0'), + ).join('') +} + +async function putIfAbsent( + blobs: SnapshotBlobStore, + bytes: Uint8Array, + keys: Map, +): Promise<{ key: string; size: number }> { + const key = `sandbox-files/sha256/${await hash(bytes)}` + if (!keys.has(key)) { + if (!(await blobs.head(key))) await blobs.put(key, bytes) + keys.set(key, key) + } + return { key, size: bytes.byteLength } +} + +async function redactBytes( + bytes: Uint8Array, + resolvedSecrets: Readonly>, +): Promise { + const output = bytes.slice() + const redacted = new Uint8Array(bytes.length) + const secrets = Object.values(resolvedSecrets) + .filter(Boolean) + .map((secret) => new TextEncoder().encode(secret)) + .sort((a, b) => b.length - a.length || compareBytes(a, b)) + for (const needle of secrets) { + if (!needle.length || needle.length > bytes.length) continue + for (let start = 0; start <= bytes.length - needle.length; start++) { + let match = true + for (let index = 0; index < needle.length; index++) + if (bytes[start + index] !== needle[index]) { + match = false + break + } + if (!match) continue + if (match) redacted.fill(1, start, start + needle.length) + } + } + for (let index = 0; index < output.length; index++) + if (redacted[index]) output[index] = 0 + return output +} + +export async function captureSandboxFiles( + handle: SandboxHandle, + bundle: SandboxSnapshotBundle, + policy: SandboxSnapshotPolicy = defaultSandboxSnapshotPolicy(), + resolvedSecrets: Readonly> = {}, +): Promise<{ files: Array }> { + if (!handle.fs.lstat) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_LSTAT_REQUIRED', + 'Snapshot capture requires fs.lstat', + ) + const rootPath = bundle.workspaceRoot ?? DEFAULT_ROOT + const root = await lstat(handle, rootPath) + if (!root || root.type !== 'dir') + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_WORKSPACE', + 'Snapshot workspace is missing or is not a directory', + ) + assertSupported(root, rootPath) + const files: Array = [] + const destinationKeys = new Map() + const walk = async (absolute: string, relative: string): Promise => { + const stat = await lstat(handle, absolute) + if (!stat) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_WORKSPACE', + `Snapshot entry disappeared '${relative}'`, + ) + assertSupported(stat, relative) + if (stat.type !== 'file' && stat.type !== 'dir') + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_UNSUPPORTED_ENTRY', + `Unsupported entry '${relative}'`, + ) + if ( + relative && + (isProtectedPath(relative, policy.workspaceHash) || + policy.exclude?.(relative, stat.type)) + ) + return false + if (stat.type === 'file') { + const path = normalize(relative) + if (policy.include && !policy.include(path, 'file')) return false + let bytes = await handle.fs.readBytes(absolute) + if (policy.redact) bytes = policy.redact({ path, bytes, resolvedSecrets }) + bytes = await redactBytes(bytes, resolvedSecrets) + const blob = await putIfAbsent(bundle.blobs, bytes, destinationKeys) + files.push({ path, kind: 'file', blobKey: blob.key, size: blob.size }) + return true + } + const children = await handle.fs.list(absolute) + let hasCapturedChild = false + for (const child of children) { + const childEntry = childPath(absolute, child) + const childRelative = relative + ? `${relative}/${childEntry.relative}` + : childEntry.relative + if ( + isProtectedPath(childRelative, policy.workspaceHash) || + policy.exclude?.(childRelative, child.type) + ) + continue + hasCapturedChild = + (await walk(childEntry.absolute, childRelative)) || hasCapturedChild + } + if ( + relative && + !hasCapturedChild && + (!policy.include || policy.include(relative, 'dir')) + ) + files.push({ path: normalize(relative), kind: 'dir' }) + return ( + hasCapturedChild || + (relative !== '' && (!policy.include || policy.include(relative, 'dir'))) + ) + } + await walk(rootPath, '') + files.sort((a, b) => comparePath(a.path, b.path)) + return { files } +} + +type PlannedEntry = SandboxSnapshotEntry & { path: string } +type PlannedFile = Extract + +function validateManifest( + snapshot: { + files: ReadonlyArray + }, + policy: SandboxSnapshotPolicy, +): Array { + const paths = new Map() + for (const entry of snapshot.files) { + const path = normalize(entry.path) + for (const ancestor of parents(path)) { + if ( + isProtectedPath(ancestor, policy.workspaceHash) || + policy.exclude?.(ancestor, 'dir') + ) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_PATH', + `Excluded snapshot ancestor '${ancestor}'`, + ) + } + if (isProtectedPath(path, policy.workspaceHash)) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_PATH', + `Protected snapshot path '${path}'`, + ) + if ( + !included(path, entry.kind, policy) || + (entry.kind === 'dir' && policy.include?.(path, 'dir') === false) + ) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_PATH', + `Excluded snapshot path '${path}'`, + ) + if (paths.has(path)) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_PATH', + `Duplicate path '${path}'`, + ) + if (entry.kind === 'file') { + const { blobKey, size } = entry + if ( + typeof blobKey !== 'string' || + !FILE_BLOB_KEY.test(blobKey) || + typeof size !== 'number' || + !Number.isSafeInteger(size) || + size < 0 + ) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_PATH', + `Invalid file entry '${path}'`, + ) + paths.set(path, { path, kind: 'file', blobKey, size }) + } else if (entry.kind === 'dir') paths.set(path, { path, kind: 'dir' }) + else + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_PATH', + `Unknown snapshot entry '${path}'`, + ) + } + for (const [path] of paths) + for ( + let index = path.indexOf('/'); + index !== -1; + index = path.indexOf('/', index + 1) + ) { + const parent = paths.get(path.slice(0, index)) + if (parent?.kind === 'file') + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_PATH', + `File ancestor '${parent.path}'`, + ) + } + return [...paths.values()] +} + +async function loadBlobs( + entries: ReadonlyArray, + bundle: SandboxSnapshotBundle, +): Promise> { + const blobs = new Map() + for (const entry of entries) + if (entry.kind === 'file' && !blobs.has(entry.blobKey)) { + const object = await bundle.blobs.get(entry.blobKey) + if (!object) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_MISSING_BLOB', + `Missing snapshot blob '${entry.blobKey}'`, + ) + const bytes = new Uint8Array(await object.arrayBuffer()) + const expectedKey = `sandbox-files/sha256/${await hash(bytes)}` + if (entry.blobKey !== expectedKey) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_BLOB', + `Invalid content for snapshot blob '${entry.blobKey}'`, + ) + blobs.set(entry.blobKey, bytes) + } + for (const entry of entries) + if (entry.kind === 'file') { + const bytes = blobs.get(entry.blobKey) + if (!bytes || bytes.byteLength !== entry.size) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_BLOB', + `Wrong size for snapshot blob '${entry.blobKey}'`, + ) + } + return blobs +} + +type CurrentEntry = { path: string; kind: 'file' | 'dir'; protected?: boolean } + +async function scanCurrent( + handle: SandboxHandle, + absolute: string, + relative: string, + policy: SandboxSnapshotPolicy, +): Promise> { + const stat = await lstat(handle, absolute) + if (!stat) return [] + assertSupported(stat, relative) + if (stat.type === 'file') return [{ path: relative, kind: 'file' }] + const paths: Array = [] + for (const child of await handle.fs.list(absolute)) { + const childEntry = childPath(absolute, child) + const childRelative = relative + ? `${relative}/${childEntry.relative}` + : childEntry.relative + if (isProtectedPath(childRelative, policy.workspaceHash)) { + // Preserve the marker and its parent directories without inspecting the + // protected tree. The restore planner uses this marker to avoid removing + // an ancestor directory that contains it. + paths.push({ path: childRelative, kind: 'dir', protected: true }) + continue + } + if (!included(childRelative, child.type, policy)) { + // Excluded entries are outside the portable snapshot. Keep a protected + // marker so removing a parent directory cannot remove them either. + paths.push({ path: childRelative, kind: child.type, protected: true }) + continue + } + paths.push( + ...(await scanCurrent( + handle, + childEntry.absolute, + childRelative, + policy, + )), + ) + } + return relative ? [{ path: relative, kind: 'dir' }, ...paths] : paths +} + +async function scanDestination( + handle: SandboxHandle, + policy: SandboxSnapshotPolicy, + rootPath: string, +): Promise> { + const root = await lstat(handle, rootPath) + if (!root || root.type !== 'dir') + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_WORKSPACE', + 'Snapshot workspace is missing or is not a directory', + ) + assertSupported(root, rootPath) + return scanCurrent(handle, rootPath, '', policy) +} + +function parents(path: string): Array { + const values: Array = [] + const parts = path.split('/') + for (let length = 1; length < parts.length; length++) + values.push(parts.slice(0, length).join('/')) + return values +} + +function comparePath(a: string, b: string): number { + return compareBytes(new TextEncoder().encode(a), new TextEncoder().encode(b)) +} + +function compareBytes(a: Uint8Array, b: Uint8Array): number { + for (let i = 0; i < Math.min(a.length, b.length); i++) { + const left = a[i] + const right = b[i] + if (left !== right) return (left ?? 0) - (right ?? 0) + } + return a.length - b.length +} + +function getRequiredBlob( + blobs: ReadonlyMap, + key: string, +): Uint8Array { + const bytes = blobs.get(key) + if (!bytes) { + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_MISSING_BLOB', + `Missing snapshot blob '${key}'`, + ) + } + return bytes +} + +function depth(path: string): number { + return path.split('/').length +} + +function buildRestorePlan( + entries: ReadonlyArray, + current: ReadonlyArray, +): { + removes: Array + mkdirs: Array + writes: Array +} { + const desired = new Map() + for (const entry of entries) { + desired.set(entry.path, entry.kind) + for (const parent of parents(entry.path)) desired.set(parent, 'dir') + } + const currentKinds = new Map(current.map((entry) => [entry.path, entry.kind])) + const candidates = current + .filter((entry) => desired.get(entry.path) !== entry.kind) + .map((entry) => entry.path) + .sort((a, b) => depth(a) - depth(b) || comparePath(a, b)) + const removes: Array = [] + for (const path of candidates) { + const containsProtectedEntry = current.some( + (entry) => + entry.protected && + (entry.path === path || entry.path.startsWith(`${path}/`)), + ) + if (containsProtectedEntry) continue + if (removes.some((ancestor) => path.startsWith(`${ancestor}/`))) continue + removes.push(path) + } + const mkdirs = [...desired] + .filter( + ([path, kind]) => kind === 'dir' && currentKinds.get(path) !== 'dir', + ) + .map(([path]) => path) + .sort((a, b) => depth(a) - depth(b) || comparePath(a, b)) + const writes = entries + .filter((entry): entry is PlannedFile => entry.kind === 'file') + .sort((a, b) => comparePath(a.path, b.path)) + return { removes, mkdirs, writes } +} + +export async function restoreSandboxFiles( + handle: SandboxHandle, + bundle: SandboxSnapshotBundle, + snapshot: { files: ReadonlyArray }, + policy: SandboxSnapshotPolicy = defaultSandboxSnapshotPolicy(), +): Promise { + if (!handle.fs.lstat) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_LSTAT_REQUIRED', + 'Snapshot restore requires fs.lstat', + ) + const entries = validateManifest(snapshot, policy) + const rootPath = bundle.workspaceRoot ?? DEFAULT_ROOT + const current = await scanDestination(handle, policy, rootPath) + for (const entry of entries) { + if ( + entry.kind === 'file' && + current.some( + (currentEntry) => + currentEntry.protected && + currentEntry.path.startsWith(`${entry.path}/`), + ) + ) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_PATH', + `Protected current descendant conflicts with '${entry.path}'`, + ) + } + const blobs = await loadBlobs(entries, bundle) + const plan = buildRestorePlan(entries, current) + // Restore is called only while a new sandbox is private to setup. It is not + // safe for a caller that allows another process to change the workspace. + for (const path of plan.removes) await handle.fs.remove(`${rootPath}/${path}`) + for (const path of plan.mkdirs) await handle.fs.mkdir(`${rootPath}/${path}`) + for (const entry of plan.writes) + await handle.fs.write( + `${rootPath}/${entry.path}`, + getRequiredBlob(blobs, entry.blobKey), + ) +} + +export async function captureSandboxArtifacts( + bundle: SandboxSnapshotBundle, + threadId: string, + resolvedSecrets: Readonly> = {}, +): Promise> { + if (!bundle.artifacts) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_ARTIFACT_SUPPORT_REQUIRED', + 'Snapshot artifact capture requires an artifact store', + ) + const records = await bundle.artifacts.listForThread(threadId) + const loaded = new Map() + const destinationKeys = new Map() + const resolveBlobKey = + bundle.resolveArtifactBlobKey ?? + ((record: ArtifactRecord) => + record.blobKey ?? `artifacts/${record.runId}/${record.artifactId}`) + for (const record of records) { + const sourceKey = resolveBlobKey(record) + if (loaded.has(sourceKey)) continue + const source = await bundle.blobs.get(sourceKey) + if (!source) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_MISSING_ARTIFACT_BLOB', + `Missing artifact source blob '${sourceKey}'`, + ) + loaded.set(sourceKey, new Uint8Array(await source.arrayBuffer())) + } + const output = [] + for (const record of records) { + const sourceKey = resolveBlobKey(record) + let bytes = getRequiredBlob(loaded, sourceKey) + bytes = await redactBytes(bytes, resolvedSecrets) + const key = `sandbox-artifacts/sha256/${await hash(bytes)}` + if (!destinationKeys.has(key)) { + if (!(await bundle.blobs.head(key))) await bundle.blobs.put(key, bytes) + destinationKeys.set(key, key) + } + output.push({ + artifactId: record.artifactId, + name: record.name, + mimeType: record.mimeType, + size: bytes.byteLength, + blobKey: key, + createdAt: record.createdAt, + }) + } + output.sort( + (a, b) => + a.createdAt - b.createdAt || comparePath(a.artifactId, b.artifactId), + ) + return Object.freeze(output.map((artifact) => Object.freeze(artifact))) +} diff --git a/packages/ai-sandbox/src/testkit/checkpoint-conformance.ts b/packages/ai-sandbox/src/testkit/checkpoint-conformance.ts new file mode 100644 index 0000000000..c9fe00bf2a --- /dev/null +++ b/packages/ai-sandbox/src/testkit/checkpoint-conformance.ts @@ -0,0 +1,472 @@ +import { describe, expect, it } from 'vitest' +import type { + SandboxCheckpoint, + SandboxCheckpointStore, + SandboxCheckpointStoreOptions, + SandboxCheckpointWriter, +} from '../checkpoint-store' + +const writers = new WeakMap< + SandboxCheckpointStore, + Map> +>() +async function writerFor(store: SandboxCheckpointStore, threadId: string) { + let storeWriters = writers.get(store) + if (!storeWriters) { + storeWriters = new Map() + writers.set(store, storeWriters) + } + let writer = storeWriters.get(threadId) + if (!writer) { + writer = store.acquireWriter(threadId) + storeWriters.set(threadId, writer) + } + return writer +} +async function append( + store: SandboxCheckpointStore, + input: { checkpoint: SandboxCheckpoint; expectedHeadId: string | null }, +) { + return store.append({ + ...input, + writer: await writerFor(store, input.checkpoint.threadId), + }) +} +async function deleteHead( + store: SandboxCheckpointStore, + input: { threadId: string; checkpointId: string }, +) { + return store.deleteHead({ + ...input, + writer: await writerFor(store, input.threadId), + }) +} + +async function expectThreadState( + store: SandboxCheckpointStore, + staleCheckpointId: string, + expected: { + head: string | null + list: Array + checkpoint: SandboxCheckpoint | null + references: Array<{ key: string; references: number }> + }, +) { + expect(await store.getHead('thread-a')).toBe(expected.head) + expect(await store.list('thread-a')).toEqual(expected.list) + expect(await store.get(staleCheckpointId)).toEqual(expected.checkpoint) + expect(await store.listBlobReferences()).toEqual(expected.references) +} + +function checkpoint( + id: string, + parentCheckpointId: string | null = null, + threadId = 'thread-a', +): SandboxCheckpoint { + const hash = (Number(id.replace(/\D/g, '')) || 1) + .toString(16) + .padStart(64, '0') + return { + id, + threadId, + parentCheckpointId, + createdAt: Number(id.replace(/\D/g, '')) || 1, + reason: 'automatic', + files: [ + { + path: `${id}.txt`, + kind: 'file', + blobKey: `sandbox-files/sha256/${hash}`, + size: 1, + }, + ], + conversation: [{ role: 'user', content: id }], + artifacts: [], + } +} + +export function runSandboxCheckpointStoreConformance( + name: string, + makeStore: ( + options?: SandboxCheckpointStoreOptions, + ) => SandboxCheckpointStore | Promise, + options?: SandboxCheckpointStoreOptions, +): void { + describe(`SandboxCheckpointStore conformance: ${name}`, () => { + it('returns null and an empty list for an unknown thread', async () => { + const store = await makeStore(options) + expect(await store.get('missing')).toBeNull() + expect(await store.getHead('missing-thread')).toBeNull() + expect(await store.list('missing-thread')).toEqual([]) + }) + + it('appends a root and then a parent-linked checkpoint', async () => { + const store = await makeStore(options) + await expect( + append(store, { checkpoint: checkpoint('1'), expectedHeadId: null }), + ).resolves.toEqual({ headId: '1' }) + await expect( + append(store, { + checkpoint: checkpoint('2', '1'), + expectedHeadId: '1', + }), + ).resolves.toEqual({ headId: '2' }) + expect(await store.getHead('thread-a')).toBe('2') + expect((await store.list('thread-a')).map((entry) => entry.id)).toEqual([ + '1', + '2', + ]) + }) + + it('enforces expected-head and parent compare-and-swap rules', async () => { + const store = await makeStore(options) + await append(store, { checkpoint: checkpoint('1'), expectedHeadId: null }) + await expect( + append(store, { checkpoint: checkpoint('2'), expectedHeadId: null }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_STALE_HEAD' }) + await expect( + append(store, { checkpoint: checkpoint('3'), expectedHeadId: '1' }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_PARENT_MISMATCH' }) + }) + + it('rejects duplicate checkpoint ids', async () => { + const store = await makeStore(options) + await append(store, { checkpoint: checkpoint('1'), expectedHeadId: null }) + await expect( + append(store, { checkpoint: checkpoint('1'), expectedHeadId: '1' }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_DUPLICATE_ID' }) + }) + + it('isolates checkpoints and heads by thread', async () => { + const store = await makeStore(options) + await append(store, { + checkpoint: checkpoint('a', null, 'thread-a'), + expectedHeadId: null, + }) + await append(store, { + checkpoint: checkpoint('b', null, 'thread-b'), + expectedHeadId: null, + }) + expect(await store.getHead('thread-a')).toBe('a') + expect(await store.getHead('thread-b')).toBe('b') + expect((await store.list('thread-a')).map((entry) => entry.id)).toEqual([ + 'a', + ]) + expect((await store.list('thread-b')).map((entry) => entry.id)).toEqual([ + 'b', + ]) + }) + + it('lists checkpoints by createdAt and then checkpoint id', async () => { + const store = await makeStore(options) + const first = { ...checkpoint('z'), createdAt: 2 } + const second = { ...checkpoint('a', 'z'), createdAt: 1 } + const third = { ...checkpoint('b', 'a'), createdAt: 1 } + await append(store, { checkpoint: first, expectedHeadId: null }) + await append(store, { checkpoint: second, expectedHeadId: 'z' }) + await append(store, { checkpoint: third, expectedHeadId: 'a' }) + expect((await store.list('thread-a')).map((entry) => entry.id)).toEqual([ + 'a', + 'b', + 'z', + ]) + }) + + it('orders checkpoint ids and blob keys by UTF-8 bytes', async () => { + const store = await makeStore(options) + const ids = ['é', 'a', 'B', '😀'] + const keys = ['1', '2', '3', '4'].map( + (value) => `sandbox-files/sha256/${value.padStart(64, '0')}`, + ) + let parent: string | null = null + for (const [index, id] of ids.entries()) { + const blobKey = keys[index] + if (!blobKey) throw new Error(`Missing blob key for '${id}'`) + await append(store, { + checkpoint: { + ...checkpoint(id, parent), + files: [ + { + path: `${id}.txt`, + kind: 'file', + blobKey, + size: 1, + }, + ], + }, + expectedHeadId: parent, + }) + parent = id + } + expect((await store.list('thread-a')).map((entry) => entry.id)).toEqual([ + 'B', + 'a', + 'é', + '😀', + ]) + expect( + (await store.listBlobReferences()).map((entry) => entry.key), + ).toEqual([ + `sandbox-files/sha256/${'0'.repeat(63)}1`, + `sandbox-files/sha256/${'0'.repeat(63)}2`, + `sandbox-files/sha256/${'0'.repeat(63)}3`, + `sandbox-files/sha256/${'0'.repeat(63)}4`, + ]) + }) + + it('rejects malformed Unicode checkpoint ids and blob keys', async () => { + const store = await makeStore(options) + const malformed = '\uD800' + await expect( + append(store, { + checkpoint: checkpoint(malformed), + expectedHeadId: null, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ID' }) + await expect( + append(store, { + checkpoint: { + ...checkpoint('root'), + files: [ + { + path: 'root.txt', + kind: 'file', + blobKey: `sandbox-files/sha256/${malformed}${'0'.repeat(63)}`, + size: 1, + }, + ], + }, + expectedHeadId: null, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ENTRY' }) + }) + + it.each([Number.NaN, Number.POSITIVE_INFINITY, Number.NEGATIVE_INFINITY])( + 'rejects non-finite createdAt without changing state: %s', + async (createdAt) => { + const store = await makeStore(options) + const before = { + head: await store.getHead('thread-a'), + list: await store.list('thread-a'), + references: await store.listBlobReferences(), + } + await expect( + append(store, { + checkpoint: { ...checkpoint('invalid'), createdAt }, + expectedHeadId: null, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ENTRY' }) + expect(await store.getHead('thread-a')).toBe(before.head) + expect(await store.list('thread-a')).toEqual(before.list) + expect(await store.listBlobReferences()).toEqual(before.references) + }, + ) + + it('keeps checkpoints immutable and reference counts distinct blob keys', async () => { + const store = await makeStore(options) + const value = { + ...checkpoint('1'), + files: [ + ...checkpoint('1').files, + { + path: 'same.txt', + kind: 'file' as const, + blobKey: `sandbox-files/sha256/${'0'.repeat(63)}1`, + size: 1, + }, + ], + } + await append(store, { checkpoint: value, expectedHeadId: null }) + const inputFile = value.files[0] + if (!inputFile) throw new Error('Expected input file') + inputFile.path = 'input-mutated' + const loaded = await store.get('1') + expect(loaded).not.toBe(value) + if (!loaded) throw new Error('Expected stored checkpoint') + const loadedFile = loaded.files[0] + if (!loadedFile) throw new Error('Expected stored file') + expect(loadedFile.path).toBe('1.txt') + loadedFile.path = 'changed' + expect((await store.get('1'))?.files[0]?.path).toBe('1.txt') + const listed = await store.list('thread-a') + const listedCheckpoint = listed[0] + const listedFile = listedCheckpoint?.files[0] + if (!listedFile) throw new Error('Expected listed file') + listedFile.path = 'list-mutated' + expect((await store.list('thread-a'))[0]?.files[0]?.path).toBe('1.txt') + expect(await store.listBlobReferences()).toEqual([ + { key: `sandbox-files/sha256/${'0'.repeat(63)}1`, references: 1 }, + ]) + }) + + it('counts one reference per checkpoint for shared blob keys', async () => { + const store = await makeStore(options) + const shared = `sandbox-files/sha256/${'f'.repeat(64)}` + await append(store, { + checkpoint: { + ...checkpoint('1'), + files: [{ path: 'a', kind: 'file', blobKey: shared, size: 1 }], + }, + expectedHeadId: null, + }) + await append(store, { + checkpoint: { + ...checkpoint('2', '1'), + files: [{ path: 'b', kind: 'file', blobKey: shared, size: 1 }], + }, + expectedHeadId: '1', + }) + expect(await store.listBlobReferences()).toEqual([ + { key: shared, references: 2 }, + ]) + await deleteHead(store, { threadId: 'thread-a', checkpointId: '2' }) + expect(await store.listBlobReferences()).toEqual([ + { key: shared, references: 1 }, + ]) + }) + + it('allows one non-waiting writer lease per thread', async () => { + const store = await makeStore(options) + const lease = await store.acquireWriter('thread-a') + try { + await expect(store.acquireWriter('thread-a')).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_WRITER_CONFLICT', + }) + const secondLease = await store.acquireWriter('thread-b') + try { + expect(secondLease).toBeDefined() + } finally { + await secondLease.release() + } + } finally { + await lease.release() + } + const nextLease = await store.acquireWriter('thread-a') + await nextLease.release() + }) + + it('extends a lease when renew succeeds', async () => { + const config = options ?? {} + let current = config.now?.() ?? 1_000 + const store = await makeStore({ ...config, now: () => current }) + const lease = await store.acquireWriter('thread-a') + const originalExpiry = lease.expiresAt + + current += lease.renewAfterMs + await expect(lease.renew()).resolves.toEqual({ + expiresAt: current + (config.leaseDurationMs ?? 120_000), + }) + expect(lease.expiresAt).toBeGreaterThan(originalExpiry) + await lease.release() + }) + + it('does not let a stale lease release a replacement lease', async () => { + const config = options ?? {} + let current = config.now?.() ?? 1_000 + const store = await makeStore({ ...config, now: () => current }) + const old = await store.acquireWriter('thread-a') + current += (config.leaseDurationMs ?? 120_000) + 1 + const replacement = await store.acquireWriter('thread-a') + + await old.release() + await expect(store.acquireWriter('thread-a')).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_WRITER_CONFLICT', + }) + await replacement.release() + }) + + it('rejects stale leases after expiry takeover without changing state', async () => { + const config = options ?? {} + let current = config.now?.() ?? Date.now() + const takeoverStore = await makeStore({ ...config, now: () => current }) + const old = await takeoverStore.acquireWriter('thread-a') + await takeoverStore.append({ + checkpoint: checkpoint('root'), + expectedHeadId: null, + writer: old, + }) + current += (config.leaseDurationMs ?? 120_000) + 1 + const beforeExpiry = { + head: await takeoverStore.getHead('thread-a'), + list: await takeoverStore.list('thread-a'), + checkpoint: await takeoverStore.get('root'), + references: await takeoverStore.listBlobReferences(), + } + await expect(old.renew()).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_WRITER_LOST', + }) + await expect( + takeoverStore.append({ + checkpoint: checkpoint('stale', 'root'), + expectedHeadId: 'root', + writer: old, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_WRITER_LOST' }) + await expectThreadState(takeoverStore, 'stale', { + ...beforeExpiry, + checkpoint: null, + }) + await expect( + takeoverStore.deleteHead({ + threadId: 'thread-a', + checkpointId: 'root', + writer: old, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_WRITER_LOST' }) + await expectThreadState(takeoverStore, 'root', beforeExpiry) + const next = await takeoverStore.acquireWriter('thread-a') + expect(next.fence).toBeGreaterThan(old.fence) + const before = { + head: await takeoverStore.getHead('thread-a'), + list: await takeoverStore.list('thread-a'), + checkpoint: await takeoverStore.get('root'), + references: await takeoverStore.listBlobReferences(), + } + await expect(old.renew()).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_WRITER_LOST', + }) + await expect( + takeoverStore.append({ + checkpoint: checkpoint('stale', 'root'), + expectedHeadId: 'root', + writer: old, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_WRITER_LOST' }) + await expectThreadState(takeoverStore, 'stale', { + ...before, + checkpoint: null, + }) + await expect( + takeoverStore.deleteHead({ + threadId: 'thread-a', + checkpointId: 'root', + writer: old, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_WRITER_LOST' }) + await expectThreadState(takeoverStore, 'root', before) + await old.release() + await expectThreadState(takeoverStore, 'stale', { + ...before, + checkpoint: null, + }) + await next.release() + }) + + it('deletes only the current head and transitions to its parent', async () => { + const store = await makeStore(options) + await append(store, { checkpoint: checkpoint('1'), expectedHeadId: null }) + await append(store, { + checkpoint: checkpoint('2', '1'), + expectedHeadId: '1', + }) + await expect( + deleteHead(store, { threadId: 'thread-a', checkpointId: '1' }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_NOT_HEAD' }) + await deleteHead(store, { threadId: 'thread-a', checkpointId: '2' }) + expect(await store.getHead('thread-a')).toBe('1') + await deleteHead(store, { threadId: 'thread-a', checkpointId: '1' }) + expect(await store.getHead('thread-a')).toBeNull() + expect(await store.listBlobReferences()).toEqual([]) + }) + }) +} diff --git a/packages/ai-sandbox/src/testkit/checkpoint-fork-conformance.ts b/packages/ai-sandbox/src/testkit/checkpoint-fork-conformance.ts new file mode 100644 index 0000000000..ebc0f6df2b --- /dev/null +++ b/packages/ai-sandbox/src/testkit/checkpoint-fork-conformance.ts @@ -0,0 +1,299 @@ +import { describe, expect, it } from 'vitest' +import type { ModelMessage } from '@tanstack/ai' +import type { + ForkCapableSandboxCheckpointStore, + SandboxCheckpoint, + SandboxCheckpointForkInput, + SandboxCheckpointWriter, + SandboxCheckpointWriterLease, +} from '../checkpoint-store' +import { InMemorySandboxCheckpointStore } from '../checkpoint-store' + +/** Combined stores required to exercise an atomic checkpoint fork. */ +export interface SandboxCheckpointForkConformanceInput { + persistence: { + stores: { + messages: { + loadThread: (threadId: string) => Promise> + saveThread: ( + threadId: string, + messages: Array, + ) => Promise + } + } + } + checkpoints: ForkCapableSandboxCheckpointStore +} + +export interface SandboxCheckpointForkConformanceFactory { + (): + | SandboxCheckpointForkConformanceInput + | Promise +} + +type ForkInputWithoutWriter = Omit + +function sourceCheckpoint(): SandboxCheckpoint { + return { + id: 'source-root', + threadId: 'source', + parentCheckpointId: null, + createdAt: 10, + reason: 'named', + label: 'source label', + sourceRunId: 'run-1', + files: [ + { + path: 'a.txt', + kind: 'file', + blobKey: `sandbox-files/sha256/${'a'.repeat(64)}`, + size: 1, + }, + ], + conversation: [{ role: 'user', content: 'checkpoint conversation' }], + artifacts: [], + } +} + +async function appendSource( + checkpoints: ForkCapableSandboxCheckpointStore, +): Promise<{ + source: SandboxCheckpoint + sourceWriter: SandboxCheckpointWriterLease +}> { + const sourceWriter = await checkpoints.acquireWriter('source') + const source = sourceCheckpoint() + await checkpoints.append({ + checkpoint: source, + expectedHeadId: null, + writer: sourceWriter, + }) + return { source, sourceWriter } +} + +function forkInput( + writer: SandboxCheckpointWriter, + overrides: Partial = {}, +): SandboxCheckpointForkInput { + return { + sourceThreadId: 'source', + sourceCheckpointId: 'source-root', + destinationThreadId: 'destination', + destinationCheckpointId: 'fork-root', + createdAt: 20, + writer, + ...overrides, + } +} + +async function destinationState( + persistence: SandboxCheckpointForkConformanceInput['persistence'], + checkpoints: ForkCapableSandboxCheckpointStore, +) { + return { + transcript: await persistence.stores.messages.loadThread('destination'), + head: await checkpoints.getHead('destination'), + checkpoints: await checkpoints.list('destination'), + references: await checkpoints.listBlobReferences(), + } +} + +async function expectRejectedWithoutDestinationChanges( + persistence: SandboxCheckpointForkConformanceInput['persistence'], + checkpoints: ForkCapableSandboxCheckpointStore, + operation: Promise, + code: string, +): Promise { + const before = await destinationState(persistence, checkpoints) + await expect(operation).rejects.toMatchObject({ code }) + expect(await destinationState(persistence, checkpoints)).toEqual(before) +} + +export function runSandboxCheckpointForkConformance( + name: string, + makeSnapshots: SandboxCheckpointForkConformanceFactory, +): void { + describe(`Sandbox checkpoint fork conformance: ${name}`, () => { + it('copies a selected historical checkpoint and creates an exact fork root', async () => { + const { persistence, checkpoints } = await makeSnapshots() + const { source, sourceWriter } = await appendSource(checkpoints) + await checkpoints.append({ + checkpoint: { + id: 'source-head', + threadId: 'source', + parentCheckpointId: source.id, + createdAt: 11, + reason: 'automatic', + files: [], + conversation: [{ role: 'user', content: 'newer conversation' }], + artifacts: [], + }, + expectedHeadId: source.id, + writer: sourceWriter, + }) + await persistence.stores.messages.saveThread('source', [ + { role: 'user', content: 'current source message' }, + ]) + const destinationWriter = await checkpoints.acquireWriter('destination') + const result = await checkpoints.forkFromCheckpoint( + forkInput(destinationWriter), + ) + expect(result.checkpoint).toEqual({ + id: 'fork-root', + threadId: 'destination', + parentCheckpointId: null, + createdAt: 20, + reason: 'fork-root', + files: source.files, + conversation: source.conversation, + artifacts: [], + }) + expect( + await persistence.stores.messages.loadThread('destination'), + ).toEqual(source.conversation) + expect(await checkpoints.getHead('source')).toBe('source-head') + }) + + it('rejects a plain checkpoint store because it has no fork capability', () => { + const store = new InMemorySandboxCheckpointStore() + expect('forkFromCheckpoint' in store).toBe(false) + }) + + it('rejects a missing source without changing the destination', async () => { + const { persistence, checkpoints } = await makeSnapshots() + const writer = await checkpoints.acquireWriter('destination') + await expectRejectedWithoutDestinationChanges( + persistence, + checkpoints, + checkpoints.forkFromCheckpoint( + forkInput(writer, { sourceCheckpointId: 'missing' }), + ), + 'SANDBOX_SNAPSHOT_FORK_SOURCE_NOT_FOUND', + ) + }) + + it('rejects a source thread mismatch without changing the destination', async () => { + const { persistence, checkpoints } = await makeSnapshots() + await appendSource(checkpoints) + const writer = await checkpoints.acquireWriter('destination') + await expectRejectedWithoutDestinationChanges( + persistence, + checkpoints, + checkpoints.forkFromCheckpoint( + forkInput(writer, { sourceThreadId: 'another-source' }), + ), + 'SANDBOX_SNAPSHOT_FORK_SOURCE_THREAD_MISMATCH', + ) + }) + + it('rejects equal source and destination threads without changing the source', async () => { + const { persistence, checkpoints } = await makeSnapshots() + const { source, sourceWriter } = await appendSource(checkpoints) + const before = { + transcript: await persistence.stores.messages.loadThread('source'), + head: await checkpoints.getHead('source'), + checkpoints: await checkpoints.list('source'), + references: await checkpoints.listBlobReferences(), + } + await expect( + checkpoints.forkFromCheckpoint( + forkInput(sourceWriter, { + destinationThreadId: source.threadId, + }), + ), + ).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_FORK_SOURCE_THREAD_MISMATCH', + }) + expect({ + transcript: await persistence.stores.messages.loadThread('source'), + head: await checkpoints.getHead('source'), + checkpoints: await checkpoints.list('source'), + references: await checkpoints.listBlobReferences(), + }).toEqual(before) + }) + + it('rejects a missing destination writer without changing the destination', async () => { + const { persistence, checkpoints } = await makeSnapshots() + await appendSource(checkpoints) + const missingWriter: SandboxCheckpointWriter = { + threadId: 'destination', + ownerToken: 'missing-owner', + fence: 1, + } + await expectRejectedWithoutDestinationChanges( + persistence, + checkpoints, + checkpoints.forkFromCheckpoint(forkInput(missingWriter)), + 'SANDBOX_SNAPSHOT_WRITER_LOST', + ) + }) + + it('rejects a wrong destination writer without changing the destination', async () => { + const { persistence, checkpoints } = await makeSnapshots() + await appendSource(checkpoints) + const wrongWriter = await checkpoints.acquireWriter('another-thread') + await expectRejectedWithoutDestinationChanges( + persistence, + checkpoints, + checkpoints.forkFromCheckpoint(forkInput(wrongWriter)), + 'SANDBOX_SNAPSHOT_WRITER_LOST', + ) + }) + + it('rejects a stale destination writer without changing the destination', async () => { + const { persistence, checkpoints } = await makeSnapshots() + await appendSource(checkpoints) + const staleWriter = await checkpoints.acquireWriter('destination') + await staleWriter.release() + await expectRejectedWithoutDestinationChanges( + persistence, + checkpoints, + checkpoints.forkFromCheckpoint(forkInput(staleWriter)), + 'SANDBOX_SNAPSHOT_WRITER_LOST', + ) + }) + + it('rejects a destination transcript without changing its checkpoint state', async () => { + const { persistence, checkpoints } = await makeSnapshots() + await appendSource(checkpoints) + await persistence.stores.messages.saveThread('destination', [ + { role: 'user', content: 'already here' }, + ]) + const writer = await checkpoints.acquireWriter('destination') + await expectRejectedWithoutDestinationChanges( + persistence, + checkpoints, + checkpoints.forkFromCheckpoint(forkInput(writer)), + 'SANDBOX_SNAPSHOT_FORK_DESTINATION_NOT_EMPTY', + ) + }) + + it('rejects an orphaned checkpoint id without changing the destination', async () => { + const { persistence, checkpoints } = await makeSnapshots() + const { source } = await appendSource(checkpoints) + const orphanWriter = await checkpoints.acquireWriter('orphan-thread') + await checkpoints.append({ + checkpoint: { + ...source, + id: 'orphaned-id', + threadId: 'orphan-thread', + parentCheckpointId: null, + createdAt: 11, + }, + expectedHeadId: null, + writer: orphanWriter, + }) + const destinationWriter = await checkpoints.acquireWriter('destination') + await expectRejectedWithoutDestinationChanges( + persistence, + checkpoints, + checkpoints.forkFromCheckpoint( + forkInput(destinationWriter, { + destinationCheckpointId: 'orphaned-id', + }), + ), + 'SANDBOX_SNAPSHOT_FORK_DESTINATION_NOT_EMPTY', + ) + }) + }) +} diff --git a/packages/ai-sandbox/src/testkit/conformance.ts b/packages/ai-sandbox/src/testkit/conformance.ts index 0086687b39..46d921be6a 100644 --- a/packages/ai-sandbox/src/testkit/conformance.ts +++ b/packages/ai-sandbox/src/testkit/conformance.ts @@ -29,6 +29,13 @@ export type { ReaperConformanceConfig } from './reaper-conformance' export { runDurableRunFieldsConformance } from './durable-run-fields-conformance' export type { MakeRunStore } from './durable-run-fields-conformance' export { makeFakeShellSpawn } from './shell-spawn' +export { runSandboxCheckpointStoreConformance } from './checkpoint-conformance' +export type { SandboxCheckpointStoreOptions } from '../checkpoint-store' +export { runSandboxCheckpointForkConformance } from './checkpoint-fork-conformance' +export type { + SandboxCheckpointForkConformanceInput, + SandboxCheckpointForkConformanceFactory, +} from './checkpoint-fork-conformance' function makeRecord( overrides?: Partial, diff --git a/packages/ai-sandbox/tests/ai-middleware-subpath.test.ts b/packages/ai-sandbox/tests/ai-middleware-subpath.test.ts new file mode 100644 index 0000000000..7e1b45eaf5 --- /dev/null +++ b/packages/ai-sandbox/tests/ai-middleware-subpath.test.ts @@ -0,0 +1,6 @@ +import { expect, it } from 'vitest' +import { CapabilityRegistry } from '@tanstack/ai/middlewares' + +it('exports a working CapabilityRegistry from the built middlewares subpath', () => { + expect(new CapabilityRegistry()).toBeInstanceOf(CapabilityRegistry) +}) diff --git a/packages/ai-sandbox/tests/checkpoint-store.conformance.test.ts b/packages/ai-sandbox/tests/checkpoint-store.conformance.test.ts new file mode 100644 index 0000000000..0aba58c279 --- /dev/null +++ b/packages/ai-sandbox/tests/checkpoint-store.conformance.test.ts @@ -0,0 +1,18 @@ +import { InMemorySandboxCheckpointStore } from '../src/checkpoint-store' +import { runSandboxCheckpointStoreConformance } from '../src/testkit/checkpoint-conformance' +import { runSandboxCheckpointForkConformance } from '../src/testkit/checkpoint-fork-conformance' +import { memorySandboxSnapshots } from '../src/memory-snapshots' +import { runPersistenceConformance } from '@tanstack/ai-persistence/testkit' + +let now = 1_000 +runSandboxCheckpointStoreConformance( + 'in-memory reference', + (options) => new InMemorySandboxCheckpointStore(options), + { now: () => now, leaseDurationMs: 100, renewAfterMs: 25 }, +) + +runSandboxCheckpointForkConformance('memory snapshots', memorySandboxSnapshots) +runPersistenceConformance('memory snapshots persistence', async () => { + const { persistence } = await memorySandboxSnapshots() + return persistence +}) diff --git a/packages/ai-sandbox/tests/checkpoint-store.test.ts b/packages/ai-sandbox/tests/checkpoint-store.test.ts new file mode 100644 index 0000000000..b25b586ded --- /dev/null +++ b/packages/ai-sandbox/tests/checkpoint-store.test.ts @@ -0,0 +1,677 @@ +import { describe, expect, it } from 'vitest' +import { + InMemorySandboxCheckpointStore, + SandboxCheckpointConflictError, + SandboxCheckpointDuplicateIdError, + SandboxCheckpointInvalidEntryError, + SandboxCheckpointInvalidIdError, + SandboxCheckpointNotHeadError, + SandboxCheckpointParentMismatchError, + SandboxCheckpointWriterConflictError, +} from '../src/checkpoint-store' +import type { + SandboxCheckpoint, + SandboxCheckpointStore, + SandboxCheckpointWriter, +} from '../src/checkpoint-store' +import { memorySandboxSnapshots } from '../src/memory-snapshots' + +const writers = new WeakMap< + SandboxCheckpointStore, + Map> +>() +const validFileKey = `sandbox-files/sha256/${'0'.repeat(64)}` +async function writerFor(store: SandboxCheckpointStore, threadId: string) { + let storeWriters = writers.get(store) + if (!storeWriters) { + storeWriters = new Map() + writers.set(store, storeWriters) + } + let writer = storeWriters.get(threadId) + if (!writer) { + writer = store.acquireWriter(threadId) + storeWriters.set(threadId, writer) + } + return writer +} +async function append( + store: SandboxCheckpointStore, + input: { checkpoint: SandboxCheckpoint; expectedHeadId: string | null }, +) { + return store.append({ + ...input, + writer: await writerFor(store, input.checkpoint.threadId), + }) +} +async function deleteHead( + store: SandboxCheckpointStore, + input: { threadId: string; checkpointId: string }, +) { + return store.deleteHead({ + ...input, + writer: await writerFor(store, input.threadId), + }) +} + +async function expectState( + store: SandboxCheckpointStore, + staleCheckpointId: string, + expected: { + head: string | null + list: Array + checkpoint: SandboxCheckpoint | null + references: Array<{ key: string; references: number }> + }, +) { + expect(await store.getHead('thread-a')).toBe(expected.head) + expect(await store.list('thread-a')).toEqual(expected.list) + expect(await store.get(staleCheckpointId)).toEqual(expected.checkpoint) + expect(await store.listBlobReferences()).toEqual(expected.references) +} + +function checkpoint( + id: string, + parentCheckpointId: string | null = null, + threadId = 'thread-a', +): SandboxCheckpoint { + const hash = (Number(id.replace(/\D/g, '')) || 1) + .toString(16) + .padStart(64, '0') + return { + id, + threadId, + parentCheckpointId, + createdAt: Number(id.replace(/\D/g, '')) || 1, + reason: 'automatic', + files: [ + { + path: 'file.txt', + kind: 'file', + blobKey: `sandbox-files/sha256/${hash}`, + size: 1, + }, + ], + conversation: [{ role: 'user', content: id }], + artifacts: [ + { + artifactId: `artifact-${id}`, + name: 'file.txt', + mimeType: 'text/plain', + size: 1, + blobKey: `sandbox-artifacts/sha256/${hash}`, + createdAt: 1, + }, + ], + } +} + +describe('InMemorySandboxCheckpointStore', () => { + it('acquires a fenced lease and rejects a second active writer', async () => { + const store = new InMemorySandboxCheckpointStore() + const lease = await store.acquireWriter('thread-a') + expect(lease.threadId).toBe('thread-a') + expect(lease.ownerToken).toBeTruthy() + expect(lease.fence).toBe(1) + expect(lease.renewAfterMs).toBeLessThan(lease.expiresAt - Date.now()) + await expect(store.acquireWriter('thread-a')).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_WRITER_CONFLICT', + }) + await lease.release() + }) + + it.each([ + { leaseDurationMs: 0 }, + { leaseDurationMs: Number.NaN }, + { renewAfterMs: 0 }, + { renewAfterMs: Number.NaN }, + { leaseDurationMs: 100, renewAfterMs: 100 }, + ])('rejects invalid lease timing options %j', (options) => { + expect(() => new InMemorySandboxCheckpointStore(options)).toThrow() + }) + + it('rejects a writer from another thread without changing state', async () => { + const store = new InMemorySandboxCheckpointStore() + const writer = await store.acquireWriter('thread-a') + const before = await store.listBlobReferences() + await expect( + store.append({ + checkpoint: checkpoint('root', null, 'thread-b'), + expectedHeadId: null, + writer, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_WRITER_LOST' }) + await expect( + store.deleteHead({ threadId: 'thread-b', checkpointId: 'root', writer }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_WRITER_LOST' }) + expect(await store.getHead('thread-a')).toBeNull() + expect(await store.listBlobReferences()).toEqual(before) + await writer.release() + }) + + it('appends only when expectedHeadId matches and increments references', async () => { + const store = new InMemorySandboxCheckpointStore() + + await expect( + append(store, { checkpoint: checkpoint('a'), expectedHeadId: null }), + ).resolves.toEqual({ headId: 'a' }) + await expect( + append(store, { checkpoint: checkpoint('b'), expectedHeadId: null }), + ).rejects.toBeInstanceOf(SandboxCheckpointConflictError) + expect(await store.listBlobReferences()).toEqual([ + { key: `sandbox-artifacts/sha256/${'0'.repeat(63)}1`, references: 1 }, + { key: `sandbox-files/sha256/${'0'.repeat(63)}1`, references: 1 }, + ]) + }) + + it('rejects a checkpoint whose parent differs from the expected head', async () => { + const store = new InMemorySandboxCheckpointStore() + await append(store, { checkpoint: checkpoint('a'), expectedHeadId: null }) + + await expect( + append(store, { checkpoint: checkpoint('b'), expectedHeadId: 'a' }), + ).rejects.toBeInstanceOf(SandboxCheckpointParentMismatchError) + }) + + it('rejects duplicate checkpoint ids', async () => { + const store = new InMemorySandboxCheckpointStore() + await append(store, { checkpoint: checkpoint('a'), expectedHeadId: null }) + await expect( + append(store, { checkpoint: checkpoint('a'), expectedHeadId: 'a' }), + ).rejects.toBeInstanceOf(SandboxCheckpointDuplicateIdError) + }) + + it('rejects empty checkpoint and parent ids', async () => { + const store = new InMemorySandboxCheckpointStore() + await expect( + append(store, { checkpoint: checkpoint(''), expectedHeadId: null }), + ).rejects.toBeInstanceOf(SandboxCheckpointInvalidIdError) + await expect( + append(store, { checkpoint: checkpoint('a', ''), expectedHeadId: null }), + ).rejects.toBeInstanceOf(SandboxCheckpointInvalidIdError) + await expect( + append(store, { checkpoint: checkpoint('b'), expectedHeadId: '' }), + ).rejects.toBeInstanceOf(SandboxCheckpointInvalidIdError) + }) + + it('rejects malformed Unicode checkpoint identities and blob keys', async () => { + const store = new InMemorySandboxCheckpointStore() + const malformed = '\uD800' + await expect( + append(store, { + checkpoint: checkpoint(malformed), + expectedHeadId: null, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ID' }) + await expect( + append(store, { + checkpoint: { + ...checkpoint('root'), + files: [ + { + path: 'file.txt', + kind: 'file', + blobKey: `sandbox-files/sha256/${malformed}${'0'.repeat(63)}`, + size: 1, + }, + ], + }, + expectedHeadId: null, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ENTRY' }) + expect(await store.getHead('thread-a')).toBeNull() + expect(await store.listBlobReferences()).toEqual([]) + }) + + it('rejects empty and non-string identifiers without changing state', async () => { + const store = new InMemorySandboxCheckpointStore() + const writer = await store.acquireWriter('thread-a') + const before = { + head: await store.getHead('thread-a'), + list: await store.list('thread-a'), + checkpoint: await store.get('root'), + references: await store.listBlobReferences(), + } + await expect( + Reflect.apply(store.acquireWriter, store, ['']), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ID' }) + await expect(Reflect.apply(store.list, store, [{}])).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_INVALID_ID', + }) + await expect( + Reflect.apply(store.getHead, store, ['']), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ID' }) + await expect(Reflect.apply(store.get, store, [''])).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_INVALID_ID', + }) + await expect( + Reflect.apply(store.append, store, [ + { + checkpoint: { ...checkpoint('invalid'), threadId: 7 }, + expectedHeadId: null, + writer, + }, + ]), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ID' }) + await expect( + Reflect.apply(store.deleteHead, store, [ + { threadId: '', checkpointId: 'missing', writer }, + ]), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ID' }) + expect(await store.getHead('thread-a')).toBe(before.head) + expect(await store.list('thread-a')).toEqual(before.list) + expect(await store.listBlobReferences()).toEqual(before.references) + await writer.release() + }) + + it.each([Number.NaN, Number.POSITIVE_INFINITY, Number.NEGATIVE_INFINITY])( + 'rejects non-finite createdAt values without changing state: %s', + async (createdAt) => { + const store = new InMemorySandboxCheckpointStore() + const writer = await store.acquireWriter('thread-a') + const before = { + head: await store.getHead('thread-a'), + list: await store.list('thread-a'), + references: await store.listBlobReferences(), + } + await expect( + store.append({ + checkpoint: { ...checkpoint('invalid'), createdAt }, + expectedHeadId: null, + writer, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ENTRY' }) + expect(await store.getHead('thread-a')).toBe(before.head) + expect(await store.list('thread-a')).toEqual(before.list) + expect(await store.listBlobReferences()).toEqual(before.references) + await writer.release() + }, + ) + + it('treats an absent parent id as null', async () => { + const store = new InMemorySandboxCheckpointStore() + const root = checkpoint('root') + Reflect.deleteProperty(root, 'parentCheckpointId') + await Reflect.apply(append, undefined, [ + store, + { checkpoint: root, expectedHeadId: null }, + ]) + expect((await store.get('root'))?.parentCheckpointId).toBeNull() + }) + + it('rejects malformed kind-specific file and directory entries', async () => { + const store = new InMemorySandboxCheckpointStore() + const malformed = (entry: unknown) => + Reflect.apply(append, undefined, [ + store, + { + checkpoint: { ...checkpoint('malformed'), files: [entry] }, + expectedHeadId: null, + }, + ]) + await expect( + malformed({ path: 'file.txt', kind: 'file' }), + ).rejects.toBeInstanceOf(SandboxCheckpointInvalidEntryError) + await expect( + malformed({ + path: 'empty', + kind: 'dir', + blobKey: 'sandbox-files/sha256/dir', + }), + ).rejects.toBeInstanceOf(SandboxCheckpointInvalidEntryError) + }) + + const invalidEntries: Array<{ name: string; entry: unknown }> = [ + { + name: 'missing size', + entry: { + path: 'file.txt', + kind: 'file' as const, + blobKey: 'sandbox-files/sha256/' + '0'.repeat(64), + }, + }, + { + name: 'fractional size', + entry: { + path: 'file.txt', + kind: 'file' as const, + blobKey: 'sandbox-files/sha256/' + '0'.repeat(64), + size: 1.5, + }, + }, + ] + it.each(invalidEntries)( + 'rejects $name without changing checkpoint state', + async ({ entry }) => { + const store = new InMemorySandboxCheckpointStore() + const writer = await store.acquireWriter('thread-a') + const before = { + head: await store.getHead('thread-a'), + list: await store.list('thread-a'), + references: await store.listBlobReferences(), + } + const invalidCheckpoint = checkpoint('invalid') + Reflect.defineProperty(invalidCheckpoint, 'files', { value: [entry] }) + await expect( + store.append({ + checkpoint: invalidCheckpoint, + expectedHeadId: null, + writer, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ENTRY' }) + expect(await store.getHead('thread-a')).toBe(before.head) + expect(await store.list('thread-a')).toEqual(before.list) + expect(await store.listBlobReferences()).toEqual(before.references) + await writer.release() + }, + ) + + it.each([ + 'bad\0name', + 'nested/../bad', + 'nested//bad', + '/absolute', + 'C:/absolute', + ])( + 'rejects unsafe checkpoint path %j without changing head, list, or refs', + async (path) => { + const store = new InMemorySandboxCheckpointStore() + const writer = await store.acquireWriter('thread-a') + const before = { + head: await store.getHead('thread-a'), + list: await store.list('thread-a'), + references: await store.listBlobReferences(), + } + await expect( + store.append({ + checkpoint: { + ...checkpoint('invalid'), + files: [ + { + path, + kind: 'file', + blobKey: 'sandbox-files/sha256/' + '0'.repeat(64), + size: 1, + }, + ], + }, + expectedHeadId: null, + writer, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ENTRY' }) + expect(await store.getHead('thread-a')).toBe(before.head) + expect(await store.list('thread-a')).toEqual(before.list) + expect(await store.listBlobReferences()).toEqual(before.references) + await writer.release() + }, + ) + + it.each([ + 'bad', + 'sandbox-files/sha256/not-hex', + 'sandbox-files/sha256/' + 'f'.repeat(63), + ])( + 'rejects invalid content-addressed blob key %j without mutation', + async (blobKey) => { + const store = new InMemorySandboxCheckpointStore() + const writer = await store.acquireWriter('thread-a') + await expect( + store.append({ + checkpoint: { + ...checkpoint('invalid'), + files: [{ path: 'file.txt', kind: 'file', blobKey, size: 1 }], + }, + expectedHeadId: null, + writer, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ENTRY' }) + expect(await store.getHead('thread-a')).toBeNull() + expect(await store.list('thread-a')).toEqual([]) + expect(await store.listBlobReferences()).toEqual([]) + await writer.release() + }, + ) + + it.each([ + '/absolute.txt', + '\\absolute.txt', + 'C:/absolute.txt', + '../outside.txt', + 'nested/../../outside.txt', + 'nested//file.txt', + './file.txt', + 'nested/', + '', + ])( + 'rejects ambiguous or non-workspace path %j without changing state', + async (path) => { + const store = new InMemorySandboxCheckpointStore() + await append(store, { + checkpoint: checkpoint('root'), + expectedHeadId: null, + }) + const before = { + head: await store.getHead('thread-a'), + list: await store.list('thread-a'), + refs: await store.listBlobReferences(), + } + await expect( + append(store, { + checkpoint: { + ...checkpoint('bad', 'root'), + files: [{ path, kind: 'file', blobKey: validFileKey, size: 1 }], + }, + expectedHeadId: 'root', + }), + ).rejects.toBeInstanceOf(SandboxCheckpointInvalidEntryError) + expect(await store.getHead('thread-a')).toBe(before.head) + expect(await store.list('thread-a')).toEqual(before.list) + expect(await store.listBlobReferences()).toEqual(before.refs) + }, + ) + + it('rejects duplicate paths and file-ancestor conflicts atomically', async () => { + const store = new InMemorySandboxCheckpointStore() + await append(store, { + checkpoint: checkpoint('root'), + expectedHeadId: null, + }) + const before = await store.listBlobReferences() + for (const files of [ + [ + { path: 'same', kind: 'file' as const, blobKey: validFileKey, size: 1 }, + { path: 'same', kind: 'dir' as const }, + ], + [ + { path: 'a', kind: 'file' as const, blobKey: validFileKey, size: 1 }, + { path: 'a/b', kind: 'file' as const, blobKey: validFileKey, size: 1 }, + ], + [ + { path: 'a/b', kind: 'file' as const, blobKey: validFileKey, size: 1 }, + { path: 'a', kind: 'file' as const, blobKey: validFileKey, size: 1 }, + ], + ]) { + await expect( + append(store, { + checkpoint: { ...checkpoint('bad', 'root'), files }, + expectedHeadId: 'root', + }), + ).rejects.toBeInstanceOf(SandboxCheckpointInvalidEntryError) + expect(await store.getHead('thread-a')).toBe('root') + expect(await store.list('thread-a')).toHaveLength(1) + expect(await store.listBlobReferences()).toEqual(before) + } + }) + + it('rejects malformed artifacts atomically', async () => { + const store = new InMemorySandboxCheckpointStore() + await append(store, { + checkpoint: checkpoint('root'), + expectedHeadId: null, + }) + const beforeHead = await store.getHead('thread-a') + const beforeList = await store.list('thread-a') + const beforeReferences = await store.listBlobReferences() + const malformed = { + ...checkpoint('bad', 'root'), + artifacts: [ + { + artifactId: '', + name: 'bad.txt', + mimeType: 'text/plain', + size: 1, + blobKey: 'sandbox-artifacts/sha256/bad', + createdAt: 2, + }, + ], + } + + await expect( + append(store, { checkpoint: malformed, expectedHeadId: 'root' }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ENTRY' }) + expect(await store.getHead('thread-a')).toBe(beforeHead) + expect(await store.list('thread-a')).toEqual(beforeList) + expect(await store.listBlobReferences()).toEqual(beforeReferences) + }) + + it('stores immutable checkpoint copies and lists by creation time', async () => { + const store = new InMemorySandboxCheckpointStore() + const first = checkpoint('a') + await append(store, { checkpoint: first, expectedHeadId: null }) + first.files[0]!.path = 'mutated' + + const loaded = await store.get('a') + expect(loaded?.files[0]?.path).toBe('file.txt') + expect(loaded).not.toBe(first) + expect(await store.list('thread-a')).toEqual([loaded]) + }) + + it('rejects a second lease without waiting and releases the first lease', async () => { + const store = new InMemorySandboxCheckpointStore() + const lease = await store.acquireWriter('thread-a') + try { + await expect(store.acquireWriter('thread-a')).rejects.toBeInstanceOf( + SandboxCheckpointWriterConflictError, + ) + const secondLease = await store.acquireWriter('thread-b') + try { + expect(secondLease).toBeDefined() + } finally { + await secondLease.release() + } + } finally { + await lease.release() + } + const nextLease = await store.acquireWriter('thread-a') + await nextLease.release() + }) + + it('takes over an expired lease with a higher fence and preserves state on stale writes', async () => { + let now = 1_000 + const store = new InMemorySandboxCheckpointStore({ + now: () => now, + leaseDurationMs: 100, + renewAfterMs: 25, + }) + const first = await store.acquireWriter('thread-a') + await store.append({ + checkpoint: checkpoint('root'), + expectedHeadId: null, + writer: first, + }) + now += 101 + const second = await store.acquireWriter('thread-a') + expect(second.fence).toBeGreaterThan(first.fence) + const before = { + head: await store.getHead('thread-a'), + list: await store.list('thread-a'), + checkpoint: await store.get('root'), + references: await store.listBlobReferences(), + } + await expect(first.renew()).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_WRITER_LOST', + }) + await expect( + store.append({ + checkpoint: checkpoint('stale', 'root'), + expectedHeadId: 'root', + writer: first, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_WRITER_LOST' }) + await expectState(store, 'stale', { ...before, checkpoint: null }) + await expect( + store.deleteHead({ + threadId: 'thread-a', + checkpointId: 'root', + writer: first, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_WRITER_LOST' }) + await expectState(store, 'root', before) + await first.release() + await expectState(store, 'stale', { ...before, checkpoint: null }) + const renewed = await second.renew() + expect(renewed.expiresAt).toBe(now + 100) + expect(second.expiresAt).toBe(renewed.expiresAt) + now += 99 + await expect(store.acquireWriter('thread-a')).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_WRITER_CONFLICT', + }) + await second.release() + }) + + it('deletes only the expected current head and moves it to its parent', async () => { + const store = new InMemorySandboxCheckpointStore() + await append(store, { checkpoint: checkpoint('a'), expectedHeadId: null }) + await append(store, { + checkpoint: checkpoint('b', 'a'), + expectedHeadId: 'a', + }) + + await expect( + deleteHead(store, { threadId: 'thread-a', checkpointId: 'a' }), + ).rejects.toBeInstanceOf(SandboxCheckpointNotHeadError) + await deleteHead(store, { threadId: 'thread-a', checkpointId: 'b' }) + expect(await store.getHead('thread-a')).toBe('a') + expect(await store.get('b')).toBeNull() + expect(await store.listBlobReferences()).toHaveLength(2) + }) + + it('deletes the root head and clears the thread head', async () => { + const store = new InMemorySandboxCheckpointStore() + await append(store, { checkpoint: checkpoint('a'), expectedHeadId: null }) + await deleteHead(store, { threadId: 'thread-a', checkpointId: 'a' }) + expect(await store.getHead('thread-a')).toBeNull() + expect(await store.list('thread-a')).toEqual([]) + expect(await store.listBlobReferences()).toEqual([]) + }) + + it('rejects reentrant outer append for both checkpoint stores', async () => { + const memory = await memorySandboxSnapshots() + const stores: Array = [ + new InMemorySandboxCheckpointStore(), + memory.checkpoints, + ] + for (const store of stores) { + const writer = await store.acquireWriter('thread-a') + let nested: Promise | undefined + const outer = checkpoint('outer') + Reflect.defineProperty(outer, 'files', { + enumerable: true, + get: () => { + nested = store.append({ + checkpoint: { ...checkpoint('nested'), files: [], artifacts: [] }, + expectedHeadId: null, + writer, + }) + return [] + }, + }) + await expect( + store.append({ checkpoint: outer, expectedHeadId: null, writer }), + ).rejects.toBeInstanceOf(SandboxCheckpointConflictError) + await nested + expect(await store.getHead('thread-a')).toBe('nested') + expect(await store.listBlobReferences()).toEqual([]) + await writer.release() + } + }) +}) diff --git a/packages/ai-sandbox/tests/fakes.test.ts b/packages/ai-sandbox/tests/fakes.test.ts new file mode 100644 index 0000000000..b0b0fb6b64 --- /dev/null +++ b/packages/ai-sandbox/tests/fakes.test.ts @@ -0,0 +1,17 @@ +import { describe, expect, it } from 'vitest' +import { createCapability } from '@tanstack/ai' +import { makeMiddlewareCtx } from './fakes' + +describe('makeMiddlewareCtx', () => { + it('tracks provided capabilities and stores their values', () => { + const ctx = makeMiddlewareCtx({ threadId: 'thread-1', runId: 'run-1' }) + const capability = createCapability<{ value: string }>()('fake-capability') + const value = { value: 'stored' } + + expect(ctx.capabilities.has(capability)).toBe(false) + ctx.provide(capability, value) + + expect(ctx.capabilities.has(capability)).toBe(true) + expect(ctx.get(capability)).toBe(value) + }) +}) diff --git a/packages/ai-sandbox/tests/fakes.ts b/packages/ai-sandbox/tests/fakes.ts index e398918bf1..249fec95b2 100644 --- a/packages/ai-sandbox/tests/fakes.ts +++ b/packages/ai-sandbox/tests/fakes.ts @@ -1,4 +1,5 @@ import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { CapabilityRegistry } from '@tanstack/ai/middlewares' import { makeFakeShellSpawn } from '../src/testkit/shell-spawn' import type { InternalLogger } from '@tanstack/ai/adapter-internals' import type { @@ -58,6 +59,7 @@ export function makeFakeHandle( return Promise.resolve() }, list: () => Promise.resolve([]), + lstat: () => Promise.resolve({ type: 'dir' as const, mode: 0 }), mkdir: () => Promise.resolve(), remove: (p) => { files.delete(p) @@ -215,16 +217,8 @@ export function captureLogger(): { * asserting on a provided capability is not exercising a stub that always * answers the same way. * - * The one exception is the `capabilities` field itself: production types it - * as the `CapabilityRegistry` class (`packages/ai/src/activities/chat/ - * middleware/capabilities.ts`), which is not exported from any public - * `@tanstack/ai` subpath (not `.`, not `/adapter-internals`), so it cannot be - * `new`'d — or even named — from this package. Every consumer of a provided - * capability only ever calls `markProvided`/`has` on that field (see - * `capabilities.ts`'s `provide`), never its private bookkeeping, so the - * minimal stand-in below is functionally equivalent for anything this fake is - * used for; only the cast on that one field is needed, not on the ctx as a - * whole. + * The capabilities field uses the public `CapabilityRegistry` export, so this + * fake exercises the same registry implementation as production code. */ export function makeMiddlewareCtx(input: { threadId: string @@ -260,11 +254,7 @@ export function makeMiddlewareCtx(input: { messages: [], createId: (prefix: string) => `${prefix}-${Math.random().toString(36).slice(2)}`, - capabilities: { - markProvided: () => {}, - has: () => false, - setOnDuplicate: () => {}, - } as unknown as ChatMiddlewareContext['capabilities'], + capabilities: new CapabilityRegistry(), get: (capability) => capability[0](ctx), getOptional: (capability) => capability[0](ctx, { optional: true }), provide: (capability, value) => capability[1](ctx, value), diff --git a/packages/ai-sandbox/tests/memory-snapshots-declaration.test-d.ts b/packages/ai-sandbox/tests/memory-snapshots-declaration.test-d.ts new file mode 100644 index 0000000000..8511956dd4 --- /dev/null +++ b/packages/ai-sandbox/tests/memory-snapshots-declaration.test-d.ts @@ -0,0 +1,29 @@ +import { expectTypeOf } from 'vitest' +import type { MemorySandboxSnapshots } from '../src' +import type { + AIPersistence, + ArtifactStore, + BlobStore, + GenerationRunStore, + InterruptStore, + MessageStore, + MetadataStore, + RunStore, +} from '@tanstack/ai-persistence' + +type MemoryPersistenceStores = { + messages: MessageStore + runs: RunStore + generationRuns: GenerationRunStore + interrupts: InterruptStore + metadata: MetadataStore + artifacts: ArtifactStore + blobs: BlobStore +} + +declare const snapshots: MemorySandboxSnapshots +expectTypeOf(snapshots.persistence).toMatchTypeOf< + AIPersistence +>() +// @ts-expect-error immutable identity fields are not patchable +snapshots.persistence.stores.generationRuns.update('run', { threadId: 'other' }) diff --git a/packages/ai-sandbox/tests/memory-snapshots-import.test.ts b/packages/ai-sandbox/tests/memory-snapshots-import.test.ts new file mode 100644 index 0000000000..2bbab650de --- /dev/null +++ b/packages/ai-sandbox/tests/memory-snapshots-import.test.ts @@ -0,0 +1,16 @@ +import { describe, expect, it, vi } from 'vitest' + +vi.mock('@tanstack/ai-persistence', () => { + throw new Error('memorySandboxSnapshots must not load the persistence root') +}) + +const { memorySandboxSnapshots } = await import('../src/memory-snapshots') + +describe('memorySandboxSnapshots runtime dependencies', () => { + it('does not load the ai-persistence root when creating a store', async () => { + await expect(memorySandboxSnapshots()).resolves.toMatchObject({ + persistence: expect.any(Object), + checkpoints: expect.any(Object), + }) + }) +}) diff --git a/packages/ai-sandbox/tests/memory-snapshots.behavior.test.ts b/packages/ai-sandbox/tests/memory-snapshots.behavior.test.ts new file mode 100644 index 0000000000..c9ea8ccfd4 --- /dev/null +++ b/packages/ai-sandbox/tests/memory-snapshots.behavior.test.ts @@ -0,0 +1,744 @@ +import { readFileSync } from 'node:fs' +import { describe, expect, it, vi } from 'vitest' +import type { + ForkCapableSandboxCheckpointStore, + SandboxCheckpoint, + SandboxCheckpointForkInput, + SandboxCheckpointWriter, +} from '../src/checkpoint-store' +import { memorySandboxSnapshots } from '../src/memory-snapshots' + +type MemoryPersistence = Awaited< + ReturnType +>['persistence'] + +function rawCheckpointRecords( + checkpoints: ForkCapableSandboxCheckpointStore, +): Map { + const state: unknown = Reflect.get(checkpoints, 'state') + if ( + state === null || + typeof state !== 'object' || + !('checkpoints' in state) + ) { + throw new Error( + 'Expected the memory snapshot store to own checkpoint state', + ) + } + const records = state.checkpoints + if (!(records instanceof Map)) { + throw new Error('Expected checkpoint records to use a private Map') + } + return records +} + +async function seedSource( + checkpoints: ForkCapableSandboxCheckpointStore, + conversation: SandboxCheckpoint['conversation'] = [], +): Promise { + const sourceWriter = await checkpoints.acquireWriter('source') + await checkpoints.append({ + checkpoint: { + id: 'source', + threadId: 'source', + parentCheckpointId: null, + createdAt: 1, + reason: 'automatic', + files: [], + conversation, + artifacts: [], + }, + expectedHeadId: null, + writer: sourceWriter, + }) + await sourceWriter.release() +} + +function forkInput( + writer: SandboxCheckpointWriter, + overrides: Partial> = {}, +): SandboxCheckpointForkInput { + return { + sourceThreadId: 'source', + sourceCheckpointId: 'source', + destinationThreadId: 'destination', + destinationCheckpointId: 'fork', + createdAt: 2, + ...overrides, + writer, + } +} + +describe('memory sandbox snapshot persistence', () => { + it('owns checkpoint records instead of wrapping the public in-memory store', () => { + const source = readFileSync( + new URL('../src/memory-snapshots.ts', import.meta.url), + 'utf8', + ) + expect(source).not.toContain('InMemorySandboxCheckpointStore') + }) + + it('preserves a destination transcript saved reentrantly after fork preflight', async () => { + const { persistence, checkpoints } = await memorySandboxSnapshots() + await seedSource(checkpoints, [{ role: 'user', content: 'source' }]) + const writer = await checkpoints.acquireWriter('destination') + const saved = [{ role: 'user', content: 'concurrent' }] as const + const originalStructuredClone = structuredClone + let armed = true + const cloneSpy = vi + .spyOn(globalThis, 'structuredClone') + .mockImplementation((value, options) => { + if (armed) { + armed = false + void persistence.stores.messages.saveThread('destination', [...saved]) + } + return originalStructuredClone(value, options) + }) + try { + await expect( + checkpoints.forkFromCheckpoint(forkInput(writer)), + ).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_FORK_DESTINATION_NOT_EMPTY', + }) + } finally { + cloneSpy.mockRestore() + } + expect(await persistence.stores.messages.loadThread('destination')).toEqual( + saved, + ) + expect(await checkpoints.list('destination')).toEqual([]) + expect(await checkpoints.getHead('destination')).toBeNull() + }) + + it('reads a late-mutating destination checkpoint id only once at entry', async () => { + const { persistence, checkpoints } = await memorySandboxSnapshots() + await seedSource(checkpoints, [{ role: 'user', content: 'source' }]) + const writer = await checkpoints.acquireWriter('destination') + let reads = 0 + const input: SandboxCheckpointForkInput = { + sourceThreadId: 'source', + sourceCheckpointId: 'source', + destinationThreadId: 'destination', + get destinationCheckpointId() { + reads++ + if (reads === 4) { + void persistence.stores.messages.saveThread('destination', [ + { role: 'user', content: 'late concurrent transcript' }, + ]) + } + return 'fork' + }, + createdAt: 2, + writer, + } + + await expect(checkpoints.forkFromCheckpoint(input)).resolves.toMatchObject({ + checkpoint: { id: 'fork' }, + }) + expect(await persistence.stores.messages.loadThread('destination')).toEqual( + [{ role: 'user', content: 'source' }], + ) + expect(await checkpoints.getHead('destination')).toBe('fork') + }) + + it('rejects state written by a destination checkpoint id entry getter', async () => { + const { persistence, checkpoints } = await memorySandboxSnapshots() + await seedSource(checkpoints) + const writer = await checkpoints.acquireWriter('destination') + const concurrent = [{ role: 'user', content: 'entry transcript' }] as const + let reads = 0 + const input: SandboxCheckpointForkInput = { + sourceThreadId: 'source', + sourceCheckpointId: 'source', + destinationThreadId: 'destination', + get destinationCheckpointId() { + reads++ + void persistence.stores.messages.saveThread('destination', [ + ...concurrent, + ]) + return 'fork' + }, + createdAt: 2, + writer, + } + + await expect(checkpoints.forkFromCheckpoint(input)).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_FORK_DESTINATION_NOT_EMPTY', + }) + expect(reads).toBe(1) + expect(await persistence.stores.messages.loadThread('destination')).toEqual( + concurrent, + ) + expect(await checkpoints.get('fork')).toBeNull() + expect(await checkpoints.getHead('destination')).toBeNull() + }) + + it('reads a late-throwing destination thread id only once at entry', async () => { + const { persistence, checkpoints } = await memorySandboxSnapshots() + await seedSource(checkpoints, [{ role: 'user', content: 'source' }]) + const writer = await checkpoints.acquireWriter('destination') + const beforeReferences = await checkpoints.listBlobReferences() + let reads = 0 + const input: SandboxCheckpointForkInput = { + sourceThreadId: 'source', + sourceCheckpointId: 'source', + get destinationThreadId() { + reads++ + if (reads === 10) throw new Error('late destination thread read') + return 'destination' + }, + destinationCheckpointId: 'fork', + createdAt: 2, + writer, + } + + const failure = await checkpoints + .forkFromCheckpoint(input) + .then(() => null) + .catch((error: unknown) => error) + expect(await persistence.stores.messages.loadThread('destination')).toEqual( + [{ role: 'user', content: 'source' }], + ) + expect(await checkpoints.get('fork')).toMatchObject({ + id: 'fork', + threadId: 'destination', + }) + expect(await checkpoints.getHead('destination')).toBe('fork') + expect(await checkpoints.listBlobReferences()).toEqual(beforeReferences) + expect(reads).toBe(1) + expect(failure).toBeNull() + }) + + it('keeps destination state empty when its thread id throws at entry', async () => { + const { persistence, checkpoints } = await memorySandboxSnapshots() + await seedSource(checkpoints, [{ role: 'user', content: 'source' }]) + const writer = await checkpoints.acquireWriter('destination') + const beforeReferences = await checkpoints.listBlobReferences() + const input: SandboxCheckpointForkInput = { + sourceThreadId: 'source', + sourceCheckpointId: 'source', + get destinationThreadId(): string { + throw new Error('entry destination thread read') + }, + destinationCheckpointId: 'fork', + createdAt: 2, + writer, + } + + await expect(checkpoints.forkFromCheckpoint(input)).rejects.toThrow( + 'entry destination thread read', + ) + expect(await persistence.stores.messages.loadThread('destination')).toEqual( + [], + ) + expect(await checkpoints.get('fork')).toBeNull() + expect(await checkpoints.getHead('destination')).toBeNull() + expect(await checkpoints.listBlobReferences()).toEqual(beforeReferences) + }) + + it('reads every fork input and writer primitive exactly once', async () => { + const { checkpoints } = await memorySandboxSnapshots() + await seedSource(checkpoints) + const lease = await checkpoints.acquireWriter('destination') + const reads = { + sourceThreadId: 0, + sourceCheckpointId: 0, + destinationThreadId: 0, + destinationCheckpointId: 0, + createdAt: 0, + writer: 0, + writerThreadId: 0, + writerOwnerToken: 0, + writerFence: 0, + } + const suppliedWriter: SandboxCheckpointWriter = { + get threadId() { + reads.writerThreadId++ + return lease.threadId + }, + get ownerToken() { + reads.writerOwnerToken++ + return lease.ownerToken + }, + get fence() { + reads.writerFence++ + return lease.fence + }, + } + const input: SandboxCheckpointForkInput = { + get sourceThreadId() { + reads.sourceThreadId++ + return 'source' + }, + get sourceCheckpointId() { + reads.sourceCheckpointId++ + return 'source' + }, + get destinationThreadId() { + reads.destinationThreadId++ + return 'destination' + }, + get destinationCheckpointId() { + reads.destinationCheckpointId++ + return 'fork' + }, + get createdAt() { + reads.createdAt++ + return 2 + }, + get writer() { + reads.writer++ + return suppliedWriter + }, + } + + await checkpoints.forkFromCheckpoint(input) + + expect(reads).toEqual({ + sourceThreadId: 1, + sourceCheckpointId: 1, + destinationThreadId: 1, + destinationCheckpointId: 1, + createdAt: 1, + writer: 1, + writerThreadId: 1, + writerOwnerToken: 1, + writerFence: 1, + }) + }) + + it('rejects a true orphan destination checkpoint record', async () => { + const { checkpoints } = await memorySandboxSnapshots() + await seedSource(checkpoints) + rawCheckpointRecords(checkpoints).set('orphan', { + id: 'orphan', + threadId: 'destination', + parentCheckpointId: null, + createdAt: 1, + reason: 'automatic', + files: [], + conversation: [], + artifacts: [], + }) + const writer = await checkpoints.acquireWriter('destination') + await expect( + checkpoints.forkFromCheckpoint(forkInput(writer)), + ).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_FORK_DESTINATION_NOT_EMPTY', + }) + expect(await checkpoints.getHead('destination')).toBeNull() + expect(await checkpoints.get('fork')).toBeNull() + }) + + it('rejects a destination with an existing checkpoint head', async () => { + const { checkpoints } = await memorySandboxSnapshots() + await seedSource(checkpoints) + const writer = await checkpoints.acquireWriter('destination') + await checkpoints.append({ + checkpoint: { + id: 'existing', + threadId: 'destination', + parentCheckpointId: null, + createdAt: 1, + reason: 'automatic', + files: [], + conversation: [], + artifacts: [], + }, + expectedHeadId: null, + writer, + }) + await expect( + checkpoints.forkFromCheckpoint(forkInput(writer)), + ).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_FORK_DESTINATION_NOT_EMPTY', + }) + expect(await checkpoints.getHead('destination')).toBe('existing') + }) + + it('does not invoke public store methods or blob I/O while forking', async () => { + const { persistence, checkpoints } = await memorySandboxSnapshots() + await seedSource(checkpoints) + const writer = await checkpoints.acquireWriter('destination') + const get = vi.spyOn(checkpoints, 'get') + const list = vi.spyOn(checkpoints, 'list') + const getHead = vi.spyOn(checkpoints, 'getHead') + const append = vi.spyOn(checkpoints, 'append') + const deleteHead = vi.spyOn(checkpoints, 'deleteHead') + const acquireWriter = vi.spyOn(checkpoints, 'acquireWriter') + const listBlobReferences = vi.spyOn(checkpoints, 'listBlobReferences') + const saveThread = vi.spyOn(persistence.stores.messages, 'saveThread') + const loadThread = vi.spyOn(persistence.stores.messages, 'loadThread') + const blobPut = vi.spyOn(persistence.stores.blobs, 'put') + const blobGet = vi.spyOn(persistence.stores.blobs, 'get') + const blobHead = vi.spyOn(persistence.stores.blobs, 'head') + const blobDelete = vi.spyOn(persistence.stores.blobs, 'delete') + const blobList = vi.spyOn(persistence.stores.blobs, 'list') + + await checkpoints.forkFromCheckpoint(forkInput(writer)) + + expect( + [ + get, + list, + getHead, + append, + deleteHead, + acquireWriter, + listBlobReferences, + saveThread, + loadThread, + blobPut, + blobGet, + blobHead, + blobDelete, + blobList, + ].map((spy) => spy.mock.calls.length), + ).toEqual([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) + }) + + it.each([1, 2, 3])( + 'keeps the destination empty when staged clone %i throws and permits retry', + async (throwAt) => { + const { persistence, checkpoints } = await memorySandboxSnapshots() + await seedSource(checkpoints, [{ role: 'user', content: 'source' }]) + const writer = await checkpoints.acquireWriter('destination') + const beforeReferences = await checkpoints.listBlobReferences() + const originalStructuredClone = structuredClone + let cloneCall = 0 + const cloneSpy = vi + .spyOn(globalThis, 'structuredClone') + .mockImplementation((value, options) => { + cloneCall++ + if (cloneCall === throwAt) throw new Error(`clone stage ${throwAt}`) + return originalStructuredClone(value, options) + }) + try { + await expect( + checkpoints.forkFromCheckpoint(forkInput(writer)), + ).rejects.toThrow(`clone stage ${throwAt}`) + } finally { + cloneSpy.mockRestore() + } + expect( + await persistence.stores.messages.loadThread('destination'), + ).toEqual([]) + expect(await checkpoints.list('destination')).toEqual([]) + expect(await checkpoints.getHead('destination')).toBeNull() + expect(await checkpoints.listBlobReferences()).toEqual(beforeReferences) + await expect( + checkpoints.forkFromCheckpoint(forkInput(writer)), + ).resolves.toMatchObject({ checkpoint: { id: 'fork' } }) + }, + ) + + it('forks a selected historical checkpoint and counts each distinct blob once', async () => { + const { persistence, checkpoints } = await memorySandboxSnapshots() + const sharedFileKey = `sandbox-files/sha256/${'a'.repeat(64)}` + const sharedArtifactKey = `sandbox-artifacts/sha256/${'b'.repeat(64)}` + const sourceWriter = await checkpoints.acquireWriter('source') + await checkpoints.append({ + checkpoint: { + id: 'historical', + threadId: 'source', + parentCheckpointId: null, + createdAt: 1, + reason: 'named', + files: [ + { path: 'a.txt', kind: 'file', blobKey: sharedFileKey, size: 1 }, + { path: 'b.txt', kind: 'file', blobKey: sharedFileKey, size: 1 }, + ], + conversation: [{ role: 'user', content: 'historical' }], + artifacts: [ + { + artifactId: 'artifact-a', + name: 'a.txt', + mimeType: 'text/plain', + size: 1, + blobKey: sharedArtifactKey, + createdAt: 1, + }, + { + artifactId: 'artifact-b', + name: 'b.txt', + mimeType: 'text/plain', + size: 1, + blobKey: sharedArtifactKey, + createdAt: 1, + }, + ], + }, + expectedHeadId: null, + writer: sourceWriter, + }) + await checkpoints.append({ + checkpoint: { + id: 'current', + threadId: 'source', + parentCheckpointId: 'historical', + createdAt: 2, + reason: 'automatic', + files: [], + conversation: [{ role: 'user', content: 'current' }], + artifacts: [], + }, + expectedHeadId: 'historical', + writer: sourceWriter, + }) + const writer = await checkpoints.acquireWriter('destination') + const result = await checkpoints.forkFromCheckpoint( + forkInput(writer, { sourceCheckpointId: 'historical' }), + ) + + expect(result.checkpoint.parentCheckpointId).toBeNull() + expect(result.checkpoint.conversation).toEqual([ + { role: 'user', content: 'historical' }, + ]) + expect(await persistence.stores.messages.loadThread('destination')).toEqual( + result.checkpoint.conversation, + ) + expect(await checkpoints.getHead('source')).toBe('current') + expect(await checkpoints.listBlobReferences()).toEqual([ + { key: sharedArtifactKey, references: 2 }, + { key: sharedFileKey, references: 2 }, + ]) + }) + + it('returns precise errors for every invalid fork identity and writer case', async () => { + const { checkpoints } = await memorySandboxSnapshots() + await seedSource(checkpoints) + const destinationWriter = await checkpoints.acquireWriter('destination') + const otherWriter = await checkpoints.acquireWriter('other') + + await expect( + checkpoints.forkFromCheckpoint( + forkInput(destinationWriter, { sourceCheckpointId: 'missing' }), + ), + ).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_FORK_SOURCE_NOT_FOUND', + }) + await expect( + checkpoints.forkFromCheckpoint( + forkInput(destinationWriter, { sourceThreadId: 'wrong-source' }), + ), + ).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_FORK_SOURCE_THREAD_MISMATCH', + }) + await expect( + checkpoints.forkFromCheckpoint( + forkInput(destinationWriter, { destinationThreadId: 'source' }), + ), + ).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_FORK_SOURCE_THREAD_MISMATCH', + }) + await expect( + checkpoints.forkFromCheckpoint(forkInput(otherWriter)), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_WRITER_LOST' }) + await expect( + checkpoints.forkFromCheckpoint( + forkInput(destinationWriter, { destinationCheckpointId: '' }), + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ID' }) + await expect( + checkpoints.forkFromCheckpoint( + forkInput(destinationWriter, { createdAt: Number.NaN }), + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ENTRY' }) + + await destinationWriter.release() + await expect( + checkpoints.forkFromCheckpoint(forkInput(destinationWriter)), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_WRITER_LOST' }) + expect(await checkpoints.getHead('destination')).toBeNull() + }) + + it('blocks every destination-owned state shape without publishing', async () => { + const stateSetups = [ + async (p: MemoryPersistence) => + p.stores.messages.saveThread('destination', []), + async (p: MemoryPersistence) => + p.stores.runs.createOrResume({ + runId: 'run', + threadId: 'destination', + startedAt: 1, + }), + async (p: MemoryPersistence) => + p.stores.generationRuns.createOrResume({ + runId: 'gen', + threadId: 'destination', + activity: 'generateText', + provider: 'test', + model: 'test', + startedAt: 1, + }), + async (p: MemoryPersistence) => + p.stores.interrupts.create({ + interruptId: 'interrupt', + runId: 'run', + threadId: 'destination', + requestedAt: 1, + payload: {}, + }), + async (p: MemoryPersistence) => + p.stores.artifacts.save({ + artifactId: 'artifact', + runId: 'run', + threadId: 'destination', + name: 'a', + mimeType: 'text/plain', + size: 0, + blobKey: `sandbox-artifacts/sha256/${'1'.repeat(64)}`, + createdAt: 1, + }), + ] + for (const setup of stateSetups) { + const { persistence, checkpoints } = await memorySandboxSnapshots() + await seedSource(checkpoints) + await setup(persistence) + const writer = await checkpoints.acquireWriter('destination') + await expect( + checkpoints.forkFromCheckpoint(forkInput(writer)), + ).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_FORK_DESTINATION_NOT_EMPTY', + }) + expect(await checkpoints.getHead('destination')).toBeNull() + } + }) + + it('keeps source and fork snapshots deeply independent', async () => { + const { persistence, checkpoints } = await memorySandboxSnapshots() + await seedSource(checkpoints, [{ role: 'user', content: 'checkpoint' }]) + await persistence.stores.messages.saveThread('source', [ + { role: 'user', content: 'later' }, + ]) + const writer = await checkpoints.acquireWriter('destination') + const result = await checkpoints.forkFromCheckpoint(forkInput(writer)) + expect(result.checkpoint.parentCheckpointId).toBeNull() + expect(result.checkpoint.conversation).toEqual([ + { role: 'user', content: 'checkpoint' }, + ]) + expect(await persistence.stores.messages.loadThread('destination')).toEqual( + [{ role: 'user', content: 'checkpoint' }], + ) + expect(await persistence.stores.messages.loadThread('source')).toEqual([ + { role: 'user', content: 'later' }, + ]) + const returnedMessage = result.checkpoint.conversation[0] + if (returnedMessage === undefined) { + throw new Error('Expected the fork result to contain its conversation') + } + Reflect.set(returnedMessage, 'content', 'mutated result') + expect((await checkpoints.get('fork'))?.conversation).toEqual([ + { role: 'user', content: 'checkpoint' }, + ]) + expect(await persistence.stores.messages.loadThread('destination')).toEqual( + [{ role: 'user', content: 'checkpoint' }], + ) + }) + + it('allows one concurrent fork for a shared valid lease', async () => { + const { checkpoints } = await memorySandboxSnapshots() + await seedSource(checkpoints) + const writer = await checkpoints.acquireWriter('destination') + const input = forkInput(writer) + const results = await Promise.allSettled([ + checkpoints.forkFromCheckpoint(input), + checkpoints.forkFromCheckpoint(input), + ]) + expect( + results.filter((value) => value.status === 'fulfilled'), + ).toHaveLength(1) + expect(await checkpoints.list('destination')).toHaveLength(1) + expect(await checkpoints.getHead('destination')).toBe('fork') + }) + + it('supports blob streams, ranges, pagination, and artifact ordering', async () => { + const { persistence } = await memorySandboxSnapshots() + const first = await persistence.stores.blobs.put('b', new Blob(['abcdef'])) + await persistence.stores.blobs.put('a', 'abc') + const object = await persistence.stores.blobs.get('b', { + range: { offset: 1, length: 3 }, + }) + expect(first.contentType).toBeUndefined() + expect(object === null ? undefined : await object.text()).toBe('bcd') + const page = await persistence.stores.blobs.list({ limit: 1 }) + expect(page.objects.map((value) => value.key)).toEqual(['a']) + expect(page.truncated).toBe(true) + const next = await persistence.stores.blobs.list( + page.cursor === undefined ? {} : { cursor: page.cursor }, + ) + expect(next.objects.map((value) => value.key)).toEqual(['b']) + }) + + it('rejects an empty saved destination transcript without mutation', async () => { + const { persistence, checkpoints } = await memorySandboxSnapshots() + await seedSource(checkpoints) + await persistence.stores.messages.saveThread('destination', []) + const destinationWriter = await checkpoints.acquireWriter('destination') + await expect( + checkpoints.forkFromCheckpoint(forkInput(destinationWriter)), + ).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_FORK_DESTINATION_NOT_EMPTY', + }) + expect(await persistence.stores.messages.loadThread('destination')).toEqual( + [], + ) + expect(await checkpoints.getHead('destination')).toBeNull() + }) + + it('leaves a new destination empty when fork input reuses a checkpoint id', async () => { + const { persistence, checkpoints } = await memorySandboxSnapshots() + await seedSource(checkpoints, [{ role: 'user', content: 'source' }]) + const destinationWriter = await checkpoints.acquireWriter('destination') + + await expect( + checkpoints.forkFromCheckpoint( + forkInput(destinationWriter, { destinationCheckpointId: 'source' }), + ), + ).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_FORK_DESTINATION_NOT_EMPTY', + }) + + expect(await persistence.stores.messages.loadThread('destination')).toEqual( + [], + ) + expect(await checkpoints.list('destination')).toEqual([]) + expect(await checkpoints.getHead('destination')).toBeNull() + }) + + it('returns interrupt records directly without cloning payloads or responses', async () => { + const { persistence } = await memorySandboxSnapshots() + const payload = { callback: () => 'payload' } + const response = () => 'response' + await persistence.stores.interrupts.create({ + interruptId: 'pending-interrupt', + runId: 'run', + threadId: 'thread', + requestedAt: 1, + payload, + }) + await persistence.stores.interrupts.create({ + interruptId: 'resolved-interrupt', + runId: 'run', + threadId: 'thread', + requestedAt: 2, + payload, + }) + await persistence.stores.interrupts.resolve('resolved-interrupt', response) + + const pending = await persistence.stores.interrupts.get('pending-interrupt') + const resolved = + await persistence.stores.interrupts.get('resolved-interrupt') + const all = await persistence.stores.interrupts.listByRun('run') + const pendingByRun = + await persistence.stores.interrupts.listPendingByRun('run') + + expect(all).toEqual([pending, resolved]) + expect(all[0]).toBe(pending) + expect(all[1]).toBe(resolved) + expect(pendingByRun).toEqual([pending]) + expect(pendingByRun[0]).toBe(pending) + expect(pending?.payload).toBe(payload) + expect(resolved?.response).toBe(response) + }) +}) diff --git a/packages/ai-sandbox/tests/root-declaration-consumer.test.ts b/packages/ai-sandbox/tests/root-declaration-consumer.test.ts new file mode 100644 index 0000000000..bb525045a9 --- /dev/null +++ b/packages/ai-sandbox/tests/root-declaration-consumer.test.ts @@ -0,0 +1,89 @@ +import { readdirSync, readFileSync } from 'node:fs' +import { join } from 'node:path' +import { describe, expect, it } from 'vitest' + +const declarationsRoot = join(import.meta.dirname, '..', 'dist', 'esm') +const sourceRoot = join(import.meta.dirname, '..', 'src') + +const prohibitedProductionTestSeams = [ + 'memorySandboxSnapshotsForTest', + 'MemorySandboxSnapshotsTestOptions', + 'failAfterTranscriptStage', + 'failAfterCheckpointStage', + 'failAfterReferenceStage', + 'failAfterHeadStage', + 'createInMemoryCheckpointCoordinator', + 'InMemoryCheckpointCoordinator', + 'PreparedCheckpointFork', +] + +function filesUnder(directory: string, extension: string): Array { + return readdirSync(directory, { withFileTypes: true }).flatMap((entry) => { + const path = join(directory, entry.name) + if (entry.isDirectory()) return filesUnder(path, extension) + return entry.isFile() && path.endsWith(extension) ? [path] : [] + }) +} + +function rootDeclarationGraph(entry: string): Array { + const seen = new Set() + const visit = (file: string): void => { + if (seen.has(file)) return + seen.add(file) + const source = readFileSync(file, 'utf8') + for (const specifier of source.matchAll(/from ['"](\.\/[^'"]+)['"]/g)) { + const moduleSpecifier = specifier[1] + if (moduleSpecifier === undefined) continue + const imported = join(file, '..', moduleSpecifier) + const declaration = imported.endsWith('.d.ts') + ? imported + : imported.endsWith('.js') + ? imported.slice(0, -3) + '.d.ts' + : `${imported}.d.ts` + visit(declaration) + } + } + visit(entry) + return [...seen] +} + +describe('root declarations', () => { + function expectNoPersistenceDeclarationReference(entry: string): void { + const files = rootDeclarationGraph(entry) + expect(files.length).toBeGreaterThan(0) + + const persistenceSpecifiers = files.flatMap((file) => { + const source = readFileSync(file, 'utf8') + return source.includes('@tanstack/ai-persistence') ? [file] : [] + }) + + expect(persistenceSpecifiers).toEqual([]) + } + + it('do not require persistence type resolution', () => { + expectNoPersistenceDeclarationReference( + join(declarationsRoot, 'index.d.ts'), + ) + }) + + it('ships a testkit declaration graph that does not require persistence', () => { + expectNoPersistenceDeclarationReference( + join(declarationsRoot, 'testkit', 'conformance.d.ts'), + ) + }) + + it('do not contain production test seams in source or emitted declarations', () => { + const files = [ + ...filesUnder(sourceRoot, '.ts'), + ...filesUnder(declarationsRoot, '.d.ts'), + ] + const matches = files.flatMap((file) => { + const source = readFileSync(file, 'utf8') + return prohibitedProductionTestSeams.flatMap((seam) => + source.includes(seam) ? [`${file}: ${seam}`] : [], + ) + }) + + expect(matches).toEqual([]) + }) +}) diff --git a/packages/ai-sandbox/tests/snapshot-lifecycle.test.ts b/packages/ai-sandbox/tests/snapshot-lifecycle.test.ts new file mode 100644 index 0000000000..5c3bf39183 --- /dev/null +++ b/packages/ai-sandbox/tests/snapshot-lifecycle.test.ts @@ -0,0 +1,2146 @@ +import { describe, expect, it, vi } from 'vitest' +import { EventType, chat } from '@tanstack/ai' +import { provideRunDisconnect } from '@tanstack/ai/adapter-internals' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' +import { + PersistenceCompletionCapability, + memoryPersistence, + withPersistence, +} from '@tanstack/ai-persistence' +import { defineSandbox } from '../src/sandbox' +import { withSandbox } from '../src/middleware' +import { memorySandboxSnapshots } from '../src/memory-snapshots' +import { SandboxCapability } from '../src/capabilities' +import { InMemorySandboxInstanceStore } from '../src/instance-store' +import { InMemorySandboxCheckpointStore } from '../src/checkpoint-store' +import type { + SandboxCheckpointStore, + SandboxCheckpointWriterLease, +} from '../src/checkpoint-store' +import type { + SandboxCapabilities, + SandboxHandle, + SandboxProvider, +} from '../src/contracts' +import { fakeLog, makeMiddlewareCtx } from './fakes' + +const caps: SandboxCapabilities = { + fs: true, + exec: true, + env: true, + ports: true, + backgroundProcesses: true, + writableStdin: true, + killableProcesses: true, + snapshots: true, + networkPolicy: true, + durableFilesystem: true, + fork: true, +} +const adapter: AnyTextAdapter = { + kind: 'text', + name: 'snapshot-test', + model: 'snapshot-test-model', + '~types': { + providerOptions: undefined, + inputModalities: undefined, + toolCapabilities: undefined, + toolCallMetadata: undefined, + systemPromptMetadata: undefined, + messageMetadataByModality: undefined, + }, + chatStream: async function* (): AsyncGenerator { + yield { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: 1, + } + yield { + type: EventType.RUN_FINISHED, + runId: 'run-1', + threadId: 'thread-1', + finishReason: 'stop', + timestamp: 1, + } + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), +} + +type Event = { method: string; id?: string; order: number } +type WorkspaceSeed = + | { path: string; type: 'dir' } + | { path: string; type: 'file'; data: Uint8Array } +type FixtureOptions = { + persistence?: ReturnType + checkpoints?: SandboxCheckpointStore + workspace?: Array + onWorkspaceList?: (path: string) => void | Promise + watch?: WatchProbe + nativeSnapshotError?: Error +} +type WatchProbe = { + emit?: (event: { type: string; path: string }) => void + stops: number +} +type Fixture = { + provider: SandboxProvider + events: Array + instances: InMemorySandboxInstanceStore + checkpoints: SandboxCheckpointStore + persistence: ReturnType + definition: ReturnType + resumed?: SandboxHandle +} + +class RenewalFailingCheckpointStore extends InMemorySandboxCheckpointStore { + override async acquireWriter( + threadId: string, + ): Promise { + const lease = await super.acquireWriter(threadId) + return { + ...lease, + renew: async () => { + throw new Error('writer lease was lost') + }, + } + } +} + +class ReleaseFailingCheckpointStore extends InMemorySandboxCheckpointStore { + releases = 0 + + override async acquireWriter( + threadId: string, + ): Promise { + const lease = await super.acquireWriter(threadId) + return { + ...lease, + release: async () => { + this.releases++ + await lease.release() + throw new Error('writer release failed') + }, + } + } +} + +class ReleaseCountingCheckpointStore extends InMemorySandboxCheckpointStore { + releases = 0 + + override async acquireWriter( + threadId: string, + ): Promise { + const lease = await super.acquireWriter(threadId) + return { + ...lease, + release: async () => { + this.releases++ + await lease.release() + }, + } + } +} + +function afterRunDefinition(f: Fixture): ReturnType { + return defineSandbox({ + id: 'fixture', + provider: f.provider, + lifecycle: { snapshot: 'after-run', destroyOnComplete: true }, + workspace: { source: { type: 'none' } }, + fileEvents: false, + }) +} + +async function runTerminalSnapshot( + f: Fixture, + checkpoints: InMemorySandboxCheckpointStore, +): Promise { + await drain( + chat({ + adapter, + messages: [{ role: 'user', content: 'hello' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [ + withPersistence(f.persistence), + withSandbox(afterRunDefinition(f), { + instances: f.instances, + snapshots: { persistence: f.persistence, checkpoints }, + }), + ], + }), + ) +} + +class RenewalCountingCheckpointStore extends InMemorySandboxCheckpointStore { + renewals = 0 + override async acquireWriter( + threadId: string, + ): Promise { + const lease = await super.acquireWriter(threadId) + return { + ...lease, + renew: async () => { + this.renewals++ + return lease.renew() + }, + } + } +} + +function deferred(): { promise: Promise; resolve: () => void } { + let resolve: () => void = () => {} + const promise = new Promise((complete) => { + resolve = complete + }) + return { promise, resolve } +} + +class DeferredRenewalCheckpointStore extends InMemorySandboxCheckpointStore { + renewals = 0 + releases = 0 + readonly renewalGate = deferred() + readonly releaseGate = deferred() + readonly releaseStarted = deferred() + + override async acquireWriter( + threadId: string, + ): Promise { + const lease = await super.acquireWriter(threadId) + return { + ...lease, + renew: async () => { + this.renewals++ + await this.renewalGate.promise + return lease.renew() + }, + release: async () => { + this.releases++ + this.releaseStarted.resolve() + await this.releaseGate.promise + await lease.release() + }, + } + } +} + +class ControlledRenewalLossStore extends InMemorySandboxCheckpointStore { + readonly renewalStarted = deferred() + readonly renewalGate = deferred() + releases = 0 + failRelease = false + + override async acquireWriter( + threadId: string, + ): Promise { + const lease = await super.acquireWriter(threadId) + return { + ...lease, + renew: async () => { + this.renewalStarted.resolve() + await this.renewalGate.promise + throw new Error('writer renewal was lost') + }, + release: async () => { + this.releases++ + await lease.release() + if (this.failRelease) throw new Error('writer release also failed') + }, + } + } +} + +function handle( + id: string, + events?: Array, + options: Pick< + FixtureOptions, + 'workspace' | 'onWorkspaceList' | 'watch' | 'nativeSnapshotError' + > = {}, +): SandboxHandle { + const watchProbe = options.watch + const entries = new Map([ + ['/workspace', { type: 'dir' }], + ]) + const normalize = (path: string) => path.replace(/\/+$/, '') || '/' + for (const entry of options.workspace ?? []) { + entries.set(normalize(entry.path), { + type: entry.type, + ...(entry.type === 'file' ? { data: entry.data.slice() } : {}), + }) + } + const children = (path: string) => { + const prefix = `${normalize(path)}/` + return [...entries].flatMap(([entryPath, entry]) => { + const rest = entryPath.startsWith(prefix) + ? entryPath.slice(prefix.length) + : '' + return rest && !rest.includes('/') + ? [{ name: rest, path: entryPath, type: entry.type }] + : [] + }) + } + return { + id, + provider: 'fixture', + capabilities: caps, + fs: { + read: async (p) => + new TextDecoder().decode( + entries.get(normalize(p))?.data ?? new Uint8Array(), + ), + readBytes: async (p) => + entries.get(normalize(p))?.data?.slice() ?? new Uint8Array(), + write: async (p, d) => { + events?.push({ method: 'fs.write', id: p, order: events.length + 1 }) + entries.set(normalize(p), { + type: 'file', + data: typeof d === 'string' ? new TextEncoder().encode(d) : d.slice(), + }) + }, + list: async (p) => { + events?.push({ method: 'fs.list', id: p, order: events.length + 1 }) + await options.onWorkspaceList?.(p) + return children(p) + }, + lstat: async (p) => { + const entry = entries.get(normalize(p)) + if (!entry) return undefined + return entry.type === 'dir' + ? { type: 'dir', mode: 0o755 } + : { type: 'file', mode: 0o644, size: entry.data?.byteLength ?? 0 } + }, + mkdir: async (p) => { + events?.push({ method: 'fs.mkdir', id: p, order: events.length + 1 }) + entries.set(normalize(p), { type: 'dir' }) + }, + remove: async (p) => { + const path = normalize(p) + for (const entryPath of entries.keys()) { + if (entryPath === path || entryPath.startsWith(`${path}/`)) + entries.delete(entryPath) + } + }, + rename: async (from, to) => { + const source = normalize(from) + const target = normalize(to) + for (const [entryPath, entry] of [...entries]) { + if (entryPath === source || entryPath.startsWith(`${source}/`)) { + entries.delete(entryPath) + entries.set(`${target}${entryPath.slice(source.length)}`, entry) + } + } + }, + exists: async (p) => entries.has(normalize(p)), + ...(watchProbe !== undefined + ? { + watch: async ( + _path: string, + onEvent: (event: { type: string; path: string }) => void, + ) => { + let active = true + watchProbe.emit = (event) => { + if (active) onEvent(event) + } + return { + stop: async () => { + if (!active) return + active = false + watchProbe.stops++ + }, + } + }, + } + : {}), + }, + git: { + clone: async () => {}, + status: async () => '', + add: async () => {}, + commit: async () => {}, + push: async () => {}, + pull: async () => {}, + branch: async () => 'main', + }, + process: { + exec: async () => ({ stdout: '', stderr: '', exitCode: 0 }), + spawn: async () => { + throw new Error('unused') + }, + }, + ports: { connect: async (port) => ({ url: `http://localhost:${port}` }) }, + env: { set: async () => {} }, + snapshot: async (label) => { + events?.push({ + method: 'handle.snapshot', + id: label, + order: events.length + 1, + }) + if (options.nativeSnapshotError) throw options.nativeSnapshotError + return { id: `snapshot-${id}`, label } + }, + destroy: async () => { + events?.push({ method: 'handle.destroy', id, order: events.length + 1 }) + }, + } +} + +function fixture(options: FixtureOptions = {}): Fixture { + const events: Array = [] + let order = 0 + let resumed: SandboxHandle | undefined + const provider: SandboxProvider = { + name: 'fixture', + capabilities: () => caps, + create: async (input) => { + const id = input.id ?? 'created' + events.push({ method: 'create', id, order: ++order }) + return handle(id, events, options) + }, + resume: async (input) => { + events.push({ method: 'resume', id: input.id, order: ++order }) + resumed = handle(input.id, events, options) + resumed.fs.read = async () => { + throw new Error('filesystem touched') + } + resumed.fs.readBytes = async () => { + throw new Error('filesystem touched') + } + resumed.fs.write = async () => { + throw new Error('filesystem touched') + } + resumed.fs.list = async () => { + throw new Error('filesystem touched') + } + resumed.fs.lstat = async () => { + throw new Error('filesystem touched') + } + resumed.fs.mkdir = async () => { + throw new Error('filesystem touched') + } + resumed.fs.remove = async () => { + throw new Error('filesystem touched') + } + resumed.fs.rename = async () => { + throw new Error('filesystem touched') + } + resumed.fs.exists = async () => { + throw new Error('filesystem touched') + } + return resumed + }, + restoreSnapshot: async (input) => { + events.push({ + method: 'restoreSnapshot', + id: input.snapshotId, + order: ++order, + }) + return handle('restored', events, options) + }, + destroy: async (input) => { + events.push({ method: 'destroy', id: input.id, order: ++order }) + }, + } + const persistence = options.persistence ?? memoryPersistence() + const checkpoints = + options.checkpoints ?? new InMemorySandboxCheckpointStore() + const instances = new InMemorySandboxInstanceStore() + const definition = defineSandbox({ + id: 'fixture', + provider, + lifecycle: { reuse: 'thread', snapshot: 'after-setup' }, + workspace: { source: { type: 'none' } }, + fileEvents: false, + }) + return { + provider, + events, + instances, + checkpoints, + persistence, + definition, + resumed, + } +} + +async function drain(stream: AsyncIterable): Promise { + for await (const _chunk of stream) { + } +} + +async function drainIterator( + iterator: AsyncIterator, +): Promise { + while (!(await iterator.next()).done) {} +} + +async function seedCheckpoint( + f: Fixture, + input: { + id?: string + files?: Array<{ path: string; blobKey: string; size: number }> + } = {}, +): Promise { + const id = input.id ?? 'seed-checkpoint' + const writer = await f.checkpoints.acquireWriter('thread-1') + try { + await f.checkpoints.append({ + checkpoint: { + id, + threadId: 'thread-1', + parentCheckpointId: null, + createdAt: 1, + reason: 'automatic', + files: (input.files ?? []).map((file) => ({ + kind: 'file' as const, + ...file, + })), + conversation: [], + artifacts: [], + }, + expectedHeadId: null, + writer, + }) + } finally { + await writer.release() + } + return id +} + +async function sandboxFileBlobKey(bytes: Uint8Array): Promise { + const digest = await globalThis.crypto.subtle.digest( + 'SHA-256', + new Uint8Array(bytes), + ) + return `sandbox-files/sha256/${Array.from(new Uint8Array(digest), (value) => + value.toString(16).padStart(2, '0'), + ).join('')}` +} + +function snapshotMiddleware( + f: Fixture, + options: { + definition?: ReturnType + checkpoints?: SandboxCheckpointStore + durability?: ReturnType + } = {}, +) { + return withSandbox(options.definition ?? f.definition, { + instances: f.instances, + ...(options.durability !== undefined + ? { + runs: f.persistence.stores.runs, + durability: { + adapter: options.durability, + detachOnDisconnect: true, + }, + } + : {}), + snapshots: { + persistence: f.persistence, + checkpoints: options.checkpoints ?? f.checkpoints, + }, + }) +} + +async function flushMicrotasks(): Promise { + await Promise.resolve() + await Promise.resolve() +} + +async function* failingStream(message: string): AsyncGenerator { + yield* [] + throw new Error(message) +} + +type CaptureBoundary = + | 'completion' + | 'conversation' + | 'files' + | 'artifacts' + | 'head' + +async function startRenewalLossAtBoundary( + boundary: CaptureBoundary, + options: { failRelease?: boolean } = {}, +) { + const boundaryStarted = deferred() + const boundaryGate = deferred() + let boundaryArmed = false + const checkpoints = new ControlledRenewalLossStore({ + leaseDurationMs: 120_000, + renewAfterMs: 10, + }) + checkpoints.failRelease = options.failRelease === true + const f = fixture({ + checkpoints, + onWorkspaceList: async () => { + if (!boundaryArmed || boundary !== 'files') return + boundaryStarted.resolve() + await boundaryGate.promise + }, + }) + const ctx = makeMiddlewareCtx({ threadId: 'thread-1', runId: 'run-1' }) + const persistence = withPersistence(f.persistence) + const sandbox = snapshotMiddleware(f, { + definition: afterRunDefinition(f), + checkpoints, + }) + await persistence.setup?.(ctx) + await sandbox.setup?.(ctx) + const finish = { duration: 1, finishReason: 'stop', content: '' } + await persistence.onFinish?.(ctx, finish) + + const block = async (): Promise => { + boundaryStarted.resolve() + await boundaryGate.promise + } + if (boundary === 'completion') { + vi.spyOn( + ctx.get(PersistenceCompletionCapability), + 'waitForRunCompletion', + ).mockImplementation(block) + } else if (boundary === 'conversation') { + const loadThread = f.persistence.stores.messages.loadThread.bind( + f.persistence.stores.messages, + ) + vi.spyOn(f.persistence.stores.messages, 'loadThread').mockImplementation( + async (threadId) => { + await block() + return loadThread(threadId) + }, + ) + } else if (boundary === 'artifacts') { + const listForThread = f.persistence.stores.artifacts.listForThread.bind( + f.persistence.stores.artifacts, + ) + vi.spyOn( + f.persistence.stores.artifacts, + 'listForThread', + ).mockImplementation(async (threadId) => { + await block() + return listForThread(threadId) + }) + } else if (boundary === 'head') { + const getHead = checkpoints.getHead.bind(checkpoints) + vi.spyOn(checkpoints, 'getHead').mockImplementation(async (threadId) => { + await block() + return getHead(threadId) + }) + } + boundaryArmed = true + const append = vi.spyOn(checkpoints, 'append') + const terminal = Promise.resolve(sandbox.onFinish?.(ctx, finish)) + await boundaryStarted.promise + + vi.advanceTimersByTime(10) + await checkpoints.renewalStarted.promise + checkpoints.renewalGate.resolve() + await flushMicrotasks() + boundaryGate.resolve() + + return { append, checkpoints, f, terminal } +} + +async function startRenewalLossDuringAppend(outcome: 'resolve' | 'reject') { + const appendStarted = deferred() + const appendGate = deferred() + const checkpoints = new ControlledRenewalLossStore({ + leaseDurationMs: 120_000, + renewAfterMs: 10, + }) + checkpoints.failRelease = true + const f = fixture({ checkpoints }) + const ctx = makeMiddlewareCtx({ threadId: 'thread-1', runId: 'run-1' }) + const persistence = withPersistence(f.persistence) + const sandbox = snapshotMiddleware(f, { + definition: afterRunDefinition(f), + checkpoints, + }) + await persistence.setup?.(ctx) + await sandbox.setup?.(ctx) + const finish = { duration: 1, finishReason: 'stop', content: '' } + await persistence.onFinish?.(ctx, finish) + + const append = checkpoints.append.bind(checkpoints) + vi.spyOn(checkpoints, 'append').mockImplementation(async (input) => { + appendStarted.resolve() + await appendGate.promise + if (outcome === 'reject') throw new Error('checkpoint append failed') + return append(input) + }) + const terminal = Promise.resolve(sandbox.onFinish?.(ctx, finish)) + await appendStarted.promise + + vi.advanceTimersByTime(10) + await checkpoints.renewalStarted.promise + checkpoints.renewalGate.resolve() + await flushMicrotasks() + appendGate.resolve() + + return { checkpoints, f, terminal } +} + +describe('sandbox snapshot lifecycle foundation', () => { + it('requires the same persistence object to be installed before snapshot sandbox setup', async () => { + const f = fixture() + const ctx = makeMiddlewareCtx({ threadId: 'thread-1', runId: 'run-1' }) + await expect( + withSandbox(f.definition, { + instances: f.instances, + snapshots: { persistence: f.persistence, checkpoints: f.checkpoints }, + }).setup?.(ctx), + ).rejects.toThrow( + 'Sandbox snapshots require withPersistence(snapshots.persistence) before withSandbox', + ) + expect(f.events).toEqual([]) + }) + + it('rejects a different installed persistence object before snapshot setup', async () => { + const f = fixture() + const otherPersistence = memoryPersistence() + const ctx = makeMiddlewareCtx({ threadId: 'thread-1', runId: 'run-1' }) + await withPersistence(otherPersistence).setup?.(ctx) + + await expect( + withSandbox(f.definition, { + instances: f.instances, + snapshots: { persistence: f.persistence, checkpoints: f.checkpoints }, + }).setup?.(ctx), + ).rejects.toThrow( + 'Sandbox snapshots require the same persistence instance passed to withPersistence', + ) + expect(f.events).toEqual([]) + }) + + it('does not acquire a writer or touch snapshot stores when snapshots are disabled', async () => { + const f = fixture() + const acquire = vi.spyOn(f.checkpoints, 'acquireWriter') + const getHead = vi.spyOn(f.checkpoints, 'getHead') + const ctx = makeMiddlewareCtx({ threadId: 'thread-1', runId: 'run-1' }) + await withSandbox(f.definition, { instances: f.instances }).setup?.(ctx) + expect(acquire).not.toHaveBeenCalled() + expect(getHead).not.toHaveBeenCalled() + }) + + it('fails terminal completion and releases the writer when renewal loses its fence', async () => { + vi.useFakeTimers() + try { + const f = fixture() + const checkpoints = new RenewalFailingCheckpointStore({ + leaseDurationMs: 120_000, + renewAfterMs: 10, + }) + let releaseAdapter: () => void = () => {} + const waitForFinish = new Promise((resolve) => { + releaseAdapter = resolve + }) + const stalledAdapter: AnyTextAdapter = { + ...adapter, + chatStream: async function* () { + yield { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: 1, + } + await waitForFinish + yield { + type: EventType.RUN_FINISHED, + runId: 'run-1', + threadId: 'thread-1', + finishReason: 'stop', + timestamp: 1, + } + }, + } + const iterator = chat({ + adapter: stalledAdapter, + messages: [{ role: 'user', content: 'hello' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [ + withPersistence(f.persistence), + withSandbox(f.definition, { + instances: f.instances, + snapshots: { persistence: f.persistence, checkpoints }, + }), + ], + })[Symbol.asyncIterator]() + await iterator.next() + await vi.advanceTimersByTimeAsync(10) + releaseAdapter() + await expect(drainIterator(iterator)).rejects.toThrow( + 'writer lease was lost', + ) + const lease = await checkpoints.acquireWriter('thread-1') + await lease.release() + } finally { + vi.useRealTimers() + } + }) + + it.each([ + 'completion', + 'conversation', + 'files', + 'artifacts', + 'head', + ])('preserves renewal loss at the %s boundary', async (boundary) => { + vi.useFakeTimers() + try { + const result = await startRenewalLossAtBoundary(boundary) + await expect(result.terminal).rejects.toThrow('writer renewal was lost') + expect(result.append).not.toHaveBeenCalled() + expect( + result.f.events.some((event) => event.method === 'handle.snapshot'), + ).toBe(false) + expect( + result.f.events.filter((event) => event.method === 'destroy'), + ).toHaveLength(1) + expect(result.checkpoints.releases).toBe(1) + } finally { + vi.useRealTimers() + } + }) + + it('does not mask renewal loss when release also fails', async () => { + vi.useFakeTimers() + try { + const result = await startRenewalLossAtBoundary('artifacts', { + failRelease: true, + }) + await expect(result.terminal).rejects.toThrow('writer renewal was lost') + expect(result.append).not.toHaveBeenCalled() + expect(result.checkpoints.releases).toBe(1) + } finally { + vi.useRealTimers() + } + }) + + it.each(['resolve', 'reject'] as const)( + 'preserves renewal loss when checkpoint append %ss', + async (outcome) => { + vi.useFakeTimers() + try { + const result = await startRenewalLossDuringAppend(outcome) + await expect(result.terminal).rejects.toThrow('writer renewal was lost') + expect( + result.f.events.some((event) => event.method === 'handle.snapshot'), + ).toBe(false) + expect( + result.f.events.filter((event) => event.method === 'destroy'), + ).toHaveLength(1) + expect(result.checkpoints.releases).toBe(1) + } finally { + vi.useRealTimers() + } + }, + ) + + it('keeps one renewal timer and stops it after a successful terminal snapshot', async () => { + vi.useFakeTimers() + try { + const f = fixture() + const checkpoints = new RenewalCountingCheckpointStore({ + leaseDurationMs: 120_000, + renewAfterMs: 10, + }) + let finishAdapter: () => void = () => {} + const waitForFinish = new Promise((resolve) => { + finishAdapter = resolve + }) + const stalledAdapter: AnyTextAdapter = { + ...adapter, + chatStream: async function* () { + yield { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: 1, + } + await waitForFinish + yield { + type: EventType.RUN_FINISHED, + runId: 'run-1', + threadId: 'thread-1', + finishReason: 'stop', + timestamp: 1, + } + }, + } + const iterator = chat({ + adapter: stalledAdapter, + messages: [{ role: 'user', content: 'hello' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [ + withPersistence(f.persistence), + withSandbox(f.definition, { + instances: f.instances, + snapshots: { persistence: f.persistence, checkpoints }, + }), + ], + })[Symbol.asyncIterator]() + + await iterator.next() + await vi.advanceTimersByTimeAsync(10) + expect(checkpoints.renewals).toBe(1) + expect(vi.getTimerCount()).toBe(1) + + finishAdapter() + await drainIterator(iterator) + expect(vi.getTimerCount()).toBe(0) + await vi.advanceTimersByTimeAsync(100) + expect(checkpoints.renewals).toBe(1) + const lease = await checkpoints.acquireWriter('thread-1') + await lease.release() + } finally { + vi.useRealTimers() + } + }) + + it('waits for one in-flight renewal before concurrent stops release the writer', async () => { + vi.useFakeTimers() + try { + const f = fixture() + const checkpoints = new DeferredRenewalCheckpointStore({ + leaseDurationMs: 120_000, + renewAfterMs: 10, + }) + const ctx = makeMiddlewareCtx({ threadId: 'thread-1', runId: 'run-1' }) + const persistence = withPersistence(f.persistence) + const sandbox = withSandbox(f.definition, { + instances: f.instances, + snapshots: { persistence: f.persistence, checkpoints }, + }) + await persistence.setup?.(ctx) + await sandbox.setup?.(ctx) + + vi.advanceTimersByTime(10) + await Promise.resolve() + expect(checkpoints.renewals).toBe(1) + + let pauseStopped = false + let abortStopped = false + const pause = Promise.resolve( + sandbox.onChunk?.(ctx, { + type: EventType.RUN_FINISHED, + runId: 'run-1', + threadId: 'thread-1', + finishReason: 'stop', + timestamp: 1, + outcome: { type: 'interrupt', interrupts: [] }, + }), + ).then(() => { + pauseStopped = true + }) + const abort = Promise.resolve( + sandbox.onAbort?.(ctx, { + cancelRequested: true, + reason: 'test abort', + duration: 1, + }), + ).then(() => { + abortStopped = true + }) + + await Promise.resolve() + expect(pauseStopped).toBe(false) + expect(abortStopped).toBe(false) + expect(checkpoints.releases).toBe(0) + + checkpoints.renewalGate.resolve() + await checkpoints.releaseStarted.promise + expect(checkpoints.releases).toBe(1) + expect(pauseStopped).toBe(false) + expect(abortStopped).toBe(false) + expect(vi.getTimerCount()).toBe(0) + + checkpoints.releaseGate.resolve() + await Promise.all([pause, abort]) + await vi.advanceTimersByTimeAsync(100) + expect(checkpoints.renewals).toBe(1) + expect(checkpoints.releases).toBe(1) + expect(vi.getTimerCount()).toBe(0) + } finally { + vi.useRealTimers() + } + }) + + it('rejects a competing writer before sandbox setup', async () => { + const f = fixture() + const writer = await f.checkpoints.acquireWriter('thread-1') + try { + await expect( + drain( + chat({ + adapter, + messages: [{ role: 'user', content: 'hello' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [ + withPersistence(f.persistence), + withSandbox(f.definition, { + instances: f.instances, + snapshots: { + persistence: f.persistence, + checkpoints: f.checkpoints, + }, + }), + ], + }), + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_WRITER_CONFLICT' }) + expect(f.events).toEqual([]) + } finally { + await writer.release() + } + }) + + it('releases the writer without capture when a run aborts', async () => { + const f = fixture() + const ctx = makeMiddlewareCtx({ threadId: 'thread-1', runId: 'run-1' }) + const persistence = withPersistence(f.persistence) + const sandbox = withSandbox(f.definition, { + instances: f.instances, + snapshots: { persistence: f.persistence, checkpoints: f.checkpoints }, + }) + const append = vi.spyOn(f.checkpoints, 'append') + await persistence.setup?.(ctx) + await sandbox.setup?.(ctx) + await sandbox.onAbort?.(ctx, { + cancelRequested: true, + reason: 'test abort', + duration: 1, + }) + expect(append).not.toHaveBeenCalled() + const lease = await f.checkpoints.acquireWriter('thread-1') + await lease.release() + }) + + it('releases the writer without capture when an actionable interrupt pauses', async () => { + const f = fixture() + const ctx = makeMiddlewareCtx({ threadId: 'thread-1', runId: 'run-1' }) + const persistence = withPersistence(f.persistence) + const sandbox = withSandbox(f.definition, { + instances: f.instances, + snapshots: { persistence: f.persistence, checkpoints: f.checkpoints }, + }) + const append = vi.spyOn(f.checkpoints, 'append') + const loadThread = vi.spyOn(f.persistence.stores.messages, 'loadThread') + await persistence.setup?.(ctx) + await sandbox.setup?.(ctx) + await sandbox.onChunk?.(ctx, { + type: EventType.RUN_FINISHED, + runId: 'run-1', + threadId: 'thread-1', + finishReason: 'stop', + timestamp: 1, + outcome: { type: 'interrupt', interrupts: [] }, + }) + expect(append).not.toHaveBeenCalled() + const lease = await f.checkpoints.acquireWriter('thread-1') + await lease.release() + const finish = { duration: 1, finishReason: 'stop', content: '' } + await persistence.onFinish?.(ctx, finish) + const persistenceLoads = loadThread.mock.calls.length + await sandbox.onFinish?.(ctx, finish) + expect(loadThread).toHaveBeenCalledTimes(persistenceLoads) + expect(append).not.toHaveBeenCalled() + }) + + it('closes portable ownership for an interrupt outcome streamed through chat', async () => { + vi.useFakeTimers() + try { + const f = fixture() + const checkpoints = new RenewalCountingCheckpointStore({ + leaseDurationMs: 120_000, + renewAfterMs: 10, + }) + const interruptAdapter: AnyTextAdapter = { + ...adapter, + chatStream: async function* () { + yield { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: 1, + } + yield { + type: EventType.RUN_FINISHED, + runId: 'run-1', + threadId: 'thread-1', + finishReason: 'stop', + timestamp: 1, + outcome: { type: 'interrupt', interrupts: [] }, + } + }, + } + const append = vi.spyOn(checkpoints, 'append') + await drain( + chat({ + adapter: interruptAdapter, + messages: [{ role: 'user', content: 'hello' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [ + withPersistence(f.persistence), + withSandbox(f.definition, { + instances: f.instances, + snapshots: { persistence: f.persistence, checkpoints }, + }), + ], + }), + ) + await vi.advanceTimersByTimeAsync(100) + expect(checkpoints.renewals).toBe(0) + expect(append).not.toHaveBeenCalled() + } finally { + vi.useRealTimers() + } + }) + + it('releases portable ownership on a real durable disconnect and later skips capture', async () => { + const f = fixture() + const ctx = makeMiddlewareCtx({ threadId: 'thread-1', runId: 'run-1' }) + const listeners: Array<() => void | Promise> = [] + provideRunDisconnect(ctx, { + subscribe: (listener) => listeners.push(listener), + }) + const persistence = withPersistence(f.persistence) + const sandbox = withSandbox(f.definition, { + instances: f.instances, + runs: f.persistence.stores.runs, + durability: { adapter: fakeLog(), detachOnDisconnect: true }, + snapshots: { persistence: f.persistence, checkpoints: f.checkpoints }, + }) + const append = vi.spyOn(f.checkpoints, 'append') + const loadThread = vi.spyOn(f.persistence.stores.messages, 'loadThread') + await persistence.setup?.(ctx) + await sandbox.setup?.(ctx) + await Promise.all(listeners.map((listener) => listener())) + const lease = await f.checkpoints.acquireWriter('thread-1') + await lease.release() + const finish = { duration: 1, finishReason: 'stop', content: '' } + await persistence.onFinish?.(ctx, finish) + const persistenceLoads = loadThread.mock.calls.length + await sandbox.onFinish?.(ctx, finish) + expect(loadThread).toHaveBeenCalledTimes(persistenceLoads) + expect(append).not.toHaveBeenCalled() + }) + + it('closes portable ownership before awaited disconnect bookkeeping can race onFinish', async () => { + const detachWriteGate = deferred() + const detachWriteStarted = deferred() + const f = fixture() + const originalUpdate = f.persistence.stores.runs.update.bind( + f.persistence.stores.runs, + ) + let blockDetachWrite = false + vi.spyOn(f.persistence.stores.runs, 'update').mockImplementation( + async (runId, patch) => { + if (blockDetachWrite && patch.detachedSince !== undefined) { + detachWriteStarted.resolve() + await detachWriteGate.promise + } + await originalUpdate(runId, patch) + }, + ) + const append = vi.spyOn(f.checkpoints, 'append') + const blobPut = vi.spyOn(f.persistence.stores.blobs, 'put') + const durableLog = fakeLog() + const ctx = makeMiddlewareCtx({ threadId: 'thread-1', runId: 'run-1' }) + const listeners: Array<() => void | Promise> = [] + provideRunDisconnect(ctx, { + subscribe: (listener) => listeners.push(listener), + }) + const persistence = withPersistence(f.persistence) + const sandbox = snapshotMiddleware(f, { durability: durableLog }) + await persistence.setup?.(ctx) + await sandbox.setup?.(ctx) + + blockDetachWrite = true + const disconnect = Promise.all(listeners.map((listener) => listener())) + await detachWriteStarted.promise + const finish = { duration: 1, finishReason: 'stop', content: '' } + await persistence.onFinish?.(ctx, finish) + await sandbox.onFinish?.(ctx, finish) + + expect(f.events.some((event) => event.method === 'fs.list')).toBe(false) + expect(blobPut).not.toHaveBeenCalled() + expect(append).not.toHaveBeenCalled() + + detachWriteGate.resolve() + await disconnect + }) + + it('waits for an active checkpoint append before disconnect releases its lease', async () => { + const appendStarted = deferred() + const appendGate = deferred() + const checkpoints = new ReleaseCountingCheckpointStore() + const f = fixture({ checkpoints }) + const ctx = makeMiddlewareCtx({ threadId: 'thread-1', runId: 'run-1' }) + const listeners: Array<() => void | Promise> = [] + provideRunDisconnect(ctx, { + subscribe: (listener) => listeners.push(listener), + }) + const persistence = withPersistence(f.persistence) + const sandbox = snapshotMiddleware(f, { + checkpoints, + durability: fakeLog(), + }) + await persistence.setup?.(ctx) + await sandbox.setup?.(ctx) + const finish = { duration: 1, finishReason: 'stop', content: '' } + await persistence.onFinish?.(ctx, finish) + const append = checkpoints.append.bind(checkpoints) + vi.spyOn(checkpoints, 'append').mockImplementation(async (input) => { + appendStarted.resolve() + await appendGate.promise + return append(input) + }) + + const terminal = Promise.resolve(sandbox.onFinish?.(ctx, finish)) + await appendStarted.promise + const disconnect = Promise.all(listeners.map((listener) => listener())) + await flushMicrotasks() + expect(checkpoints.releases).toBe(0) + + appendGate.resolve() + await terminal + await disconnect + expect(checkpoints.releases).toBe(1) + expect(await checkpoints.getHead('thread-1')).toBe('checkpoint-run-1') + }) + + it('stops the file watcher after detach and emits no later file event', async () => { + const watch: WatchProbe = { stops: 0 } + let fileEvents = 0 + const f = fixture({ watch }) + const definition = defineSandbox({ + id: 'fixture', + provider: f.provider, + workspace: { source: { type: 'none' } }, + fileEvents: true, + hooks: { onFile: () => void fileEvents++ }, + }) + const ctx = makeMiddlewareCtx({ threadId: 'thread-1', runId: 'run-1' }) + const listeners: Array<() => void | Promise> = [] + provideRunDisconnect(ctx, { + subscribe: (listener) => listeners.push(listener), + }) + const persistence = withPersistence(f.persistence) + const sandbox = snapshotMiddleware(f, { + definition, + durability: fakeLog(), + }) + await persistence.setup?.(ctx) + await sandbox.setup?.(ctx) + await Promise.all(listeners.map((listener) => listener())) + await persistence.onFinish?.(ctx, { + duration: 1, + finishReason: 'stop', + content: '', + }) + await sandbox.onFinish?.(ctx, { + duration: 1, + finishReason: 'stop', + content: '', + }) + watch.emit?.({ type: 'change', path: '/workspace/app.ts' }) + await flushMicrotasks() + + expect(watch.stops).toBe(1) + expect(fileEvents).toBe(0) + }) + + it('stops the file watcher when an actionable interrupt pauses the run', async () => { + const watch: WatchProbe = { stops: 0 } + let fileEvents = 0 + const f = fixture({ watch }) + const definition = defineSandbox({ + id: 'fixture', + provider: f.provider, + workspace: { source: { type: 'none' } }, + fileEvents: true, + hooks: { onFile: () => void fileEvents++ }, + }) + const ctx = makeMiddlewareCtx({ threadId: 'thread-1', runId: 'run-1' }) + const persistence = withPersistence(f.persistence) + const sandbox = snapshotMiddleware(f, { definition }) + await persistence.setup?.(ctx) + await sandbox.setup?.(ctx) + await sandbox.onChunk?.(ctx, { + type: EventType.RUN_FINISHED, + runId: 'run-1', + threadId: 'thread-1', + finishReason: 'stop', + timestamp: 1, + outcome: { type: 'interrupt', interrupts: [] }, + }) + watch.emit?.({ type: 'change', path: '/workspace/app.ts' }) + await flushMicrotasks() + + expect(watch.stops).toBe(1) + expect(fileEvents).toBe(0) + }) + + it('still destroys on abort when releasing the writer fails', async () => { + const f = fixture() + const checkpoints = new ReleaseFailingCheckpointStore() + const ctx = makeMiddlewareCtx({ threadId: 'thread-1', runId: 'run-1' }) + const persistence = withPersistence(f.persistence) + const sandbox = withSandbox(f.definition, { + instances: f.instances, + snapshots: { persistence: f.persistence, checkpoints }, + }) + await persistence.setup?.(ctx) + await sandbox.setup?.(ctx) + await expect( + sandbox.onAbort?.(ctx, { + cancelRequested: true, + reason: 'test abort', + duration: 1, + }), + ).rejects.toThrow('writer release failed') + expect(f.events.some((event) => event.method === 'destroy')).toBe(true) + }) + + it('does not capture or append when a created run adapter fails', async () => { + const f = fixture() + const append = vi.spyOn(f.checkpoints, 'append') + const failingAdapter: AnyTextAdapter = { + ...adapter, + chatStream: () => failingStream('adapter stream failed'), + } + + await expect( + drain( + chat({ + adapter: failingAdapter, + messages: [{ role: 'user', content: 'hello' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [ + withPersistence(f.persistence), + withSandbox(afterRunDefinition(f), { + instances: f.instances, + snapshots: { + persistence: f.persistence, + checkpoints: f.checkpoints, + }, + }), + ], + }), + ), + ).rejects.toThrow('adapter stream failed') + + expect(append).not.toHaveBeenCalled() + expect(f.events.some((event) => event.method === 'fs.list')).toBe(false) + expect(f.events.filter((event) => event.method === 'destroy')).toHaveLength( + 1, + ) + const lease = await f.checkpoints.acquireWriter('thread-1') + await lease.release() + }) + + it('preserves an adapter failure when writer release also fails after cleanup', async () => { + const f = fixture() + const checkpoints = new ReleaseFailingCheckpointStore() + const failingAdapter: AnyTextAdapter = { + ...adapter, + chatStream: () => failingStream('adapter stream failed'), + } + + await expect( + drain( + chat({ + adapter: failingAdapter, + messages: [{ role: 'user', content: 'hello' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [ + withPersistence(f.persistence), + withSandbox(afterRunDefinition(f), { + instances: f.instances, + snapshots: { persistence: f.persistence, checkpoints }, + }), + ], + }), + ), + ).rejects.toThrow('adapter stream failed') + + expect(checkpoints.releases).toBe(1) + expect(f.events.filter((event) => event.method === 'destroy')).toHaveLength( + 1, + ) + }) + + it('preserves an append failure when the terminal writer release also fails', async () => { + const f = fixture() + const checkpoints = new ReleaseFailingCheckpointStore() + vi.spyOn(checkpoints, 'append').mockRejectedValue( + new Error('append failed'), + ) + await expect( + drain( + chat({ + adapter, + messages: [{ role: 'user', content: 'hello' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [ + withPersistence(f.persistence), + withSandbox(f.definition, { + instances: f.instances, + snapshots: { persistence: f.persistence, checkpoints }, + }), + ], + }), + ), + ).rejects.toThrow('append failed') + }) + + it('preserves an artifact capture failure when terminal lease release also fails', async () => { + const f = fixture() + const checkpoints = new ReleaseFailingCheckpointStore() + vi.spyOn(f.persistence.stores.artifacts, 'listForThread').mockRejectedValue( + new Error('artifact capture failed'), + ) + + await expect(runTerminalSnapshot(f, checkpoints)).rejects.toThrow( + 'artifact capture failed', + ) + + expect(checkpoints.releases).toBe(1) + expect(f.events.some((event) => event.method === 'handle.snapshot')).toBe( + false, + ) + expect(f.events.filter((event) => event.method === 'destroy')).toHaveLength( + 1, + ) + }) + + it('preserves a checkpoint append failure when terminal lease release also fails', async () => { + const f = fixture() + const checkpoints = new ReleaseFailingCheckpointStore() + vi.spyOn(checkpoints, 'append').mockRejectedValue( + new Error('checkpoint append failed'), + ) + + await expect(runTerminalSnapshot(f, checkpoints)).rejects.toThrow( + 'checkpoint append failed', + ) + + expect(checkpoints.releases).toBe(1) + expect(f.events.some((event) => event.method === 'handle.snapshot')).toBe( + false, + ) + expect(f.events.filter((event) => event.method === 'destroy')).toHaveLength( + 1, + ) + }) + + it('captures portable state before native after-run snapshot and destroy', async () => { + const f = fixture() + const definition = defineSandbox({ + id: 'fixture', + provider: f.provider, + lifecycle: { snapshot: 'after-run', destroyOnComplete: true }, + workspace: { source: { type: 'none' } }, + fileEvents: false, + }) + const appendGate = deferred() + const appendStarted = deferred() + const append = f.checkpoints.append.bind(f.checkpoints) + vi.spyOn(f.checkpoints, 'append').mockImplementation(async (input) => { + appendStarted.resolve() + await appendGate.promise + const checkpoint = await append(input) + f.events.push({ + method: 'checkpoint.append.complete', + order: f.events.length + 1, + }) + return checkpoint + }) + const terminal = drain( + chat({ + adapter, + messages: [{ role: 'user', content: 'hello' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [ + withPersistence(f.persistence), + withSandbox(definition, { + instances: f.instances, + snapshots: { + persistence: f.persistence, + checkpoints: f.checkpoints, + }, + }), + ], + }), + ) + await appendStarted.promise + expect(f.events.some((event) => event.method === 'handle.snapshot')).toBe( + false, + ) + expect(f.events.some((event) => event.method === 'destroy')).toBe(false) + appendGate.resolve() + await terminal + const nativeSnapshot = f.events.findIndex( + (event) => event.method === 'handle.snapshot', + ) + const destroy = f.events.findIndex((event) => event.method === 'destroy') + const capture = f.events.findIndex((event) => event.method === 'fs.list') + const appendComplete = f.events.findIndex( + (event) => event.method === 'checkpoint.append.complete', + ) + expect(capture).toBeGreaterThanOrEqual(0) + expect(appendComplete).toBeGreaterThan(capture) + expect(nativeSnapshot).toBeGreaterThan(appendComplete) + expect(nativeSnapshot).toBeGreaterThanOrEqual(0) + expect(destroy).toBeGreaterThan(nativeSnapshot) + }) + + it('destroys after native snapshot failure and preserves that failure', async () => { + const f = fixture({ + nativeSnapshotError: new Error('native snapshot failed'), + }) + + await expect( + drain( + chat({ + adapter, + messages: [{ role: 'user', content: 'hello' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [ + withPersistence(f.persistence), + snapshotMiddleware(f, { definition: afterRunDefinition(f) }), + ], + }), + ), + ).rejects.toThrow('native snapshot failed') + + expect(await f.checkpoints.getHead('thread-1')).toBe('checkpoint-run-1') + expect( + f.events.filter((event) => event.method === 'handle.snapshot'), + ).toHaveLength(1) + expect(f.events.filter((event) => event.method === 'destroy')).toHaveLength( + 1, + ) + }) + + it('publishes a portable checkpoint only after persistence completes', async () => { + const f = fixture() + const completion = vi.spyOn(f.persistence.stores.messages, 'loadThread') + const append = vi.spyOn(f.checkpoints, 'append') + await drain( + chat({ + adapter, + messages: [{ role: 'user', content: 'hello' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [ + withPersistence(f.persistence), + withSandbox(f.definition, { + instances: f.instances, + snapshots: { + persistence: f.persistence, + checkpoints: f.checkpoints, + }, + }), + ], + }), + ) + expect(append).toHaveBeenCalledTimes(1) + expect(completion).toHaveBeenCalledBefore(append) + expect(await f.checkpoints.getHead('thread-1')).toBe('checkpoint-run-1') + }) + + it('persists the completed transcript, files, and artifacts after completion', async () => { + type SnapshotStage = + | 'persistence-transcript-commit' + | 'persistence-completion' + | 'snapshot-transcript-read' + | 'file-capture' + | 'artifact-capture' + | 'checkpoint-append' + const stages: Array = [] + const snapshots = await memorySandboxSnapshots() + const sourceArtifactBytes = new TextEncoder().encode('generated artifact') + const f = fixture({ + persistence: snapshots.persistence, + checkpoints: snapshots.checkpoints, + workspace: [ + { + path: '/workspace/app.ts', + type: 'file', + data: new TextEncoder().encode('export const app = true'), + }, + { path: '/workspace/empty-dir', type: 'dir' }, + ], + onWorkspaceList: (path) => { + if (path === '/workspace') stages.push('file-capture') + }, + }) + const sourceArtifactKey = 'artifacts/run-1/generated.txt' + await f.persistence.stores.blobs.put(sourceArtifactKey, sourceArtifactBytes) + await f.persistence.stores.artifacts.save({ + artifactId: 'generated.txt', + runId: 'run-1', + threadId: 'thread-1', + blobKey: sourceArtifactKey, + name: 'generated.txt', + mimeType: 'text/plain', + size: sourceArtifactBytes.byteLength, + createdAt: 1, + }) + + const saveThread = f.persistence.stores.messages.saveThread.bind( + f.persistence.stores.messages, + ) + vi.spyOn(f.persistence.stores.messages, 'saveThread').mockImplementation( + async (threadId, messages) => { + await saveThread(threadId, messages) + stages.push('persistence-transcript-commit') + }, + ) + const updateRun = f.persistence.stores.runs.update.bind( + f.persistence.stores.runs, + ) + vi.spyOn(f.persistence.stores.runs, 'update').mockImplementation( + async (runId, patch) => { + await updateRun(runId, patch) + if (patch.status === 'completed') stages.push('persistence-completion') + }, + ) + const loadThread = f.persistence.stores.messages.loadThread.bind( + f.persistence.stores.messages, + ) + vi.spyOn(f.persistence.stores.messages, 'loadThread').mockImplementation( + async (threadId) => { + const messages = await loadThread(threadId) + if (stages.includes('persistence-completion')) + stages.push('snapshot-transcript-read') + return messages + }, + ) + const listForThread = f.persistence.stores.artifacts.listForThread.bind( + f.persistence.stores.artifacts, + ) + vi.spyOn( + f.persistence.stores.artifacts, + 'listForThread', + ).mockImplementation(async (threadId) => { + const artifacts = await listForThread(threadId) + stages.push('artifact-capture') + return artifacts + }) + const append = f.checkpoints.append.bind(f.checkpoints) + vi.spyOn(f.checkpoints, 'append').mockImplementation(async (input) => { + const result = await append(input) + stages.push('checkpoint-append') + return result + }) + + const finalAdapter: AnyTextAdapter = { + ...adapter, + chatStream: async function* (): AsyncGenerator { + yield { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: 1, + } + yield { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'message-1', + delta: 'The completed answer', + } + yield { + type: EventType.RUN_FINISHED, + runId: 'run-1', + threadId: 'thread-1', + finishReason: 'stop', + timestamp: 1, + } + }, + } + await drain( + chat({ + adapter: finalAdapter, + messages: [{ role: 'user', content: 'Persist this answer' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [ + withPersistence(f.persistence), + withSandbox(afterRunDefinition(f), { + instances: f.instances, + snapshots: { + persistence: f.persistence, + checkpoints: f.checkpoints, + }, + }), + ], + }), + ) + + const checkpoint = await f.checkpoints.get('checkpoint-run-1') + expect(checkpoint).not.toBeNull() + expect(checkpoint?.conversation).toEqual([ + { role: 'user', content: 'Persist this answer' }, + { role: 'assistant', content: 'The completed answer' }, + ]) + expect(checkpoint?.files).toEqual([ + { kind: 'file', path: 'app.ts', size: 23, blobKey: expect.any(String) }, + { kind: 'dir', path: 'empty-dir' }, + ]) + const capturedFile = checkpoint?.files.find( + (entry) => entry.kind === 'file' && entry.path === 'app.ts', + ) + if (!capturedFile || capturedFile.kind !== 'file') + throw new Error('Expected the app.ts file to be captured') + const capturedFileBlob = await f.persistence.stores.blobs.get( + capturedFile.blobKey, + ) + if (!capturedFileBlob) + throw new Error('Expected the captured app.ts blob to exist') + expect(await capturedFileBlob.text()).toBe('export const app = true') + expect(checkpoint?.artifacts).toEqual([ + { + artifactId: 'generated.txt', + name: 'generated.txt', + mimeType: 'text/plain', + size: sourceArtifactBytes.byteLength, + blobKey: expect.stringMatching(/^sandbox-artifacts\/sha256\//), + createdAt: 1, + }, + ]) + const artifactBlobKey = checkpoint?.artifacts[0]?.blobKey + expect(artifactBlobKey).toBeDefined() + const capturedArtifactBlob = await f.persistence.stores.blobs.get( + artifactBlobKey ?? '', + ) + if (!capturedArtifactBlob) + throw new Error('Expected the captured artifact blob to exist') + expect(await capturedArtifactBlob.text()).toBe('generated artifact') + const terminalCommit = stages.lastIndexOf('persistence-transcript-commit') + expect(stages.slice(0, terminalCommit)).not.toContain('file-capture') + expect(stages.slice(terminalCommit)).toEqual([ + 'persistence-transcript-commit', + 'persistence-completion', + 'snapshot-transcript-read', + 'file-capture', + 'artifact-capture', + 'checkpoint-append', + ]) + }) + + it('created run exposes SandboxCapability', async () => { + const f = fixture() + const ctx = makeMiddlewareCtx({ threadId: 'thread-1', runId: 'run-1' }) + const middleware = withSandbox(f.definition, { instances: f.instances }) + await middleware.setup?.(ctx) + expect(ctx.capabilities.has(SandboxCapability)).toBe(true) + const created = f.events.find((event) => event.method === 'create') + expect(created).toBeDefined() + expect(ctx.get(SandboxCapability).id).toBe(created?.id) + }) + + it('restores a portable head before onReady exposes a freshly created sandbox', async () => { + const f = fixture() + const bytes = new TextEncoder().encode('restored') + const blobKey = await sandboxFileBlobKey(bytes) + await f.persistence.stores.blobs.put(blobKey, bytes) + await seedCheckpoint(f, { + files: [{ path: 'app.ts', blobKey, size: 8 }], + }) + const definition = defineSandbox({ + id: 'fixture', + provider: f.provider, + workspace: { source: { type: 'none' } }, + fileEvents: false, + hooks: { + onReady: async () => { + f.events.push({ + method: 'onReady', + order: f.events.length + 1, + }) + }, + }, + }) + await drain( + chat({ + adapter, + messages: [{ role: 'user', content: 'hello' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [ + withPersistence(f.persistence), + withSandbox(definition, { + instances: f.instances, + snapshots: { + persistence: f.persistence, + checkpoints: f.checkpoints, + }, + }), + ], + }), + ) + const restore = f.events.findIndex((event) => event.method === 'fs.write') + const ready = f.events.findIndex((event) => event.method === 'onReady') + expect(restore).toBeGreaterThanOrEqual(0) + expect(ready).toBeGreaterThan(restore) + }) + + it('restores before onReady after a native provider restore', async () => { + const f = fixture() + const bytes = new TextEncoder().encode('restored') + const blobKey = await sandboxFileBlobKey(bytes) + await f.persistence.stores.blobs.put(blobKey, bytes) + await seedCheckpoint(f, { + files: [{ path: 'app.ts', blobKey, size: 8 }], + }) + const key = f.definition.key({ + threadId: 'thread-1', + runId: 'run-1', + store: f.instances, + }) + await f.instances.upsert({ + key, + provider: 'fixture', + providerSandboxId: 'gone', + latestSnapshotId: 'native-snapshot', + threadId: 'thread-1', + updatedAt: Date.now(), + }) + f.provider.resume = async () => null + const definition = defineSandbox({ + id: 'fixture', + provider: f.provider, + workspace: { source: { type: 'none' } }, + fileEvents: false, + hooks: { + onReady: async () => { + f.events.push({ method: 'onReady', order: f.events.length + 1 }) + }, + }, + }) + await drain( + chat({ + adapter, + messages: [{ role: 'user', content: 'hello' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [ + withPersistence(f.persistence), + withSandbox(definition, { + instances: f.instances, + snapshots: { + persistence: f.persistence, + checkpoints: f.checkpoints, + }, + }), + ], + }), + ) + const nativeRestore = f.events.findIndex( + (event) => event.method === 'restoreSnapshot', + ) + const fileRestore = f.events.findIndex( + (event) => event.method === 'fs.write', + ) + const ready = f.events.findIndex((event) => event.method === 'onReady') + expect(nativeRestore).toBeGreaterThanOrEqual(0) + expect(fileRestore).toBeGreaterThan(nativeRestore) + expect(ready).toBeGreaterThan(fileRestore) + }) + + it('destroys only the private created sandbox and releases the lease when restore fails', async () => { + const f = fixture() + await seedCheckpoint(f, { + files: [ + { + path: 'app.ts', + blobKey: `sandbox-files/sha256/${'e'.repeat(64)}`, + size: 1, + }, + ], + }) + await expect( + drain( + chat({ + adapter, + messages: [{ role: 'user', content: 'hello' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [ + withPersistence(f.persistence), + withSandbox(f.definition, { + instances: f.instances, + snapshots: { + persistence: f.persistence, + checkpoints: f.checkpoints, + }, + }), + ], + }), + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_MISSING_BLOB' }) + expect(f.events.map((event) => event.method)).toEqual([ + 'create', + 'handle.snapshot', + 'fs.list', + 'destroy', + ]) + const lease = await f.checkpoints.acquireWriter('thread-1') + await lease.release() + }) + + it('destroys a private native-restored sandbox and releases its writer when portable restore fails', async () => { + const f = fixture() + await seedCheckpoint(f, { + files: [ + { + path: 'app.ts', + blobKey: `sandbox-files/sha256/${'e'.repeat(64)}`, + size: 1, + }, + ], + }) + const key = f.definition.key({ + threadId: 'thread-1', + runId: 'run-1', + store: f.instances, + }) + await f.instances.upsert({ + key, + provider: 'fixture', + providerSandboxId: 'gone', + latestSnapshotId: 'native-snapshot', + threadId: 'thread-1', + updatedAt: Date.now(), + }) + f.provider.resume = async () => null + + await expect( + drain( + chat({ + adapter, + messages: [{ role: 'user', content: 'hello' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [ + withPersistence(f.persistence), + withSandbox(f.definition, { + instances: f.instances, + snapshots: { + persistence: f.persistence, + checkpoints: f.checkpoints, + }, + }), + ], + }), + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_MISSING_BLOB' }) + + expect(f.events.map((event) => event.method)).toEqual([ + 'restoreSnapshot', + 'fs.list', + 'destroy', + ]) + const lease = await f.checkpoints.acquireWriter('thread-1') + await lease.release() + }) + + it.each([false, true])( + 'fails and cleans up when the checkpoint head is missing (%s)', + async (nativeRestore) => { + const f = fixture() + class MissingHeadCheckpointStore extends InMemorySandboxCheckpointStore { + override async getHead(): Promise { + return 'missing-head' + } + + override async get() { + return null + } + } + const checkpoints: SandboxCheckpointStore = + new MissingHeadCheckpointStore() + if (nativeRestore) { + const key = f.definition.key({ + threadId: 'thread-1', + runId: 'run-1', + store: f.instances, + }) + await f.instances.upsert({ + key, + provider: 'fixture', + providerSandboxId: 'gone', + latestSnapshotId: 'native-snapshot', + threadId: 'thread-1', + updatedAt: Date.now(), + }) + f.provider.resume = async () => null + } + + await expect( + drain( + chat({ + adapter, + messages: [{ role: 'user', content: 'hello' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [ + withPersistence(f.persistence), + withSandbox(f.definition, { + instances: f.instances, + snapshots: { + persistence: f.persistence, + checkpoints, + }, + }), + ], + }), + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_CHECKPOINT_NOT_FOUND' }) + expect(f.events.map((event) => event.method)).toContain('destroy') + const lease = await checkpoints.acquireWriter('thread-1') + await lease.release() + }, + ) + + it('does not destroy a previously live resumed sandbox when setup fails after resume', async () => { + const f = fixture() + const definition = defineSandbox({ + id: 'fixture', + provider: f.provider, + lifecycle: { reuse: 'thread', snapshot: 'none' }, + workspace: { source: { type: 'none' } }, + fileEvents: false, + hooks: { + onReady: async () => { + throw new Error('onReady failed') + }, + }, + }) + const key = definition.key({ + threadId: 'thread-1', + runId: 'run-1', + store: f.instances, + }) + await f.instances.upsert({ + key, + provider: 'fixture', + providerSandboxId: 'still-live', + threadId: 'thread-1', + updatedAt: Date.now(), + }) + + await expect( + drain( + chat({ + adapter, + messages: [{ role: 'user', content: 'hello' }], + runId: 'run-1', + threadId: 'thread-1', + middleware: [ + withPersistence(f.persistence), + withSandbox(definition, { + instances: f.instances, + snapshots: { + persistence: f.persistence, + checkpoints: f.checkpoints, + }, + }), + ], + }), + ), + ).rejects.toThrow('onReady failed') + + expect(f.events.map((event) => event.method)).toEqual(['resume']) + const lease = await f.checkpoints.acquireWriter('thread-1') + await lease.release() + }) + + it('seeded resumed instance reports resumed without snapshot lookup', async () => { + const f = fixture() + const definition = defineSandbox({ + id: 'fixture', + provider: f.provider, + lifecycle: { reuse: 'thread', snapshot: 'none' }, + workspace: { source: { type: 'none' } }, + fileEvents: false, + }) + const key = definition.key({ + threadId: 'thread-1', + runId: 'run-1', + store: f.instances, + }) + await f.instances.upsert({ + key, + provider: 'fixture', + providerSandboxId: 'native-1', + latestSnapshotId: 'checkpoint-1', + threadId: 'thread-1', + updatedAt: Date.now(), + }) + const middleware = [ + withPersistence(f.persistence), + withSandbox(definition, { + instances: f.instances, + snapshots: { persistence: f.persistence, checkpoints: f.checkpoints }, + }), + ] + const getHead = vi.spyOn(f.checkpoints, 'getHead') + const get = vi.spyOn(f.checkpoints, 'get') + const loadThread = vi.spyOn(f.persistence.stores.messages, 'loadThread') + const listForThread = vi + .spyOn(f.persistence.stores.artifacts, 'listForThread') + .mockRejectedValue(new Error('artifacts touched')) + const blobGet = vi + .spyOn(f.persistence.stores.blobs, 'get') + .mockRejectedValue(new Error('blobs touched')) + const blobHead = vi + .spyOn(f.persistence.stores.blobs, 'head') + .mockRejectedValue(new Error('blobs touched')) + const blobPut = vi + .spyOn(f.persistence.stores.blobs, 'put') + .mockRejectedValue(new Error('blobs touched')) + let baseline = 0 + const liveAdapter: AnyTextAdapter = { + ...adapter, + chatStream: async function* () { + baseline = loadThread.mock.calls.length + yield* failingStream('stop') + }, + } + await expect( + drain( + chat({ + adapter: liveAdapter, + messages: [{ role: 'user', content: 'hello' }], + runId: 'run-1', + threadId: 'thread-1', + middleware, + }), + ), + ).rejects.toThrow('stop') + expect(loadThread).toHaveBeenCalledTimes(baseline) + expect(f.events.map((event) => event.method)).toEqual(['resume']) + expect(f.events.some((event) => event.method === 'restoreSnapshot')).toBe( + false, + ) + expect(getHead).not.toHaveBeenCalled() + expect(get).not.toHaveBeenCalled() + expect(listForThread).not.toHaveBeenCalled() + expect(blobGet).not.toHaveBeenCalled() + expect(blobHead).not.toHaveBeenCalled() + expect(blobPut).not.toHaveBeenCalled() + }) +}) diff --git a/packages/ai-sandbox/tests/snapshot-operations.test-d.ts b/packages/ai-sandbox/tests/snapshot-operations.test-d.ts new file mode 100644 index 0000000000..dc00aeb8b1 --- /dev/null +++ b/packages/ai-sandbox/tests/snapshot-operations.test-d.ts @@ -0,0 +1,33 @@ +import { expectTypeOf } from 'vitest' +import { memorySandboxSnapshots, SandboxSnapshotError } from '../src' +import type { SandboxSnapshotErrorCode, SandboxSnapshots } from '../src' + +type ExpectedSandboxSnapshotErrorCode = + | 'SANDBOX_SNAPSHOT_MISSING_REUSABLE_SANDBOX' + | 'SANDBOX_SNAPSHOT_REUSE_NONE' + | 'SANDBOX_SNAPSHOT_MISSING_CHECKPOINT_ARTIFACT' + | 'SANDBOX_SNAPSHOT_FOREIGN_CHECKPOINT_ARTIFACT' + | 'SANDBOX_SNAPSHOT_INVALID_ARTIFACT_BYTES' + | 'SANDBOX_SNAPSHOT_FORK_UNAVAILABLE' + | 'SANDBOX_SNAPSHOT_INVALID_PATH' + | 'SANDBOX_SNAPSHOT_INVALID_WORKSPACE' + | 'SANDBOX_SNAPSHOT_LSTAT_REQUIRED' + | 'SANDBOX_SNAPSHOT_UNSUPPORTED_ENTRY' + | 'SANDBOX_SNAPSHOT_MISSING_BLOB' + | 'SANDBOX_SNAPSHOT_INVALID_BLOB' + | 'SANDBOX_SNAPSHOT_ARTIFACT_SUPPORT_REQUIRED' + | 'SANDBOX_SNAPSHOT_MISSING_ARTIFACT_BLOB' + +expectTypeOf().toEqualTypeOf() +const sourceError = new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_MISSING_REUSABLE_SANDBOX', + 'missing', +) +expectTypeOf(sourceError.code).toEqualTypeOf() + +async function assignActualMemorySnapshots(): Promise { + const snapshots = await memorySandboxSnapshots() + const structuralSnapshots: SandboxSnapshots = snapshots + void structuralSnapshots +} +void assignActualMemorySnapshots diff --git a/packages/ai-sandbox/tests/snapshot-operations.test.ts b/packages/ai-sandbox/tests/snapshot-operations.test.ts new file mode 100644 index 0000000000..09db241f71 --- /dev/null +++ b/packages/ai-sandbox/tests/snapshot-operations.test.ts @@ -0,0 +1,1338 @@ +import { InMemoryLockStore } from '@tanstack/ai/locks' +import { describe, expect, it, vi } from 'vitest' +import { + computeSandboxKey, + computeWorkspaceHash, + createSecrets, + defineSandbox, + forkFromSandboxSnapshot, + InMemorySandboxCheckpointStore, + InMemorySandboxInstanceStore, + memorySandboxSnapshots, + resolveSnapshotArtifact, + SandboxSnapshotError, + saveNamedSandboxSnapshot, +} from '../src' +import { makeFakeHandle, makeFakeProvider } from './fakes' +import type { + MemorySandboxSnapshots, + SandboxCheckpoint, + SandboxCheckpointStore, + SandboxCheckpointWriterLease, + SandboxDefinition, + SandboxEnsureContext, + SandboxHandle, + SandboxSnapshotPolicy, + SandboxSnapshots, + WorkspaceDefinition, +} from '../src' +import type { FakeProvider } from './fakes' + +const THREAD = 'thread' +const RUN = 'run' +const LABEL = 'named' +const invalidArtifactKey = + 'sandbox-artifacts/sha256/0000000000000000000000000000000000000000000000000000000000000000' + +function deferred() { + let resolve: (value: T) => void = () => {} + const promise = new Promise((complete) => { + resolve = complete + }) + return { promise, resolve } +} + +class LeaseProbeStore extends InMemorySandboxCheckpointStore { + renewals = 0 + activeRenewals = 0 + maxActiveRenewals = 0 + releases = 0 + appends = 0 + renewalGate?: Promise + renewalError?: Error + releaseError?: Error + appendError?: Error + readonly renewalStarted = deferred() + readonly renewalFinished = deferred() + + constructor() { + super({ leaseDurationMs: 100, renewAfterMs: 10 }) + } + + override async acquireWriter( + threadId: string, + ): Promise { + const lease = await super.acquireWriter(threadId) + return { + ...lease, + renew: async () => { + this.renewals++ + this.activeRenewals++ + this.maxActiveRenewals = Math.max( + this.maxActiveRenewals, + this.activeRenewals, + ) + this.renewalStarted.resolve() + try { + await this.renewalGate + if (this.renewalError) throw this.renewalError + return await lease.renew() + } finally { + this.activeRenewals-- + this.renewalFinished.resolve() + } + }, + release: async () => { + this.releases++ + await lease.release() + if (this.releaseError) throw this.releaseError + }, + } + } + + override async append( + input: Parameters[0], + ): Promise<{ headId: string }> { + this.appends++ + if (this.appendError) throw this.appendError + return super.append(input) + } +} + +class ForkProbeStore implements SandboxCheckpointStore { + releases = 0 + renewals = 0 + forks = 0 + releaseError?: Error + forkError?: Error + private readonly store: SandboxCheckpointStore + readonly get: SandboxCheckpointStore['get'] + readonly list: SandboxCheckpointStore['list'] + readonly getHead: SandboxCheckpointStore['getHead'] + readonly append: SandboxCheckpointStore['append'] + readonly deleteHead: SandboxCheckpointStore['deleteHead'] + readonly listBlobReferences: SandboxCheckpointStore['listBlobReferences'] + + constructor(store: SandboxCheckpointStore) { + this.store = store + this.get = store.get.bind(store) + this.list = store.list.bind(store) + this.getHead = store.getHead.bind(store) + this.append = store.append.bind(store) + this.deleteHead = store.deleteHead.bind(store) + this.listBlobReferences = store.listBlobReferences.bind(store) + } + + async acquireWriter(threadId: string): Promise { + const lease = await this.store.acquireWriter(threadId) + return { + ...lease, + renew: async () => { + this.renewals++ + return lease.renew() + }, + release: async () => { + this.releases++ + await lease.release() + if (this.releaseError) throw this.releaseError + }, + } + } + + async forkFromCheckpoint( + input: Parameters< + NonNullable + >[0], + ): Promise<{ checkpoint: SandboxCheckpoint }> { + this.forks++ + if (this.forkError) throw this.forkError + const fork = this.store.forkFromCheckpoint + if (!fork) throw new Error('test fork capability is missing') + return fork.call(this.store, input) + } +} + +type NamedFixture = { + definition: SandboxDefinition + instances: InMemorySandboxInstanceStore + provider: FakeProvider + memory: MemorySandboxSnapshots + snapshots: SandboxSnapshots + locks: InMemoryLockStore +} + +async function namedFixture( + options: { + checkpoints?: SandboxCheckpointStore + lifecycle?: Parameters[0]['lifecycle'] + policy?: SandboxSnapshotPolicy + seedInstance?: boolean + workspace?: WorkspaceDefinition + } = {}, +): Promise { + const memory = await memorySandboxSnapshots() + const instances = new InMemorySandboxInstanceStore() + const provider = makeFakeProvider() + const definition = defineSandbox({ + id: 'sandbox', + provider, + lifecycle: options.lifecycle, + workspace: options.workspace, + }) + if (options.seedInstance !== false) { + await instances.upsert({ + key: definition.key({ threadId: THREAD, runId: 'old' }), + provider: provider.name, + providerSandboxId: 'existing', + threadId: THREAD, + updatedAt: Date.now(), + }) + } + return { + definition, + instances, + provider, + memory, + snapshots: { + persistence: memory.persistence, + checkpoints: options.checkpoints ?? memory.checkpoints, + ...(options.policy === undefined ? {} : { policy: options.policy }), + }, + locks: new InMemoryLockStore(), + } +} + +type NamedSaveInput = Parameters[0] + +function namedSave( + fixture: NamedFixture, + changes: Partial = {}, +) { + const input: NamedSaveInput = { + definition: fixture.definition, + threadId: THREAD, + runId: RUN, + instances: fixture.instances, + snapshots: fixture.snapshots, + label: LABEL, + locks: fixture.locks, + ...changes, + } + return saveNamedSandboxSnapshot(input) +} + +function checkpoint( + input: Partial & + Pick, +): SandboxCheckpoint { + return { + id: input.id, + threadId: input.threadId, + parentCheckpointId: input.parentCheckpointId ?? null, + createdAt: input.createdAt ?? 1, + reason: input.reason ?? 'named', + files: input.files ?? [], + conversation: input.conversation ?? [], + artifacts: input.artifacts ?? [], + ...(input.label === undefined ? {} : { label: input.label }), + ...(input.sourceRunId === undefined + ? {} + : { sourceRunId: input.sourceRunId }), + } +} + +async function appendCheckpoint( + snapshots: SandboxSnapshots, + value: SandboxCheckpoint, +): Promise { + const writer = await snapshots.checkpoints.acquireWriter(value.threadId) + await snapshots.checkpoints.append({ + checkpoint: value, + expectedHeadId: value.parentCheckpointId, + writer, + }) + await writer.release() +} + +function withoutFork(store: SandboxCheckpointStore): SandboxCheckpointStore { + return { + get: store.get.bind(store), + list: store.list.bind(store), + getHead: store.getHead.bind(store), + append: store.append.bind(store), + deleteHead: store.deleteHead.bind(store), + acquireWriter: store.acquireWriter.bind(store), + listBlobReferences: store.listBlobReferences.bind(store), + } +} + +async function artifactKey(bytes: Uint8Array): Promise { + const digest = await crypto.subtle.digest('SHA-256', new Uint8Array(bytes)) + return `sandbox-artifacts/sha256/${Array.from(new Uint8Array(digest), (byte) => byte.toString(16).padStart(2, '0')).join('')}` +} + +async function artifactSnapshot(input: { + blobKey?: string + includeArtifact?: boolean + size?: number +}): Promise { + const snapshots = await memorySandboxSnapshots() + const artifacts = + input.includeArtifact === false + ? [] + : [ + { + artifactId: 'artifact', + name: 'file', + mimeType: 'text/plain', + createdAt: 1, + blobKey: input.blobKey ?? invalidArtifactKey, + size: input.size ?? 1, + }, + ] + await appendCheckpoint( + snapshots, + checkpoint({ id: 'checkpoint', threadId: THREAD, artifacts }), + ) + return snapshots +} + +function resolveArtifact( + snapshots: SandboxSnapshots, + changes: Partial[0]> = {}, +) { + return resolveSnapshotArtifact({ + threadId: THREAD, + checkpointId: 'checkpoint', + artifactId: 'artifact', + snapshots, + ...changes, + }) +} + +function workspaceHandle(input: { + files: ReadonlyArray<{ path: string; content: string }> + listCalls?: Array + root: string +}): SandboxHandle { + type Entry = { kind: 'dir' } | { kind: 'file'; bytes: Uint8Array } + const entries = new Map([[input.root, { kind: 'dir' }]]) + for (const file of input.files) { + const path = `${input.root}/${file.path}` + const parts = path.split('/') + for (let length = 2; length < parts.length; length++) { + const parent = parts.slice(0, length).join('/') + if (parent) entries.set(parent, { kind: 'dir' }) + } + entries.set(path, { + kind: 'file', + bytes: new TextEncoder().encode(file.content), + }) + } + const handle = makeFakeHandle('existing', 'fake') + handle.fs.lstat = async (path) => { + const entry = entries.get(path) + if (!entry) return undefined + return entry.kind === 'dir' + ? { type: 'dir', mode: 0o755 } + : { type: 'file', mode: 0o644, size: entry.bytes.byteLength } + } + handle.fs.list = async (path) => { + input.listCalls?.push(path) + const prefix = `${path}/` + return [...entries].flatMap(([entryPath, entry]) => { + const relative = entryPath.startsWith(prefix) + ? entryPath.slice(prefix.length) + : '' + return relative && !relative.includes('/') + ? [{ name: relative, path: entryPath, type: entry.kind }] + : [] + }) + } + handle.fs.read = async (path) => { + const entry = entries.get(path) + return entry?.kind === 'file' ? new TextDecoder().decode(entry.bytes) : '' + } + handle.fs.readBytes = async (path) => { + const entry = entries.get(path) + return entry?.kind === 'file' ? entry.bytes.slice() : new Uint8Array() + } + return handle +} + +function resumeWithHandle(fixture: NamedFixture, handle: SandboxHandle): void { + const resume = fixture.provider.resume.bind(fixture.provider) + fixture.provider.resume = async (input) => { + await resume(input) + return handle + } +} + +describe('public sandbox snapshot operations', () => { + describe('input staging', () => { + it('stages named-save dependencies exactly once before writer acquisition', async () => { + const fixture = await namedFixture() + let acquired = false + let lifecycleReads = 0 + let instanceGetReads = 0 + const acquireWriter = fixture.snapshots.checkpoints.acquireWriter.bind( + fixture.snapshots.checkpoints, + ) + vi.spyOn( + fixture.snapshots.checkpoints, + 'acquireWriter', + ).mockImplementation(async (threadId) => { + acquired = true + return acquireWriter(threadId) + }) + const instanceGet = fixture.instances.get + Object.defineProperty(fixture.instances, 'get', { + configurable: true, + get() { + instanceGetReads++ + if (acquired) throw new Error('late instance getter read') + return instanceGet + }, + }) + Object.defineProperty(fixture.definition, 'lifecycle', { + configurable: true, + get() { + lifecycleReads++ + if (acquired) throw new Error('late lifecycle read') + return { reuse: 'thread' as const } + }, + }) + + await expect(namedSave(fixture)).resolves.toMatchObject({ + reason: 'named', + }) + expect({ lifecycleReads, instanceGetReads }).toEqual({ + lifecycleReads: 1, + instanceGetReads: 1, + }) + }) + + it('does not acquire a named-save writer when first-read staging throws', async () => { + const fixture = await namedFixture() + const acquire = vi.spyOn(fixture.snapshots.checkpoints, 'acquireWriter') + Object.defineProperty(fixture.definition, 'lifecycle', { + configurable: true, + get() { + throw new Error('lifecycle staging failed') + }, + }) + + await expect(namedSave(fixture)).rejects.toThrow( + 'lifecycle staging failed', + ) + expect(acquire).not.toHaveBeenCalled() + }) + + it('does not access public resume getters after delayed writer acquisition', async () => { + const base = new LeaseProbeStore() + const acquisitionStarted = deferred() + const acquisitionGate = deferred() + const checkpoints: SandboxCheckpointStore = { + ...withoutFork(base), + acquireWriter: async (threadId) => { + acquisitionStarted.resolve() + await acquisitionGate.promise + return base.acquireWriter(threadId) + }, + } + const fixture = await namedFixture({ checkpoints }) + let ensureReads = 0 + Object.defineProperty(fixture.definition, 'ensureExisting', { + configurable: true, + get() { + ensureReads++ + throw new Error('public ensureExisting getter read') + }, + }) + + const save = namedSave(fixture) + await acquisitionStarted.promise + acquisitionGate.resolve() + await expect(save).resolves.toMatchObject({ reason: 'named' }) + expect(ensureReads).toBe(0) + }) + + it('stages a structural definition resume before writer acquisition', async () => { + const fixture = await namedFixture() + let acquired = false + let ensureReads = 0 + const acquireWriter = fixture.snapshots.checkpoints.acquireWriter.bind( + fixture.snapshots.checkpoints, + ) + vi.spyOn( + fixture.snapshots.checkpoints, + 'acquireWriter', + ).mockImplementation(async (threadId) => { + acquired = true + return acquireWriter(threadId) + }) + const baseDefinition = fixture.definition + const structuralDefinition: SandboxDefinition = { + id: baseDefinition.id, + provider: baseDefinition.provider, + workspace: baseDefinition.workspace, + policy: baseDefinition.policy, + lifecycle: baseDefinition.lifecycle, + hooks: baseDefinition.hooks, + fileEvents: baseDefinition.fileEvents, + key: baseDefinition.key, + ensure: baseDefinition.ensure, + get ensureExisting() { + ensureReads++ + if (acquired) throw new Error('late structural ensureExisting read') + return (context: SandboxEnsureContext) => + baseDefinition.ensureExisting(context) + }, + destroy: baseDefinition.destroy, + } + + await expect( + namedSave(fixture, { definition: structuralDefinition }), + ).resolves.toMatchObject({ reason: 'named' }) + expect(ensureReads).toBe(1) + }) + + it('uses one staged workspace value for hash, secrets, and custom root', async () => { + const fixture = await namedFixture() + const workspace = { + source: { type: 'none' as const }, + root: '/custom-workspace', + } + let workspaceReads = 0 + Object.defineProperty(fixture.definition, 'workspace', { + configurable: true, + get() { + workspaceReads++ + return workspaceReads === 1 ? workspace : undefined + }, + }) + await fixture.instances.upsert({ + key: computeSandboxKey({ + threadId: THREAD, + sandboxId: fixture.definition.id, + providerName: fixture.provider.name, + workspace, + }), + provider: fixture.provider.name, + providerSandboxId: 'existing-custom', + threadId: THREAD, + updatedAt: Date.now(), + }) + const listCalls: Array = [] + resumeWithHandle( + fixture, + workspaceHandle({ files: [], listCalls, root: workspace.root }), + ) + + await namedSave(fixture) + + expect(workspaceReads).toBe(1) + expect(listCalls).toEqual(['/custom-workspace']) + }) + + it('does not reread nested workspace or lifecycle getters after acquisition', async () => { + const secrets = createSecrets({ TOKEN: 'top-secret' }) + let acquired = false + let secretReads = 0 + let maxAgeReads = 0 + const workspace: WorkspaceDefinition = { source: { type: 'none' } } + Object.defineProperty(workspace, 'secrets', { + configurable: true, + enumerable: true, + get() { + secretReads++ + if (acquired) throw new Error('late secrets read') + return secrets + }, + }) + const lifecycle: NonNullable< + Parameters[0]['lifecycle'] + > = { reuse: 'thread' } + Object.defineProperty(lifecycle, 'snapshotMaxAge', { + configurable: true, + enumerable: true, + get() { + maxAgeReads++ + if (acquired) throw new Error('late max-age read') + return undefined + }, + }) + const fixture = await namedFixture({ lifecycle, workspace }) + secretReads = 0 + maxAgeReads = 0 + const acquireWriter = fixture.snapshots.checkpoints.acquireWriter.bind( + fixture.snapshots.checkpoints, + ) + vi.spyOn( + fixture.snapshots.checkpoints, + 'acquireWriter', + ).mockImplementation(async (threadId) => { + acquired = true + return acquireWriter(threadId) + }) + + await expect(namedSave(fixture)).resolves.toMatchObject({ + reason: 'named', + }) + expect({ secretReads, maxAgeReads }).toEqual({ + secretReads: 1, + maxAgeReads: 1, + }) + }) + + it('stages fork capability once before delayed writer acquisition', async () => { + const snapshots = await memorySandboxSnapshots() + await appendCheckpoint( + snapshots, + checkpoint({ id: 'source-checkpoint', threadId: 'source' }), + ) + const acquisitionStarted = deferred() + const acquisitionGate = deferred() + let forkReads = 0 + let acquisitionCalls = 0 + const checkpoints: SandboxCheckpointStore = { + ...withoutFork(snapshots.checkpoints), + acquireWriter: async (threadId) => { + acquisitionCalls++ + acquisitionStarted.resolve() + await acquisitionGate.promise + return snapshots.checkpoints.acquireWriter(threadId) + }, + get forkFromCheckpoint(): NonNullable< + SandboxCheckpointStore['forkFromCheckpoint'] + > { + forkReads++ + if (forkReads > 1) throw new Error('late fork getter read') + return snapshots.checkpoints.forkFromCheckpoint.bind( + snapshots.checkpoints, + ) + }, + } + + const fork = forkFromSandboxSnapshot({ + sourceThreadId: 'source', + sourceCheckpointId: 'source-checkpoint', + destinationThreadId: 'destination', + snapshots: { persistence: snapshots.persistence, checkpoints }, + }) + await acquisitionStarted.promise + acquisitionGate.resolve() + await expect(fork).resolves.toMatchObject({ reason: 'fork-root' }) + expect({ forkReads, acquisitionCalls }).toEqual({ + forkReads: 1, + acquisitionCalls: 1, + }) + }) + + it('does not acquire a fork writer when capability staging throws', async () => { + const snapshots = await memorySandboxSnapshots() + let acquisitionCalls = 0 + const checkpoints: SandboxCheckpointStore = { + ...withoutFork(snapshots.checkpoints), + acquireWriter: async (threadId) => { + acquisitionCalls++ + return snapshots.checkpoints.acquireWriter(threadId) + }, + get forkFromCheckpoint(): NonNullable< + SandboxCheckpointStore['forkFromCheckpoint'] + > { + throw new Error('fork staging failed') + }, + } + + await expect( + forkFromSandboxSnapshot({ + sourceThreadId: 'source', + sourceCheckpointId: 'checkpoint', + destinationThreadId: 'destination', + snapshots: { persistence: snapshots.persistence, checkpoints }, + }), + ).rejects.toThrow('fork staging failed') + expect(acquisitionCalls).toBe(0) + }) + + it('stages artifact inputs before a delayed checkpoint read', async () => { + const bytes = new TextEncoder().encode('actual') + const key = await artifactKey(bytes) + const snapshots = await artifactSnapshot({ + blobKey: key, + size: bytes.byteLength, + }) + await snapshots.persistence.stores.blobs.put(key, bytes) + const stored = await snapshots.checkpoints.get('checkpoint') + if (!stored) throw new Error('test checkpoint was not stored') + const checkpointStarted = deferred() + const checkpointGate = deferred() + vi.spyOn(snapshots.checkpoints, 'get').mockImplementation(async () => { + checkpointStarted.resolve() + return checkpointGate.promise + }) + let threadReads = 0 + let checkpointPending = false + const input: Parameters[0] = { + get threadId() { + threadReads++ + if (checkpointPending) throw new Error('late threadId read') + return THREAD + }, + checkpointId: 'checkpoint', + artifactId: 'artifact', + snapshots, + } + + const resolved = resolveSnapshotArtifact(input) + await checkpointStarted.promise + checkpointPending = true + checkpointGate.resolve(stored) + await expect(resolved).resolves.toMatchObject({ + artifact: { artifactId: 'artifact' }, + }) + expect(threadReads).toBe(1) + }) + + it('does not read a checkpoint when artifact dependency staging throws', async () => { + const snapshots = await memorySandboxSnapshots() + const getCheckpoint = vi.spyOn(snapshots.checkpoints, 'get') + Object.defineProperty(snapshots.persistence.stores.blobs, 'get', { + configurable: true, + get() { + throw new Error('blob getter staging failed') + }, + }) + + await expect(resolveArtifact(snapshots)).rejects.toThrow( + 'blob getter staging failed', + ) + expect(getCheckpoint).not.toHaveBeenCalled() + }) + }) + + describe('named-save lease operation', () => { + it('starts recursive non-overlapping renewal while provider resume is pending', async () => { + vi.useFakeTimers() + try { + const checkpoints = new LeaseProbeStore() + const renewalGate = deferred() + checkpoints.renewalGate = renewalGate.promise + const fixture = await namedFixture({ checkpoints }) + const resumeStarted = deferred() + const resumeGate = deferred() + const resume = fixture.provider.resume.bind(fixture.provider) + fixture.provider.resume = async (input) => { + resumeStarted.resolve() + await resumeGate.promise + return resume(input) + } + const save = namedSave(fixture) + await resumeStarted.promise + + await vi.advanceTimersByTimeAsync(10) + await checkpoints.renewalStarted.promise + await vi.advanceTimersByTimeAsync(50) + expect(checkpoints.renewals).toBe(1) + expect(checkpoints.maxActiveRenewals).toBe(1) + + renewalGate.resolve() + await checkpoints.renewalFinished.promise + await vi.advanceTimersByTimeAsync(10) + expect(checkpoints.renewals).toBe(2) + resumeGate.resolve() + await save + expect(checkpoints.releases).toBe(1) + } finally { + vi.useRealTimers() + } + }) + + it('awaits an in-flight renewal before one release', async () => { + vi.useFakeTimers() + try { + const checkpoints = new LeaseProbeStore() + const renewalGate = deferred() + checkpoints.renewalGate = renewalGate.promise + const fixture = await namedFixture({ checkpoints }) + const resumeStarted = deferred() + const resumeGate = deferred() + const resume = fixture.provider.resume.bind(fixture.provider) + fixture.provider.resume = async (input) => { + resumeStarted.resolve() + await resumeGate.promise + return resume(input) + } + const settled = vi.fn() + const save = namedSave(fixture) + void save.then(settled, settled) + await resumeStarted.promise + await vi.advanceTimersByTimeAsync(10) + await checkpoints.renewalStarted.promise + resumeGate.resolve() + await vi.advanceTimersByTimeAsync(0) + + expect(settled).not.toHaveBeenCalled() + expect(checkpoints.releases).toBe(0) + renewalGate.resolve() + await save + expect(checkpoints.releases).toBe(1) + } finally { + vi.useRealTimers() + } + }) + + it('keeps renewal loss primary, blocks append, and releases once', async () => { + vi.useFakeTimers() + try { + const checkpoints = new LeaseProbeStore() + const renewalGate = deferred() + checkpoints.renewalGate = renewalGate.promise + checkpoints.renewalError = new Error('renewal lost') + checkpoints.releaseError = new Error('release also failed') + const fixture = await namedFixture({ checkpoints }) + const resumeStarted = deferred() + const resumeGate = deferred() + const resume = fixture.provider.resume.bind(fixture.provider) + fixture.provider.resume = async (input) => { + resumeStarted.resolve() + await resumeGate.promise + return resume(input) + } + const save = namedSave(fixture) + await resumeStarted.promise + await vi.advanceTimersByTimeAsync(10) + await checkpoints.renewalStarted.promise + renewalGate.resolve() + resumeGate.resolve() + + await expect(save).rejects.toThrow('renewal lost') + expect(checkpoints.appends).toBe(0) + expect(checkpoints.releases).toBe(1) + } finally { + vi.useRealTimers() + } + }) + + it.each([ + { failure: 'capture', message: 'capture failed' }, + { failure: 'append', message: 'append failed' }, + ])( + 'keeps $failure error primary when release also fails', + async (testCase) => { + const checkpoints = new LeaseProbeStore() + checkpoints.releaseError = new Error('release failed') + const fixture = await namedFixture({ checkpoints }) + if (testCase.failure === 'capture') { + const resume = fixture.provider.resume.bind(fixture.provider) + fixture.provider.resume = async (input) => { + const handle = await resume(input) + if (handle) { + handle.fs.lstat = async () => { + throw new Error(testCase.message) + } + } + return handle + } + } else { + checkpoints.appendError = new Error(testCase.message) + } + + await expect(namedSave(fixture)).rejects.toThrow(testCase.message) + expect(checkpoints.releases).toBe(1) + expect(await checkpoints.getHead(THREAD)).toBeNull() + }, + ) + + it('reports release failure after successful publication', async () => { + const checkpoints = new LeaseProbeStore() + checkpoints.releaseError = new Error('release failed') + const fixture = await namedFixture({ checkpoints }) + + await expect(namedSave(fixture)).rejects.toThrow('release failed') + expect(checkpoints.releases).toBe(1) + expect(await checkpoints.getHead(THREAD)).not.toBeNull() + }) + + it('releases once after a successful named save', async () => { + const checkpoints = new LeaseProbeStore() + const fixture = await namedFixture({ checkpoints }) + + await expect(namedSave(fixture)).resolves.toMatchObject({ + reason: 'named', + }) + expect(checkpoints.releases).toBe(1) + }) + + it('rejects a stale compare-and-swap without moving the head', async () => { + const checkpoints = new LeaseProbeStore() + const fixture = await namedFixture({ checkpoints }) + await appendCheckpoint( + fixture.snapshots, + checkpoint({ id: 'existing-head', threadId: THREAD }), + ) + const getHead = vi + .spyOn(checkpoints, 'getHead') + .mockResolvedValueOnce(null) + + await expect(namedSave(fixture)).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_STALE_HEAD', + }) + getHead.mockRestore() + expect(await checkpoints.getHead(THREAD)).toBe('existing-head') + expect(checkpoints.releases).toBe(2) + }) + }) + + describe('named save behavior and policy', () => { + it('saves messages from an existing sandbox without provider create', async () => { + const fixture = await namedFixture() + await fixture.memory.persistence.stores.messages.saveThread(THREAD, [ + { role: 'user', content: 'saved' }, + ]) + + const saved = await namedSave(fixture, { label: 'before-change' }) + + expect(saved).toMatchObject({ + reason: 'named', + label: 'before-change', + sourceRunId: RUN, + conversation: [{ role: 'user', content: 'saved' }], + }) + expect(fixture.provider.calls).toMatchObject({ create: 0, resume: 1 }) + }) + + it('rejects a missing reusable sandbox without provider create', async () => { + const fixture = await namedFixture({ seedInstance: false }) + + await expect(namedSave(fixture)).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_MISSING_REUSABLE_SANDBOX', + }) + expect(fixture.provider.calls).toMatchObject({ create: 0, resume: 0 }) + }) + + it('rejects reuse none without provider work', async () => { + const fixture = await namedFixture({ lifecycle: { reuse: 'none' } }) + + await expect(namedSave(fixture)).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_REUSE_NONE', + }) + expect(fixture.provider.calls).toMatchObject({ create: 0, resume: 0 }) + }) + + it('rejects failed resume without provider create', async () => { + const fixture = await namedFixture() + fixture.provider.resume = async () => null + + await expect(namedSave(fixture)).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_MISSING_REUSABLE_SANDBOX', + }) + expect(fixture.provider.calls.create).toBe(0) + }) + + it('rejects an active writer before provider work', async () => { + const fixture = await namedFixture() + const writer = await fixture.snapshots.checkpoints.acquireWriter(THREAD) + + await expect(namedSave(fixture)).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_WRITER_CONFLICT', + }) + expect(fixture.provider.calls.resume).toBe(0) + await writer.release() + }) + + it.each([{ custom: false }, { custom: true }])( + 'protects the workspace marker with custom policy $custom', + async ({ custom }) => { + const workspace: WorkspaceDefinition = { + source: { type: 'none' }, + root: '/custom', + } + const hash = computeWorkspaceHash(workspace) + const marker = `.tanstack-projected-${hash}` + const listCalls: Array = [] + const fixture = await namedFixture({ + workspace, + ...(custom ? { policy: { exclude: () => false } } : {}), + }) + resumeWithHandle( + fixture, + workspaceHandle({ + root: workspace.root ?? '/workspace', + listCalls, + files: [ + { path: 'kept.txt', content: 'kept' }, + { path: `${marker}/private.txt`, content: 'private' }, + ], + }), + ) + + const saved = await namedSave(fixture) + + expect(saved.files.map((entry) => entry.path)).toEqual(['kept.txt']) + expect(listCalls).not.toContain(`/custom/${marker}`) + }, + ) + + it('preserves a supplied workspace hash when no workspace is defined', async () => { + const marker = '.tanstack-projected-caller-hash' + const listCalls: Array = [] + const fixture = await namedFixture({ + policy: { workspaceHash: 'caller-hash', exclude: () => false }, + }) + resumeWithHandle( + fixture, + workspaceHandle({ + root: '/workspace', + listCalls, + files: [ + { path: 'kept.txt', content: 'kept' }, + { path: `${marker}/private.txt`, content: 'private' }, + ], + }), + ) + + const saved = await namedSave(fixture) + + expect(saved.files.map((entry) => entry.path)).toEqual(['kept.txt']) + expect(listCalls).not.toContain(`/workspace/${marker}`) + }) + + it('preserves custom redaction and passes resolved workspace secrets', async () => { + const seenSecrets: Array = [] + const workspace: WorkspaceDefinition = { + source: { type: 'none' }, + secrets: createSecrets({ TOKEN: 'top-secret' }), + root: '/custom', + } + const fixture = await namedFixture({ + workspace, + policy: { + redact: ({ bytes, resolvedSecrets }) => { + seenSecrets.push(resolvedSecrets.TOKEN) + return new TextEncoder().encode( + new TextDecoder().decode(bytes).replace('visible', 'custom'), + ) + }, + }, + }) + resumeWithHandle( + fixture, + workspaceHandle({ + root: '/custom', + files: [{ path: 'secret.txt', content: 'visible top-secret' }], + }), + ) + + const saved = await namedSave(fixture) + const file = saved.files.find((entry) => entry.path === 'secret.txt') + if (!file || file.kind !== 'file') + throw new Error('captured file was not found') + const blob = await fixture.memory.persistence.stores.blobs.get( + file.blobKey, + ) + const text = await blob?.text() + + expect(seenSecrets).toEqual(['top-secret']) + expect(text).toContain('custom') + expect(text).not.toContain('top-secret') + }) + }) + + describe('ensureExisting', () => { + it('serializes resume through the supplied lock store', async () => { + const fixture = await namedFixture() + const locks = new InMemoryLockStore() + const firstResumeStarted = deferred() + const firstResumeGate = deferred() + let resumes = 0 + fixture.provider.resume = async (input) => { + resumes++ + if (resumes === 1) { + firstResumeStarted.resolve() + await firstResumeGate.promise + } + return makeFakeHandle(input.id, fixture.provider.name) + } + const context = { + threadId: THREAD, + runId: RUN, + store: fixture.instances, + locks, + } + + const first = fixture.definition.ensureExisting(context) + await firstResumeStarted.promise + const second = fixture.definition.ensureExisting(context) + await Promise.resolve() + await Promise.resolve() + expect(resumes).toBe(1) + + firstResumeGate.resolve() + await expect(Promise.all([first, second])).resolves.toHaveLength(2) + expect(resumes).toBe(2) + expect(fixture.provider.calls.create).toBe(0) + }) + + it('returns null for an expired record without resume or create', async () => { + vi.useFakeTimers() + vi.setSystemTime(100_000) + try { + const fixture = await namedFixture({ + lifecycle: { snapshotMaxAge: '1m' }, + seedInstance: false, + }) + await fixture.instances.upsert({ + key: fixture.definition.key({ threadId: THREAD, runId: 'old' }), + provider: fixture.provider.name, + providerSandboxId: 'existing', + threadId: THREAD, + updatedAt: 39_999, + }) + + await expect( + fixture.definition.ensureExisting({ + threadId: THREAD, + runId: RUN, + store: fixture.instances, + locks: new InMemoryLockStore(), + }), + ).resolves.toBeNull() + expect(fixture.provider.calls).toMatchObject({ create: 0, resume: 0 }) + } finally { + vi.useRealTimers() + } + }) + }) + + describe('fork', () => { + it('forks an older selected checkpoint without changing source state', async () => { + const snapshots = await memorySandboxSnapshots() + const selected = checkpoint({ + id: 'selected', + threadId: 'source', + conversation: [{ role: 'user', content: 'selected' }], + }) + const newest = checkpoint({ + id: 'newest', + threadId: 'source', + parentCheckpointId: selected.id, + createdAt: 2, + conversation: [{ role: 'user', content: 'newest' }], + }) + await appendCheckpoint(snapshots, selected) + await appendCheckpoint(snapshots, newest) + await snapshots.persistence.stores.messages.saveThread('source', [ + { role: 'user', content: 'newest' }, + ]) + const sourceBefore = await snapshots.checkpoints.list('source') + const sourceHeadBefore = await snapshots.checkpoints.getHead('source') + const sourceConversationBefore = + await snapshots.persistence.stores.messages.loadThread('source') + + const result = await forkFromSandboxSnapshot({ + sourceThreadId: 'source', + sourceCheckpointId: 'selected', + destinationThreadId: 'destination', + destinationCheckpointId: 'fork', + createdAt: 3, + snapshots, + }) + + expect(result).toMatchObject({ + id: 'fork', + threadId: 'destination', + parentCheckpointId: null, + reason: 'fork-root', + conversation: [{ role: 'user', content: 'selected' }], + }) + expect(await snapshots.checkpoints.list('source')).toEqual(sourceBefore) + expect(await snapshots.checkpoints.getHead('source')).toBe( + sourceHeadBefore, + ) + expect( + await snapshots.persistence.stores.messages.loadThread('source'), + ).toEqual(sourceConversationBefore) + expect( + await snapshots.persistence.stores.messages.loadThread('destination'), + ).toEqual([{ role: 'user', content: 'selected' }]) + }) + + it('releases once and does not renew a successful fork', async () => { + const snapshots = await memorySandboxSnapshots() + const checkpoints = new ForkProbeStore(snapshots.checkpoints) + const bundle: SandboxSnapshots = { + persistence: snapshots.persistence, + checkpoints, + } + await appendCheckpoint( + bundle, + checkpoint({ id: 'source-checkpoint', threadId: 'source' }), + ) + + await expect( + forkFromSandboxSnapshot({ + sourceThreadId: 'source', + sourceCheckpointId: 'source-checkpoint', + destinationThreadId: 'destination', + snapshots: bundle, + }), + ).resolves.toMatchObject({ reason: 'fork-root' }) + expect(checkpoints.releases).toBe(2) + expect(checkpoints.renewals).toBe(0) + }) + + it('releases an unavailable fork writer and keeps capability error primary', async () => { + const checkpoints = new LeaseProbeStore() + checkpoints.releaseError = new Error('release failed') + const snapshots = await memorySandboxSnapshots() + + await expect( + forkFromSandboxSnapshot({ + sourceThreadId: 'source', + sourceCheckpointId: 'source-checkpoint', + destinationThreadId: 'destination', + snapshots: { + persistence: snapshots.persistence, + checkpoints: withoutFork(checkpoints), + }, + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_FORK_UNAVAILABLE' }) + expect(checkpoints.releases).toBe(1) + }) + + it('keeps fork failure primary when release also fails', async () => { + const snapshots = await memorySandboxSnapshots() + const checkpoints = new ForkProbeStore(snapshots.checkpoints) + checkpoints.forkError = new Error('fork failed') + checkpoints.releaseError = new Error('release failed') + + await expect( + forkFromSandboxSnapshot({ + sourceThreadId: 'source', + sourceCheckpointId: 'source-checkpoint', + destinationThreadId: 'destination', + snapshots: { persistence: snapshots.persistence, checkpoints }, + }), + ).rejects.toThrow('fork failed') + expect(checkpoints.releases).toBe(1) + }) + + it('reports release failure after a successful fork publication', async () => { + const snapshots = await memorySandboxSnapshots() + const checkpoints = new ForkProbeStore(snapshots.checkpoints) + const bundle: SandboxSnapshots = { + persistence: snapshots.persistence, + checkpoints, + } + await appendCheckpoint( + bundle, + checkpoint({ id: 'source-checkpoint', threadId: 'source' }), + ) + checkpoints.releaseError = new Error('release failed') + + await expect( + forkFromSandboxSnapshot({ + sourceThreadId: 'source', + sourceCheckpointId: 'source-checkpoint', + destinationThreadId: 'destination', + snapshots: bundle, + }), + ).rejects.toThrow('release failed') + expect(await checkpoints.getHead('destination')).not.toBeNull() + expect(checkpoints.releases).toBe(2) + }) + }) + + describe('artifact resolution', () => { + it('rejects a foreign-thread checkpoint', async () => { + const snapshots = await artifactSnapshot({ includeArtifact: false }) + + await expect( + resolveArtifact(snapshots, { threadId: 'other' }), + ).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_FOREIGN_CHECKPOINT_ARTIFACT', + }) + }) + + it('rejects a missing checkpoint', async () => { + const snapshots = await memorySandboxSnapshots() + + await expect( + resolveArtifact(snapshots, { checkpointId: 'missing' }), + ).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_MISSING_CHECKPOINT_ARTIFACT', + }) + }) + + it('rejects an existing checkpoint with no selected artifact', async () => { + const snapshots = await artifactSnapshot({ includeArtifact: false }) + + await expect(resolveArtifact(snapshots)).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_MISSING_CHECKPOINT_ARTIFACT', + }) + }) + + it('rejects a missing artifact blob', async () => { + const snapshots = await artifactSnapshot({}) + + await expect(resolveArtifact(snapshots)).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_INVALID_ARTIFACT_BYTES', + }) + }) + + it('rejects artifact bytes with an invalid digest', async () => { + const bytes = new TextEncoder().encode('actual') + const snapshots = await artifactSnapshot({ + blobKey: invalidArtifactKey, + size: bytes.byteLength, + }) + await snapshots.persistence.stores.blobs.put(invalidArtifactKey, bytes) + + await expect(resolveArtifact(snapshots)).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_INVALID_ARTIFACT_BYTES', + }) + }) + + it('rejects artifact bytes with an invalid size', async () => { + const bytes = new TextEncoder().encode('actual') + const key = await artifactKey(bytes) + const snapshots = await artifactSnapshot({ + blobKey: key, + size: bytes.byteLength + 1, + }) + await snapshots.persistence.stores.blobs.put(key, bytes) + + await expect(resolveArtifact(snapshots)).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_INVALID_ARTIFACT_BYTES', + }) + }) + + it('returns metadata and bytes independent from store state', async () => { + const bytes = new TextEncoder().encode('actual') + const key = await artifactKey(bytes) + const snapshots = await artifactSnapshot({ + blobKey: key, + size: bytes.byteLength, + }) + await snapshots.persistence.stores.blobs.put(key, bytes) + + const first = await resolveArtifact(snapshots) + first.bytes[0] = 0 + first.artifact.name = 'changed' + const second = await resolveArtifact(snapshots) + + expect(new TextDecoder().decode(second.bytes)).toBe('actual') + expect(second.artifact.name).toBe('file') + }) + + it('uses the root-exported snapshot error class', async () => { + const snapshots = await memorySandboxSnapshots() + + await expect( + resolveArtifact(snapshots, { checkpointId: 'missing' }), + ).rejects.toBeInstanceOf(SandboxSnapshotError) + }) + }) +}) diff --git a/packages/ai-sandbox/tests/snapshot-policy-export.test-d.ts b/packages/ai-sandbox/tests/snapshot-policy-export.test-d.ts new file mode 100644 index 0000000000..d2086b6cfa --- /dev/null +++ b/packages/ai-sandbox/tests/snapshot-policy-export.test-d.ts @@ -0,0 +1,28 @@ +import { expectTypeOf } from 'vitest' +import type { + SandboxCheckpointStoreOptions, + SandboxSnapshotPolicy, +} from '../src' + +const checkpointStoreOptions: SandboxCheckpointStoreOptions = { + leaseDurationMs: 120_000, +} + +expectTypeOf( + checkpointStoreOptions, +).toMatchTypeOf() + +const policy: SandboxSnapshotPolicy = { + include: (path, kind) => path !== 'tmp' && kind === 'file', + exclude: (path) => path.startsWith('.git'), + redact: ({ path, bytes, resolvedSecrets }) => { + expectTypeOf(path).toBeString() + expectTypeOf(bytes).toEqualTypeOf() + expectTypeOf(resolvedSecrets).toEqualTypeOf< + Readonly> + >() + return bytes + }, +} + +expectTypeOf(policy).toMatchTypeOf() diff --git a/packages/ai-sandbox/tests/snapshots.test.ts b/packages/ai-sandbox/tests/snapshots.test.ts new file mode 100644 index 0000000000..173f779dfc --- /dev/null +++ b/packages/ai-sandbox/tests/snapshots.test.ts @@ -0,0 +1,2037 @@ +import { describe, expect, it } from 'vitest' +import { readFileSync } from 'node:fs' +import type { SandboxHandle, SandboxFsStat } from '../src/contracts' +import type { SandboxSnapshotEntry } from '../src/checkpoint-store' +import type { ArtifactRecord, ArtifactStore } from '@tanstack/ai-persistence' +import { + captureSandboxFiles, + captureSandboxArtifacts, + defaultSandboxSnapshotPolicy, + restoreSandboxFiles, + SandboxSnapshotError, +} from '../src/snapshots' + +type Entry = { + type: 'file' | 'dir' | 'symlink' | 'other' + mode: number + bytes?: Uint8Array +} + +function artifactStore( + listForThread: ArtifactStore['listForThread'], +): ArtifactStore { + return { + listForThread, + save: async () => {}, + get: async () => null, + list: async () => [], + delete: async () => {}, + deleteForRun: async () => {}, + } +} + +function artifactRecord( + overrides: Partial = {}, +): ArtifactRecord { + return { + artifactId: 'a', + runId: 'r', + threadId: 't', + name: 'a', + mimeType: 'x', + size: 1, + createdAt: 1, + ...overrides, + } +} + +function fakeHandle(entries: Record): SandboxHandle { + const normalize = (path: string) => + path.replace(/\\/g, '/').replace(/\/+/g, '/') + const parentDirs = (path: string) => { + const parts = normalize(path).split('/').filter(Boolean) + for (let i = 1; i < parts.length; i++) + entries[`/${parts.slice(0, i).join('/')}`] ??= { + type: 'dir', + mode: 0o755, + } + } + for (const path of Object.keys(entries)) parentDirs(path) + const fs = { + async read(path: string) { + return new TextDecoder().decode(await this.readBytes(path)) + }, + async readBytes(path: string) { + const e = entries[normalize(path)] + if (!e?.bytes) throw new Error(`missing ${path}`) + return e.bytes + }, + async write(path: string, data: string | Uint8Array) { + const bytes = + typeof data === 'string' ? new TextEncoder().encode(data) : data + parentDirs(path) + entries[normalize(path)] = { + type: 'file', + mode: 0o644, + bytes: bytes.slice(), + } + }, + async list(path: string) { + const root = normalize(path).replace(/\/$/, '') + const seen = new Map() + for (const [key, value] of Object.entries(entries)) { + if (!key.startsWith(`${root}/`)) continue + const rest = key.slice(root.length + 1) + if (!rest || rest.includes('/')) { + const name = rest.split('/')[0]! + seen.set( + name, + rest.includes('/') ? 'dir' : value.type === 'dir' ? 'dir' : 'file', + ) + } else seen.set(rest, value.type === 'dir' ? 'dir' : 'file') + } + return Array.from(seen, ([name, type]) => ({ + name, + path: `${root}/${name}`, + type, + })) + }, + async mkdir(path: string) { + entries[normalize(path)] = { type: 'dir', mode: 0o755 } + parentDirs(path) + }, + async remove(path: string) { + const root = normalize(path) + for (const key of Object.keys(entries)) + if (key === root || key.startsWith(`${root}/`)) delete entries[key] + }, + async rename() {}, + async exists(path: string) { + return Boolean(entries[normalize(path)]) + }, + async lstat(path: string): Promise { + const e = entries[normalize(path)] + if (!e && normalize(path) === '/workspace') + return { type: 'dir', mode: 0o755 } + if (!e) return undefined + return e.type === 'file' + ? { type: 'file', mode: e.mode, size: e.bytes?.byteLength ?? 0 } + : e.type === 'dir' + ? { type: 'dir', mode: e.mode } + : { type: e.type, mode: e.mode } + }, + } + return { + id: 'fake', + provider: 'fake', + capabilities: { + fs: true, + exec: true, + ports: false, + snapshots: false, + fork: false, + env: false, + backgroundProcesses: false, + writableStdin: false, + killableProcesses: false, + networkPolicy: false, + durableFilesystem: false, + }, + fs, + git: { + clone: async () => {}, + status: async () => '', + add: async () => {}, + commit: async () => {}, + push: async () => {}, + pull: async () => {}, + branch: async () => '', + }, + process: { + exec: async () => ({ stdout: '', stderr: '', exitCode: 0 }), + spawn: async () => { + throw new Error('unsupported') + }, + }, + ports: { connect: async () => ({ url: '' }) }, + env: { set: async () => {} }, + destroy: async () => {}, + } +} + +function blobs() { + const values = new Map() + const deletes: string[] = [] + let putCount = 0 + return { + values, + deletes, + get putCount() { + return putCount + }, + async put(key: string, body: Uint8Array) { + putCount++ + values.set(key, body) + return { key, size: body.byteLength } + }, + async get(key: string) { + const value = values.get(key) + return value + ? { + key, + size: value.byteLength, + arrayBuffer: async () => value.slice().buffer, + text: async () => new TextDecoder().decode(value), + } + : null + }, + async head(key: string) { + const value = values.get(key) + return value ? { key, size: value.byteLength } : null + }, + async delete(key: string) { + deletes.push(key) + values.delete(key) + }, + } +} + +function blobsWithAccessLog() { + const store = blobs() + const gets: string[] = [] + const heads: string[] = [] + const get = store.get + const head = store.head + store.get = async (key) => { + gets.push(key) + return get(key) + } + store.head = async (key) => { + heads.push(key) + return head(key) + } + return { store, gets, heads } +} + +function fsMutationLog(handle: SandboxHandle) { + const writes: string[] = [] + const removes: string[] = [] + const mkdirs: string[] = [] + const write = handle.fs.write + const remove = handle.fs.remove + const mkdir = handle.fs.mkdir + handle.fs.write = async (path, data) => { + writes.push(path) + await write(path, data) + } + handle.fs.remove = async (path) => { + removes.push(path) + await remove(path) + } + handle.fs.mkdir = async (path) => { + mkdirs.push(path) + await mkdir(path) + } + return { writes, removes, mkdirs } +} + +async function putSnapshotBlob( + store: ReturnType, + bytes: Uint8Array, +) { + const digest = await crypto.subtle.digest('SHA-256', bytes.slice()) + const hash = Array.from(new Uint8Array(digest), (byte) => + byte.toString(16).padStart(2, '0'), + ).join('') + const key = `sandbox-files/sha256/${hash}` + await store.put(key, bytes) + return key +} + +describe('portable sandbox snapshots', () => { + it('fails loudly when artifact persistence support is absent', async () => { + await expect( + captureSandboxArtifacts({ blobs: blobs() }, 'thread-1'), + ).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_ARTIFACT_SUPPORT_REQUIRED', + }) + }) + + it('captures thread artifacts through the real listForThread resolver and reuses source reads', async () => { + const store = blobs() + const sourceKey = 'custom/source' + await store.put(sourceKey, new Uint8Array([1, 2, 3])) + let listed = '' + let reads = 0 + const originalGet = store.get + store.get = async (key) => { + if (key === sourceKey) reads++ + return originalGet(key) + } + const artifacts = { + listForThread: async (threadId: string) => { + listed = threadId + return [ + { + artifactId: 'b', + runId: 'run', + threadId, + name: 'b', + mimeType: 'x', + size: 3, + createdAt: 2, + blobKey: sourceKey, + }, + { + artifactId: 'a', + runId: 'run', + threadId, + name: 'a', + mimeType: 'x', + size: 3, + createdAt: 1, + blobKey: sourceKey, + }, + ] + }, + } + const captured = await captureSandboxArtifacts( + { blobs: store, artifacts: artifactStore(artifacts.listForThread) }, + 'thread-9', + ) + expect(listed).toBe('thread-9') + expect(reads).toBe(1) + expect(captured.map((artifact) => artifact.artifactId)).toEqual(['a', 'b']) + expect(captured[0]?.blobKey).toBe(captured[1]?.blobKey) + }) + + it('orders tied artifact timestamps by UTF-8 bytes', async () => { + const store = blobs() + const sourceKey = 'source/blob' + await store.put(sourceKey, new Uint8Array([1])) + const artifacts = artifactStore(async () => [ + artifactRecord({ + artifactId: '\u{10000}', + name: 'high', + blobKey: sourceKey, + }), + artifactRecord({ + artifactId: '\uE000', + name: 'private', + blobKey: sourceKey, + }), + ]) + const captured = await captureSandboxArtifacts( + { blobs: store, artifacts }, + 't', + ) + expect(captured.map((artifact) => artifact.artifactId)).toEqual([ + '\uE000', + '\u{10000}', + ]) + }) + + it('verifies every source before touching destination artifact blobs', async () => { + const store = blobs() + const sourceKey = 'source/valid' + await store.put(sourceKey, new Uint8Array([1])) + const heads: string[] = [] + const puts: string[] = [] + const originalHead = store.head + const originalPut = store.put + store.head = async (key) => { + heads.push(key) + return originalHead(key) + } + store.put = async (key, body) => { + puts.push(key) + return originalPut(key, body) + } + const artifacts = artifactStore(async () => [ + artifactRecord({ + artifactId: 'valid', + name: 'valid', + blobKey: sourceKey, + }), + artifactRecord({ + artifactId: 'missing', + name: 'missing', + createdAt: 2, + blobKey: 'source/missing', + }), + ]) + await expect( + captureSandboxArtifacts({ blobs: store, artifacts }, 't'), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_MISSING_ARTIFACT_BLOB' }) + expect(heads).toEqual([]) + expect(puts).toEqual([]) + }) + + it('uses artifact resolver fallback and never writes a partial result on missing source', async () => { + const store = blobs() + const artifacts = { + listForThread: async () => [ + { + artifactId: 'missing', + runId: 'run', + threadId: 't', + name: 'x', + mimeType: 'x', + size: 1, + createdAt: 1, + }, + ], + } + await expect( + captureSandboxArtifacts( + { blobs: store, artifacts: artifactStore(artifacts.listForThread) }, + 't', + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_MISSING_ARTIFACT_BLOB' }) + expect(store.putCount).toBe(0) + }) + + it('uses head for an existing immutable artifact blob and never deletes source blobs', async () => { + const store = blobs() + const sourceKey = 'source/blob' + const bytes = new Uint8Array([1, 2, 3]) + await store.put(sourceKey, bytes) + let puts = 0 + const originalPut = store.put + store.put = async (key, body) => { + puts++ + return originalPut(key, body) + } + const artifacts = artifactStore(async () => [ + artifactRecord({ blobKey: sourceKey, size: 3 }), + ]) + await captureSandboxArtifacts({ blobs: store, artifacts }, 't') + expect(puts).toBe(1) + await captureSandboxArtifacts({ blobs: store, artifacts }, 't') + expect(puts).toBe(1) + expect(store.deletes).toEqual([]) + }) + + it('does not put an artifact when its source blob is missing', async () => { + const store = blobs() + const artifacts = artifactStore(async () => [ + artifactRecord({ artifactId: 'missing', name: 'x' }), + ]) + await expect( + captureSandboxArtifacts({ blobs: store, artifacts }, 't'), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_MISSING_ARTIFACT_BLOB' }) + expect(store.putCount).toBe(0) + }) + + it('orders Unicode file metadata by UTF-8 bytes', async () => { + const snapshot = await captureSandboxFiles( + fakeHandle({ + '/workspace/é': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + '/workspace/e\u0301': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }), + { blobs: blobs() }, + ) + expect( + snapshot.files + .filter((entry) => entry.kind === 'file') + .map((entry) => entry.path), + ).toEqual(['e\u0301', 'é']) + }) + it('uses a non-default workspace root for capture and restore', async () => { + const paths: string[] = [] + const source = fakeHandle({ + '/custom/root/a': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }) + const sourceLstat = source.fs.lstat! + source.fs.lstat = async (path) => { + paths.push(path) + return sourceLstat(path) + } + const bundle = { blobs: blobs(), workspaceRoot: '/custom/root' } + const snapshot = await captureSandboxFiles(source, bundle) + const target = fakeHandle({ '/custom/root': { type: 'dir', mode: 0o755 } }) + const targetWrite = target.fs.write + target.fs.write = async (path, bytes) => { + paths.push(path) + return targetWrite(path, bytes) + } + await restoreSandboxFiles(target, bundle, snapshot) + expect(paths.every((path) => !path.startsWith('/workspace'))).toBe(true) + expect(await target.fs.exists('/custom/root/a')).toBe(true) + }) + it('rejects unsafe manifests before mutation', async () => { + const target = fakeHandle({ + '/workspace/keep': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }) + await expect( + restoreSandboxFiles( + target, + { blobs: blobs() }, + { files: [{ path: '../escape', kind: 'dir' }] }, + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_PATH' }) + expect(await target.fs.exists('/workspace/keep')).toBe(true) + }) + + it('rejects duplicate paths after validating each path', async () => { + await expect( + restoreSandboxFiles( + fakeHandle({}), + { blobs: blobs() }, + { + files: [ + { path: 'a', kind: 'dir' }, + { path: 'a', kind: 'dir' }, + ], + }, + ), + ).rejects.toThrow('Duplicate path') + }) + + it('rejects file ancestors after validating each path', async () => { + await expect( + restoreSandboxFiles( + fakeHandle({}), + { blobs: blobs() }, + { + files: [ + { + path: 'a', + kind: 'file', + blobKey: `sandbox-files/sha256/${'0'.repeat(64)}`, + size: 1, + }, + { path: 'a/b', kind: 'dir' }, + ], + }, + ), + ).rejects.toThrow('File ancestor') + }) + + it('requires lstat for capture', async () => { + const handle = fakeHandle({}) + handle.fs.lstat = undefined + await expect( + captureSandboxFiles(handle, { blobs: blobs() }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_LSTAT_REQUIRED' }) + }) + + it('requires lstat for restore before mutation', async () => { + const target = fakeHandle({ + '/workspace/keep': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }) + target.fs.lstat = undefined + await expect( + restoreSandboxFiles(target, { blobs: blobs() }, { files: [] }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_LSTAT_REQUIRED' }) + expect(await target.fs.exists('/workspace/keep')).toBe(true) + }) + + it('rejects generated projection marker paths before mutation', async () => { + const target = fakeHandle({ + '/workspace/keep': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }) + await expect( + restoreSandboxFiles( + target, + { blobs: blobs() }, + { + files: [{ path: '.tanstack-projected-abc123', kind: 'dir' }], + }, + defaultSandboxSnapshotPolicy('abc123'), + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_PATH' }) + expect(await target.fs.exists('/workspace/keep')).toBe(true) + }) + + it('reuses an existing content-addressed blob and returns known sha256', async () => { + const store = blobs() + const bytes = new TextEncoder().encode('hello') + const first = await captureSandboxFiles( + fakeHandle({ '/workspace/a': { type: 'file', mode: 0o644, bytes } }), + { blobs: store }, + ) + const second = await captureSandboxFiles( + fakeHandle({ '/workspace/b': { type: 'file', mode: 0o644, bytes } }), + { blobs: store }, + ) + expect(first.files[0]).toMatchObject({ + blobKey: + 'sandbox-files/sha256/2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824', + }) + expect(second.files[0]).toMatchObject({ + blobKey: first.files[0]?.kind === 'file' ? first.files[0].blobKey : '', + }) + expect(store.putCount).toBe(1) + }) + + it('does not complete capture when blob put fails', async () => { + const store = blobs() + store.put = async () => { + throw new Error('put failed') + } + await expect( + captureSandboxFiles( + fakeHandle({ + '/workspace/a': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }), + { blobs: store }, + ), + ).rejects.toThrow('put failed') + }) + it('captures binary files and empty directories with content-addressed blobs', async () => { + const bundle = { blobs: blobs() } + const snapshot = await captureSandboxFiles( + fakeHandle({ + '/workspace/image.bin': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([0, 255, 1]), + }, + '/workspace/empty': { type: 'dir', mode: 0o755 }, + }), + bundle, + defaultSandboxSnapshotPolicy(), + ) + expect(snapshot.files).toContainEqual({ path: 'empty', kind: 'dir' }) + const image = snapshot.files.find((x) => x.path === 'image.bin') + expect(image?.kind).toBe('file') + expect(image && image.kind === 'file' ? image.blobKey : '').toMatch( + /^sandbox-files\/sha256\/[0-9a-f]{64}$/, + ) + }) + + it('omits non-empty directories from the manifest', async () => { + const snapshot = await captureSandboxFiles( + fakeHandle({ + '/workspace/nested/file.txt': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }), + { blobs: blobs() }, + ) + expect(snapshot.files).not.toContainEqual({ path: 'nested', kind: 'dir' }) + }) + + it('redacts resolved secrets before hashing and storing', async () => { + const bundle = { blobs: blobs() } + const policy = { + ...defaultSandboxSnapshotPolicy(), + redact: ({ + bytes, + resolvedSecrets, + }: { + path: string + bytes: Uint8Array + resolvedSecrets: Record + }) => { + let text = new TextDecoder().decode(bytes) + for (const [name, secret] of Object.entries(resolvedSecrets)) + text = text.replaceAll(secret, `[${name}]`) + return new TextEncoder().encode(text) + }, + } + const snapshot = await captureSandboxFiles( + fakeHandle({ + '/workspace/app.ts': { + type: 'file', + mode: 0o644, + bytes: new TextEncoder().encode('key=secret-value'), + }, + }), + bundle, + policy, + { API_KEY: 'secret-value' }, + ) + const entry = snapshot.files.find((x) => x.path === 'app.ts') + expect(entry?.kind).toBe('file') + if (!entry || entry.kind !== 'file') throw new Error('expected file') + expect( + await bundle.blobs + .get(entry.blobKey) + .then(async (x) => + x + ? new TextDecoder().decode(new Uint8Array(await x.arrayBuffer())) + : '', + ), + ).not.toContain('secret-value') + }) + + it('redacts secrets deterministically and keeps the longest overlapping secret hidden', async () => { + const bytes = new TextEncoder().encode('prefix-secret-long-suffix') + const first = await captureSandboxFiles( + fakeHandle({ '/workspace/a': { type: 'file', mode: 0o644, bytes } }), + { blobs: blobs() }, + defaultSandboxSnapshotPolicy(), + { short: 'secret', long: 'secret-long' }, + ) + const second = await captureSandboxFiles( + fakeHandle({ '/workspace/a': { type: 'file', mode: 0o644, bytes } }), + { blobs: blobs() }, + defaultSandboxSnapshotPolicy(), + { long: 'secret-long', short: 'secret' }, + ) + expect(first.files[0]).toEqual(second.files[0]) + const store = blobs() + await captureSandboxFiles( + fakeHandle({ '/workspace/a': { type: 'file', mode: 0o644, bytes } }), + { blobs: store }, + defaultSandboxSnapshotPolicy(), + { short: 'secret', long: 'secret-long' }, + ) + const value = [...store.values.values()][0] + expect(value && new TextDecoder().decode(value)).not.toContain('secret') + }) + + it('applies built-in secret redaction after policy redaction', async () => { + const store = blobs() + const snapshot = await captureSandboxFiles( + fakeHandle({ + '/workspace/a': { + type: 'file', + mode: 0o644, + bytes: new TextEncoder().encode('secret'), + }, + }), + { blobs: store }, + { redact: ({ bytes }) => bytes }, + { key: 'secret' }, + ) + const entry = snapshot.files[0] + if (!entry || entry.kind !== 'file') throw new Error('expected file') + const value = await store.get(entry.blobKey) + expect(value && (await value.text())).not.toContain('secret') + }) + + it('stores fixed zero bytes for secrets, not a secret fingerprint', async () => { + const store = blobs() + const snapshot = await captureSandboxFiles( + fakeHandle({ + '/workspace/a': { + type: 'file', + mode: 0o644, + bytes: new TextEncoder().encode('x=secret'), + }, + }), + { blobs: store }, + defaultSandboxSnapshotPolicy(), + { key: 'secret' }, + ) + const entry = snapshot.files.find((file) => file.path === 'a') + if (!entry || entry.kind !== 'file') throw new Error('expected file') + const value = store.values.get(entry.blobKey) + expect(value).toEqual(new Uint8Array([120, 61, 0, 0, 0, 0, 0, 0])) + }) + + it('zeros the union of overlapping file secret matches', async () => { + const store = blobs() + const snapshot = await captureSandboxFiles( + fakeHandle({ + '/workspace/a': { + type: 'file', + mode: 0o644, + bytes: new TextEncoder().encode('abcde'), + }, + }), + { blobs: store }, + defaultSandboxSnapshotPolicy(), + { first: 'abcd', second: 'bcde' }, + ) + const entry = snapshot.files.find((file) => file.path === 'a') + if (!entry || entry.kind !== 'file') throw new Error('expected file') + expect(store.values.get(entry.blobKey)).toEqual(new Uint8Array(5)) + }) + + it('redacts artifact bytes with the final configured secrets', async () => { + const store = blobs() + await store.put('source', new TextEncoder().encode('artifact-secret')) + const captured = await captureSandboxArtifacts( + { + blobs: store, + artifacts: artifactStore(async () => [ + artifactRecord({ + mimeType: 'text/plain', + size: 15, + blobKey: 'source', + }), + ]), + }, + 't', + { key: 'secret' }, + ) + const value = store.values.get(captured[0]!.blobKey) + expect(value).toEqual(new TextEncoder().encode('artifact-\0\0\0\0\0\0')) + }) + + it('zeros the union of overlapping artifact secret matches', async () => { + const store = blobs() + await store.put('source', new TextEncoder().encode('abcde')) + const captured = await captureSandboxArtifacts( + { + blobs: store, + artifacts: artifactStore(async () => [ + artifactRecord({ + mimeType: 'application/octet-stream', + size: 5, + blobKey: 'source', + }), + ]), + }, + 't', + { first: 'abcd', second: 'bcde' }, + ) + expect(store.values.get(captured[0]!.blobKey)).toEqual(new Uint8Array(5)) + }) + + it('traverses a parent directory when only nested TypeScript files are included', async () => { + const snapshot = await captureSandboxFiles( + fakeHandle({ + '/workspace/src/app.ts': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + '/workspace/src/app.js': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([2]), + }, + }), + { blobs: blobs() }, + { include: (path, kind) => kind === 'dir' || path.endsWith('.ts') }, + ) + expect(snapshot.files.map((file) => file.path)).toEqual(['src/app.ts']) + }) + + it('performs one destination head/put per equal artifact digest', async () => { + const store = blobs() + await store.put('source-a', new Uint8Array([1, 2])) + await store.put('source-b', new Uint8Array([1, 2])) + let heads = 0 + let puts = 0 + const head = store.head + const put = store.put + store.head = async (key) => { + heads++ + return head(key) + } + store.put = async (key, bytes) => { + puts++ + return put(key, bytes) + } + await captureSandboxArtifacts( + { + blobs: store, + artifacts: artifactStore(async () => [ + artifactRecord({ size: 2, blobKey: 'source-a' }), + artifactRecord({ + artifactId: 'b', + name: 'b', + size: 2, + createdAt: 2, + blobKey: 'source-b', + }), + ]), + }, + 't', + ) + expect(heads).toBe(1) + expect(puts).toBe(1) + }) + + it('performs one destination head and put for equal file bytes', async () => { + const store = blobs() + let heads = 0 + let puts = 0 + const head = store.head + const put = store.put + store.head = async (key) => { + heads++ + return head(key) + } + store.put = async (key, bytes) => { + puts++ + return put(key, bytes) + } + await captureSandboxFiles( + fakeHandle({ + '/workspace/a': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([7]), + }, + '/workspace/b': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([7]), + }, + }), + { blobs: store }, + ) + expect(heads).toBe(1) + expect(puts).toBe(1) + }) + + it('hashes payloads without calling Uint8Array.prototype.slice', async () => { + const store = blobs() + const originalSlice = Uint8Array.prototype.slice + let slices = 0 + Uint8Array.prototype.slice = function (...args) { + slices++ + return originalSlice.apply(this, args) + } + try { + await captureSandboxFiles( + fakeHandle({ + '/workspace/a': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }), + { blobs: store }, + ) + const source = new Uint8Array([2]) + await store.put('source', source) + await captureSandboxArtifacts( + { + blobs: store, + artifacts: artifactStore(async () => [ + artifactRecord({ blobKey: 'source' }), + ]), + }, + 't', + ) + } finally { + Uint8Array.prototype.slice = originalSlice + } + // One file redaction, one artifact source read, and one artifact redaction + // each use slice; hashing must not add another payload copy. + expect(slices).toBe(3) + }) + + it('keeps persistence optional and type-only in snapshot source', () => { + const packageJson = JSON.parse( + readFileSync( + new URL('../../ai-sandbox/package.json', import.meta.url), + 'utf8', + ), + ) + if ( + typeof packageJson !== 'object' || + packageJson === null || + !('peerDependenciesMeta' in packageJson) + ) + throw new Error('invalid package.json') + expect( + packageJson.peerDependenciesMeta?.['@tanstack/ai-persistence']?.optional, + ).toBe(true) + const source = readFileSync( + new URL('../src/snapshots.ts', import.meta.url), + 'utf8', + ) + expect(source).not.toMatch( + /import(?!\s+type)[^;]*['"]@tanstack\/ai-persistence['"]/, + ) + }) + + it.each(['.git/config', 'node_modules/x', '.env.local'])( + 'excludes %s', + async (path) => { + const snapshot = await captureSandboxFiles( + fakeHandle({ + [`/workspace/${path}`]: { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }), + { blobs: blobs() }, + defaultSandboxSnapshotPolicy(), + ) + expect(snapshot.files).not.toContainEqual( + expect.objectContaining({ path }), + ) + }, + ) + + it.each([ + 'src/.git/config', + 'src/node_modules/pkg/index.js', + 'src/.env.local', + ])( + 'excludes protected segments at any depth during capture: %s', + async (path) => { + const snapshot = await captureSandboxFiles( + fakeHandle({ + [`/workspace/${path}`]: { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }), + { blobs: blobs() }, + defaultSandboxSnapshotPolicy(), + ) + expect(snapshot.files).not.toContainEqual( + expect.objectContaining({ path }), + ) + }, + ) + + it('captures projection-looking user files without a workspace hash', async () => { + const snapshot = await captureSandboxFiles( + fakeHandle({ + '/workspace/.tanstack-projected-other/file': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }), + { blobs: blobs() }, + defaultSandboxSnapshotPolicy(), + ) + expect(snapshot.files).toContainEqual( + expect.objectContaining({ path: '.tanstack-projected-other/file' }), + ) + }) + + it('does not inspect an excluded symlink during capture', async () => { + const handle = fakeHandle({ + '/workspace/.git': { type: 'symlink', mode: 0o777 }, + }) + const inspected: string[] = [] + const originalLstat = handle.fs.lstat! + handle.fs.lstat = async (path) => { + inspected.push(path) + return originalLstat(path) + } + + await expect( + captureSandboxFiles(handle, { blobs: blobs() }), + ).resolves.toEqual({ files: [] }) + + expect(inspected.every((path) => path === '/workspace')).toBe(true) + }) + + it('excludes bootstrap-owned instruction symlinks during capture', async () => { + const snapshot = await captureSandboxFiles( + fakeHandle({ + '/workspace/AGENTS.md': { + type: 'file', + mode: 0o644, + bytes: new TextEncoder().encode('# Instructions'), + }, + '/workspace/CLAUDE.md': { type: 'symlink', mode: 0o777 }, + '/workspace/GEMINI.md': { type: 'symlink', mode: 0o777 }, + }), + { blobs: blobs() }, + ) + + expect(snapshot.files.map((entry) => entry.path)).toEqual(['AGENTS.md']) + }) + + it.each([ + '.claude/skills/review', + '.codex/skills/review', + '.grok/skills/review', + ])( + 'excludes a projected git-skill symlink during capture: %s', + async (path) => { + const snapshot = await captureSandboxFiles( + fakeHandle({ + '/workspace/.tanstack-skills/review/SKILL.md': { + type: 'file', + mode: 0o644, + bytes: new TextEncoder().encode('# Review'), + }, + [`/workspace/${path}`]: { type: 'symlink', mode: 0o777 }, + }), + { blobs: blobs() }, + ) + + const capturedPaths = snapshot.files.map((entry) => entry.path) + expect(capturedPaths).toContain('.tanstack-skills/review/SKILL.md') + expect(capturedPaths).not.toContain(path) + }, + ) + + it('keeps a user .tanstack directory', async () => { + const snapshot = await captureSandboxFiles( + fakeHandle({ + '/workspace/.tanstack/data': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }), + { blobs: blobs() }, + defaultSandboxSnapshotPolicy(), + ) + expect(snapshot.files).toContainEqual( + expect.objectContaining({ path: '.tanstack/data' }), + ) + }) + + it('excludes only the default protected roots while preserving similarly named files', async () => { + const snapshot = await captureSandboxFiles( + fakeHandle({ + '/workspace/.git/config': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + '/workspace/node_modules/pkg/index.js': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + '/workspace/.env.local': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + '/workspace/.gitkeep': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + '/workspace/node_modules.txt': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }), + { blobs: blobs() }, + defaultSandboxSnapshotPolicy('hash'), + ) + expect(snapshot.files.map((entry) => entry.path)).toEqual([ + '.gitkeep', + 'node_modules.txt', + ]) + }) + + it('does not let include override an explicit exclusion', async () => { + const snapshot = await captureSandboxFiles( + fakeHandle({ + '/workspace/keep': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }), + { blobs: blobs() }, + { + include: () => true, + exclude: (path) => path === 'keep', + }, + ) + expect(snapshot.files).toEqual([]) + }) + + it('does not read a non-included sibling file', async () => { + const handle = fakeHandle({ + '/workspace/keep': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + '/workspace/skip': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([2]), + }, + }) + const reads: string[] = [] + const readBytes = handle.fs.readBytes + handle.fs.readBytes = async (path) => { + reads.push(path) + return readBytes(path) + } + await captureSandboxFiles( + handle, + { blobs: blobs() }, + { include: (path) => path === 'keep' }, + ) + expect(reads).toEqual(['/workspace/keep']) + }) + + it('keeps non-harness .sandbox files', async () => { + const snapshot = await captureSandboxFiles( + fakeHandle({ + '/workspace/.sandbox/data': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }), + { blobs: blobs() }, + defaultSandboxSnapshotPolicy(), + ) + expect(snapshot.files).toContainEqual( + expect.objectContaining({ path: '.sandbox/data' }), + ) + }) + + it('rejects a user-created symlink during capture', async () => { + await expect( + captureSandboxFiles( + fakeHandle({ + '/workspace/user-link': { type: 'symlink', mode: 0o777 }, + }), + { blobs: blobs() }, + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_UNSUPPORTED_ENTRY' }) + }) + + const unsupportedEntries: Array<{ kind: string; entry: Entry }> = [ + { kind: 'other', entry: { type: 'other', mode: 0o644 } }, + { + kind: 'executable', + entry: { type: 'file', mode: 0o755, bytes: new Uint8Array([1]) }, + }, + ] + + it.each(unsupportedEntries)( + 'rejects unsupported %s entries', + async ({ entry }) => { + await expect( + captureSandboxFiles( + fakeHandle({ '/workspace/bad': entry }), + { blobs: blobs() }, + defaultSandboxSnapshotPolicy(), + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_UNSUPPORTED_ENTRY' }) + }, + ) + + it('restores after preflighting every blob and removes stale entries', async () => { + const bundle = { blobs: blobs() } + const source = fakeHandle({ + '/workspace/a.txt': { + type: 'file', + mode: 0o644, + bytes: new TextEncoder().encode('a'), + }, + }) + const snapshot = await captureSandboxFiles( + source, + bundle, + defaultSandboxSnapshotPolicy(), + ) + const target = fakeHandle({ + '/workspace/stale.txt': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([2]), + }, + }) + const writes: string[] = [] + const removed: string[] = [] + const originalWrite = target.fs.write + const originalRemove = target.fs.remove + target.fs.write = async (path, bytes) => { + writes.push(path) + await originalWrite(path, bytes) + } + target.fs.remove = async (path) => { + removed.push(path) + await originalRemove!(path) + } + await restoreSandboxFiles(target, bundle, snapshot) + expect(writes).toContain('/workspace/a.txt') + expect(removed).toContain('/workspace/stale.txt') + }) + + it('replaces an existing file with an expected directory', async () => { + const target = fakeHandle({ + '/workspace/node': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }) + await restoreSandboxFiles( + target, + { blobs: blobs() }, + { files: [{ path: 'node', kind: 'dir' }] }, + ) + expect(await target.fs.lstat!('/workspace/node')).toMatchObject({ + type: 'dir', + }) + }) + + it('replaces an existing directory with an expected file', async () => { + const bundle = { blobs: blobs() } + const blobKey = + 'sandbox-files/sha256/2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824' + await bundle.blobs.put(blobKey, new TextEncoder().encode('hello')) + const target = fakeHandle({ + '/workspace/node': { type: 'dir', mode: 0o755 }, + '/workspace/node/old': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }) + await restoreSandboxFiles(target, bundle, { + files: [{ path: 'node', kind: 'file', blobKey, size: 5 }], + }) + expect(await target.fs.lstat!('/workspace/node')).toMatchObject({ + type: 'file', + }) + }) + + it('does not mutate the workspace when a manifest blob is missing', async () => { + const target = fakeHandle({ + '/workspace/a': { type: 'file', mode: 0o644, bytes: new Uint8Array([0]) }, + }) + const writes: string[] = [] + const removed: string[] = [] + target.fs.write = async (path) => { + writes.push(path) + } + target.fs.remove = async (path) => { + removed.push(path) + } + await expect( + restoreSandboxFiles( + target, + { blobs: blobs() }, + { + files: [ + { + path: 'a', + kind: 'file', + blobKey: 'sandbox-files/sha256/missing', + size: 1, + }, + ], + }, + ), + ).rejects.toBeInstanceOf(SandboxSnapshotError) + expect(writes).toEqual([]) + expect(removed).toEqual([]) + }) + + it('propagates operational lstat errors without mutation', async () => { + const target = fakeHandle({ + '/workspace/a': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([0]), + }, + }) + const originalLstat = target.fs.lstat! + target.fs.lstat = async (path) => { + if (path === '/workspace/a') { + throw new Error('permission denied') + } + return originalLstat(path) + } + const writes: string[] = [] + const removed: string[] = [] + target.fs.write = async (path) => { + writes.push(path) + } + target.fs.remove = async (path) => { + removed.push(path) + } + const bundle = { blobs: blobs() } + const blobKey = await putSnapshotBlob(bundle.blobs, new Uint8Array([1])) + await expect( + restoreSandboxFiles(target, bundle, { + files: [{ path: 'a', kind: 'file', blobKey, size: 1 }], + }), + ).rejects.toThrow('permission denied') + expect(writes).toEqual([]) + expect(removed).toEqual([]) + }) + + it('does not mutate when a later destination preflight fails', async () => { + const target = fakeHandle({ + '/workspace/stale': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + '/workspace/b': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([0]), + }, + }) + const originalLstat = target.fs.lstat! + target.fs.lstat = async (path) => { + if (path === '/workspace/b') throw new Error('lstat failed late') + return originalLstat(path) + } + const bundle = { blobs: blobs() } + const firstKey = await putSnapshotBlob(bundle.blobs, new Uint8Array([1])) + const secondKey = await putSnapshotBlob(bundle.blobs, new Uint8Array([2])) + await expect( + restoreSandboxFiles(target, bundle, { + files: [ + { path: 'a', kind: 'file', blobKey: firstKey, size: 1 }, + { path: 'b', kind: 'file', blobKey: secondKey, size: 1 }, + ], + }), + ).rejects.toThrow('lstat failed late') + expect(await target.fs.exists('/workspace/stale')).toBe(true) + }) + + it.each(['ancestor', 'final'])( + 'rejects a %s symlink before writing', + async (position) => { + const target = fakeHandle( + position === 'ancestor' + ? { '/workspace/a': { type: 'symlink', mode: 0o777 } } + : { '/workspace/a': { type: 'symlink', mode: 0o777 } }, + ) + const bundle = { blobs: blobs() } + const blobKey = await putSnapshotBlob(bundle.blobs, new Uint8Array([1])) + const manifest: { files: SandboxSnapshotEntry[] } = + position === 'ancestor' + ? { + files: [ + { + path: 'a/b', + kind: 'file', + blobKey, + size: 1, + }, + ], + } + : { + files: [{ path: 'a', kind: 'file', blobKey, size: 1 }], + } + await expect( + restoreSandboxFiles(target, bundle, manifest), + ).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_UNSUPPORTED_ENTRY', + }) + expect(await target.fs.lstat!('/workspace/a')).toMatchObject({ + type: 'symlink', + }) + }, + ) + + it('restores binary bytes without text conversion', async () => { + const bundle = { blobs: blobs() } + const bytes = new Uint8Array([0, 255, 1, 128]) + const blobKey = await putSnapshotBlob(bundle.blobs, bytes) + const target = fakeHandle({}) + await restoreSandboxFiles(target, bundle, { + files: [ + { + path: 'nested/data.bin', + kind: 'file', + blobKey, + size: bytes.byteLength, + }, + ], + }) + expect(await target.fs.readBytes('/workspace/nested/data.bin')).toEqual( + bytes, + ) + }) + + it('creates missing nested parents during restore', async () => { + const bundle = { blobs: blobs() } + const blobKey = await putSnapshotBlob( + bundle.blobs, + new TextEncoder().encode('ok'), + ) + const target = fakeHandle({}) + await restoreSandboxFiles(target, bundle, { + files: [{ path: 'a/b/c.txt', kind: 'file', blobKey, size: 2 }], + }) + expect(await target.fs.lstat!('/workspace/a/b')).toMatchObject({ + type: 'dir', + }) + }) + + it('rejects trailing separators before restore mutation', async () => { + const target = fakeHandle({}) + await expect( + restoreSandboxFiles( + target, + { blobs: blobs() }, + { files: [{ path: 'nested/', kind: 'dir' }] }, + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_PATH' }) + }) + + it('rejects backslash manifest paths before restore mutation', async () => { + const target = fakeHandle({ + '/workspace/keep': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }) + await expect( + restoreSandboxFiles( + target, + { blobs: blobs() }, + { + files: [{ path: 'nested\\file.txt', kind: 'dir' }], + }, + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_PATH' }) + expect(await target.fs.exists('/workspace/keep')).toBe(true) + }) + + it('rejects a provider list result outside the requested directory', async () => { + const target = fakeHandle({}) + target.fs.list = async () => [ + { name: 'file', path: '/outside/file', type: 'file' }, + ] + await expect( + captureSandboxFiles(target, { blobs: blobs() }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_WORKSPACE' }) + }) + + it('rejects a NUL provider list name before inspecting the child', async () => { + const target = fakeHandle({}) + const inspected: string[] = [] + const originalLstat = target.fs.lstat! + target.fs.lstat = async (path) => { + inspected.push(path) + return originalLstat(path) + } + target.fs.list = async () => [ + { name: 'bad\0name', path: '/workspace/bad\0name', type: 'file' }, + ] + + await expect( + captureSandboxFiles(target, { blobs: blobs() }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_WORKSPACE' }) + + expect(inspected.every((path) => path === '/workspace')).toBe(true) + }) + + it('keeps protected current entries when restoring a fresh sandbox', async () => { + const bundle = { blobs: blobs() } + const blobKey = await putSnapshotBlob(bundle.blobs, new Uint8Array([2])) + const target = fakeHandle({ + '/workspace/.git/config': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + '/workspace/.env.local': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + '/workspace/node_modules/package/index.js': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + '/workspace/.tanstack-projected-workspaceHash/config': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + '/workspace/stale': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }) + await restoreSandboxFiles( + target, + bundle, + { + files: [{ path: 'restored', kind: 'file', blobKey, size: 1 }], + }, + defaultSandboxSnapshotPolicy('workspaceHash'), + ) + expect(await target.fs.exists('/workspace/.git/config')).toBe(true) + expect(await target.fs.exists('/workspace/.env.local')).toBe(true) + expect( + await target.fs.exists('/workspace/node_modules/package/index.js'), + ).toBe(true) + expect( + await target.fs.exists( + '/workspace/.tanstack-projected-workspaceHash/config', + ), + ).toBe(true) + expect(await target.fs.exists('/workspace/stale')).toBe(false) + }) + + it('removes stale descendants while preserving a protected sibling', async () => { + const target = fakeHandle({ + '/workspace/config/.env': { + type: 'file', + mode: 0o600, + bytes: new TextEncoder().encode('SECRET=value'), + }, + '/workspace/config/old.txt': { + type: 'file', + mode: 0o644, + bytes: new TextEncoder().encode('old'), + }, + }) + const mutations = fsMutationLog(target) + + await restoreSandboxFiles(target, { blobs: blobs() }, { files: [] }) + + expect(await target.fs.exists('/workspace/config/.env')).toBe(true) + expect(await target.fs.exists('/workspace/config/old.txt')).toBe(false) + expect(mutations).toEqual({ + writes: [], + removes: ['/workspace/config/old.txt'], + mkdirs: [], + }) + }) + + it.each([ + 'src/.git/config', + 'src/node_modules/package/index.js', + 'src/.env.local', + ])( + 'does not restore over protected segments at any depth: %s', + async (path) => { + const bundle = { blobs: blobs() } + const blobKey = await putSnapshotBlob(bundle.blobs, new Uint8Array([2])) + const target = fakeHandle({ + [`/workspace/${path}`]: { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + '/workspace/stale': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }) + await restoreSandboxFiles( + target, + bundle, + { files: [{ path: 'restored', kind: 'file', blobKey, size: 1 }] }, + defaultSandboxSnapshotPolicy('workspaceHash'), + ) + expect(await target.fs.exists(`/workspace/${path}`)).toBe(true) + expect(await target.fs.exists('/workspace/stale')).toBe(false) + }, + ) + + it('restores a projection-looking user path when the workspace hash is unknown', async () => { + const bundle = { blobs: blobs() } + const blobKey = await putSnapshotBlob(bundle.blobs, new Uint8Array([2])) + const target = fakeHandle({ + '/workspace/.tanstack-projected-secret/config': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }) + await restoreSandboxFiles( + target, + bundle, + { files: [{ path: 'restored', kind: 'file', blobKey, size: 1 }] }, + { exclude: () => false }, + ) + expect( + await target.fs.exists('/workspace/.tanstack-projected-secret/config'), + ).toBe(false) + expect(await target.fs.exists('/workspace/restored')).toBe(true) + }) + + it('rejects an explicit empty directory excluded by include before any access', async () => { + const { store, gets, heads } = blobsWithAccessLog() + const target = fakeHandle({}) + const mutations = fsMutationLog(target) + + await expect( + restoreSandboxFiles( + target, + { blobs: store }, + { files: [{ path: 'private', kind: 'dir' }] }, + { include: (_path, kind) => kind !== 'dir' }, + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_PATH' }) + + expect(gets).toEqual([]) + expect(heads).toEqual([]) + expect(mutations).toEqual({ writes: [], removes: [], mkdirs: [] }) + }) + + it('rejects a file under an excluded custom ancestor before any access', async () => { + const { store, gets, heads } = blobsWithAccessLog() + const target = fakeHandle({}) + const mutations = fsMutationLog(target) + const blobKey = + 'sandbox-files/sha256/2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824' + + await expect( + restoreSandboxFiles( + target, + { blobs: store }, + { + files: [ + { path: 'private/secret.txt', kind: 'file', blobKey, size: 5 }, + ], + }, + { exclude: (path) => path === 'private' }, + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_PATH' }) + + expect(gets).toEqual([]) + expect(heads).toEqual([]) + expect(mutations).toEqual({ writes: [], removes: [], mkdirs: [] }) + }) + + it.each([ + { + name: 'the exact protected projection ancestor', + path: '.tanstack-projected-workspaceHash/config.ts', + policy: defaultSandboxSnapshotPolicy('workspaceHash'), + }, + { + name: 'a projection-looking user ancestor without the workspace hash', + path: '.tanstack-projected/config.ts', + policy: defaultSandboxSnapshotPolicy('workspaceHash'), + }, + ])( + 'handles $name in the desired manifest before blob access', + async ({ path, policy }) => { + const { store, gets, heads } = blobsWithAccessLog() + const target = fakeHandle({}) + const mutations = fsMutationLog(target) + const blobKey = + 'sandbox-files/sha256/2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824' + await store.put(blobKey, new TextEncoder().encode('hello')) + const restore = restoreSandboxFiles( + target, + { blobs: store }, + { files: [{ path, kind: 'file', blobKey, size: 5 }] }, + policy, + ) + + if (path.startsWith('.tanstack-projected-workspaceHash/')) { + await expect(restore).rejects.toMatchObject({ + code: 'SANDBOX_SNAPSHOT_INVALID_PATH', + }) + expect(gets).toEqual([]) + expect(heads).toEqual([]) + expect(mutations).toEqual({ writes: [], removes: [], mkdirs: [] }) + } else { + await expect(restore).resolves.toBeUndefined() + expect(await target.fs.exists(`/workspace/${path}`)).toBe(true) + } + }, + ) + + it('rejects a desired file that would replace a protected current descendant before access', async () => { + const { store, gets, heads } = blobsWithAccessLog() + const target = fakeHandle({ + '/workspace/config/.env': { + type: 'file', + mode: 0o600, + bytes: new Uint8Array([1]), + }, + }) + const mutations = fsMutationLog(target) + const blobKey = + 'sandbox-files/sha256/2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824' + + await expect( + restoreSandboxFiles( + target, + { blobs: store }, + { files: [{ path: 'config', kind: 'file', blobKey, size: 5 }] }, + defaultSandboxSnapshotPolicy('workspaceHash'), + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_PATH' }) + + expect(gets).toEqual([]) + expect(heads).toEqual([]) + expect(mutations).toEqual({ writes: [], removes: [], mkdirs: [] }) + }) + + it('allows an implicit parent directory excluded by include for an included file', async () => { + const store = blobs() + const bytes = new TextEncoder().encode('hello') + const blobKey = await putSnapshotBlob(store, bytes) + const target = fakeHandle({}) + const mutations = fsMutationLog(target) + + await restoreSandboxFiles( + target, + { blobs: store }, + { + files: [ + { path: 'generated/index.txt', kind: 'file', blobKey, size: 5 }, + ], + }, + { include: (_path, kind) => kind === 'file' }, + ) + + expect(mutations.mkdirs).toContain('/workspace/generated') + expect(mutations.writes).toEqual(['/workspace/generated/index.txt']) + expect(await target.fs.read('/workspace/generated/index.txt')).toBe('hello') + }) + + it('rejects an invalid file blob key before any blob read or workspace mutation', async () => { + const { store, gets, heads } = blobsWithAccessLog() + const target = fakeHandle({ + '/workspace/keep': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }) + const writes: string[] = [] + const removes: string[] = [] + const write = target.fs.write + const remove = target.fs.remove + target.fs.write = async (path, data) => { + writes.push(path) + await write(path, data) + } + target.fs.remove = async (path) => { + removes.push(path) + await remove(path) + } + + await expect( + restoreSandboxFiles( + target, + { blobs: store }, + { + files: [ + { + path: 'unsafe.txt', + kind: 'file', + blobKey: 'sandbox-files/sha256/not-a-file-hash', + size: 1, + }, + ], + }, + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_PATH' }) + + expect(gets).toEqual([]) + expect(heads).toEqual([]) + expect(writes).toEqual([]) + expect(removes).toEqual([]) + expect(await target.fs.exists('/workspace/keep')).toBe(true) + }) + + it('rejects structural manifest errors before any blob read', async () => { + const { store, gets, heads } = blobsWithAccessLog() + const target = fakeHandle({}) + const blobKey = await putSnapshotBlob(store, new Uint8Array([1])) + + await expect( + restoreSandboxFiles( + target, + { blobs: store }, + { + files: [ + { path: 'parent', kind: 'file', blobKey, size: 1 }, + { path: 'parent/child', kind: 'file', blobKey, size: 1 }, + ], + }, + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_PATH' }) + + expect(gets).toEqual([]) + expect(heads).toEqual([]) + }) + + it('skips the known projection tree during custom-policy restore', async () => { + const bundle = { blobs: blobs() } + const blobKey = await putSnapshotBlob(bundle.blobs, new Uint8Array([2])) + const target = fakeHandle({ + '/workspace/.tanstack-projected-secret/config': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + '/workspace/stale': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }) + + await restoreSandboxFiles( + target, + bundle, + { files: [{ path: 'restored', kind: 'file', blobKey, size: 1 }] }, + { + workspaceHash: 'secret', + exclude: () => false, + }, + ) + + expect( + await target.fs.exists('/workspace/.tanstack-projected-secret/config'), + ).toBe(true) + expect(await target.fs.exists('/workspace/stale')).toBe(false) + }) + + it('does not read the known projection file when custom policy allows everything', async () => { + const handle = fakeHandle({ + '/workspace/.tanstack-projected-secret': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + '/workspace/keep': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([2]), + }, + }) + const reads: string[] = [] + const originalReadBytes = handle.fs.readBytes + handle.fs.readBytes = async (path) => { + reads.push(path) + return originalReadBytes(path) + } + const bundle = blobs() + const puts: string[] = [] + const originalPut = bundle.put + bundle.put = async (key, body) => { + puts.push(key) + return originalPut(key, body) + } + + const snapshot = await captureSandboxFiles( + handle, + { blobs: bundle }, + { workspaceHash: 'secret', include: () => true, exclude: () => false }, + ) + + expect(reads).toEqual(['/workspace/keep']) + expect(puts).toHaveLength(1) + expect(snapshot.files).toEqual([ + expect.objectContaining({ path: 'keep', kind: 'file' }), + ]) + }) + + it('does not inspect or persist a projection tree when custom policy allows everything', async () => { + const handle = fakeHandle({ + '/workspace/.tanstack-projected-secret/config': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + '/workspace/keep/config': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([2]), + }, + }) + const inspected: string[] = [] + const originalLstat = handle.fs.lstat! + handle.fs.lstat = async (path) => { + inspected.push(path) + return originalLstat(path) + } + const reads: string[] = [] + const originalReadBytes = handle.fs.readBytes + handle.fs.readBytes = async (path) => { + reads.push(path) + return originalReadBytes(path) + } + + const snapshot = await captureSandboxFiles( + handle, + { blobs: blobs() }, + { workspaceHash: 'secret', include: () => true, exclude: () => false }, + ) + + expect(inspected).not.toContain('/workspace/.tanstack-projected-secret') + expect(reads).toEqual(['/workspace/keep/config']) + expect(snapshot.files).toEqual([ + expect.objectContaining({ path: 'keep/config', kind: 'file' }), + ]) + }) + + it('validates every reference to a shared blob before mutation', async () => { + const bundle = { blobs: blobs() } + const blobKey = await putSnapshotBlob(bundle.blobs, new Uint8Array([1])) + const target = fakeHandle({ + '/workspace/keep': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }) + await expect( + restoreSandboxFiles(target, bundle, { + files: [ + { path: 'a', kind: 'file', blobKey, size: 1 }, + { path: 'b', kind: 'file', blobKey, size: 2 }, + ], + }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_BLOB' }) + expect(await target.fs.exists('/workspace/keep')).toBe(true) + }) + + it('rejects a blob whose bytes do not match its content-addressed key', async () => { + const bundle = { blobs: blobs() } + const captured = await captureSandboxFiles( + fakeHandle({ + '/workspace/source': { + type: 'file', + mode: 0o644, + bytes: new TextEncoder().encode('expected'), + }, + }), + bundle, + ) + const entry = captured.files[0] + if (!entry || entry.kind !== 'file') throw new Error('expected file') + await bundle.blobs.put(entry.blobKey, new TextEncoder().encode('tampered')) + const target = fakeHandle({ + '/workspace/keep': { + type: 'file', + mode: 0o644, + bytes: new Uint8Array([1]), + }, + }) + await expect( + restoreSandboxFiles(target, bundle, captured), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_BLOB' }) + expect(await target.fs.exists('/workspace/keep')).toBe(true) + }) +}) diff --git a/packages/ai-sandbox/tests/testkit-subpath.test.ts b/packages/ai-sandbox/tests/testkit-subpath.test.ts index fe789bdf1f..419ae2fb53 100644 --- a/packages/ai-sandbox/tests/testkit-subpath.test.ts +++ b/packages/ai-sandbox/tests/testkit-subpath.test.ts @@ -2,21 +2,26 @@ import { describe, expect, it } from 'vitest' import { makeFakeShellSpawn, runJournalConformance, + runSandboxCheckpointForkConformance, + runSandboxCheckpointStoreConformance, runTakeoverConformance, } from '@tanstack/ai-sandbox/testkit' +import { memorySandboxSnapshots } from '@tanstack/ai-sandbox' + +runSandboxCheckpointForkConformance( + 'published testkit consumer', + memorySandboxSnapshots, +) /** - * Proves the `@tanstack/ai-sandbox/testkit` subpath actually resolves and - * ships `makeFakeShellSpawn` — the thing that silently breaks when the - * package.json `exports` map and the build entry list disagree (a subpath - * that publint/test:build waves through but that consumers can't import). + * Proves the `@tanstack/ai-sandbox/testkit` subpath resolves as a consumer + * would resolve it. */ describe('@tanstack/ai-sandbox/testkit subpath', () => { - it('ships both provider conformance suites', () => { - // The provider packages (and third-party providers outside this repo) reach - // these ONLY through the built subpath, so an export that exists in `src` - // but never lands in `dist` is invisible until a consumer breaks. + it('ships the provider and checkpoint conformance suites', () => { expect(typeof runJournalConformance).toBe('function') + expect(typeof runSandboxCheckpointForkConformance).toBe('function') + expect(typeof runSandboxCheckpointStoreConformance).toBe('function') expect(typeof runTakeoverConformance).toBe('function') }) diff --git a/packages/ai/src/middlewares/index.ts b/packages/ai/src/middlewares/index.ts index 1470cdcff1..d3fe58ecc9 100644 --- a/packages/ai/src/middlewares/index.ts +++ b/packages/ai/src/middlewares/index.ts @@ -12,6 +12,8 @@ export { type ContentFilteredInfo, } from './content-guard' +export { CapabilityRegistry } from '../activities/chat/middleware/capabilities' + // otelMiddleware is exported from the dedicated subpath // `@tanstack/ai/middlewares/otel` so that importing the main middlewares barrel // does not eagerly require `@opentelemetry/api` (which is an optional peer diff --git a/packages/ai/tests/middlewares/index.test.ts b/packages/ai/tests/middlewares/index.test.ts new file mode 100644 index 0000000000..367bdf0e34 --- /dev/null +++ b/packages/ai/tests/middlewares/index.test.ts @@ -0,0 +1,17 @@ +import { describe, expect, it } from 'vitest' +import { CapabilityRegistry } from '../../src/middlewares' +import { createCapability } from '../../src/activities/chat/middleware/capabilities' + +describe('middlewares public exports', () => { + it('exports a working CapabilityRegistry', () => { + const registry = new CapabilityRegistry() + const capability = createCapability()('value') + const [getValue, provideValue] = capability + const context = { capabilities: registry } + + expect(registry.has(capability)).toBe(false) + provideValue(context, 42) + expect(registry.has(capability)).toBe(true) + expect(getValue(context)).toBe(42) + }) +}) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index f600b7fd6d..f203551064 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -2353,6 +2353,9 @@ importers: '@tanstack/ai': specifier: workspace:* version: link:../ai + '@tanstack/ai-persistence': + specifier: workspace:* + version: link:../ai-persistence '@vitest/coverage-v8': specifier: 4.0.14 version: 4.0.14(supports-color@7.2.0)(vitest@4.1.10) diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index c6d9853490..f8f0e6dc24 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -34,6 +34,7 @@ import { Route as ApiToolsTestRouteImport } from './routes/api.tools-test' import { Route as ApiToolCallLifecycleWireRouteImport } from './routes/api.tool-call-lifecycle-wire' import { Route as ApiSummarizeRouteImport } from './routes/api.summarize' import { Route as ApiSandboxToolHistoryRouteImport } from './routes/api.sandbox-tool-history' +import { Route as ApiSandboxFilePersistenceRouteImport } from './routes/api.sandbox-file-persistence' import { Route as ApiSandboxDurabilityRouteImport } from './routes/api.sandbox-durability' import { Route as ApiPersistenceDurabilityRouteImport } from './routes/api.persistence-durability' import { Route as ApiOtelUsageRouteImport } from './routes/api.otel-usage' @@ -205,6 +206,12 @@ const ApiSandboxToolHistoryRoute = ApiSandboxToolHistoryRouteImport.update({ path: '/api/sandbox-tool-history', getParentRoute: () => rootRouteImport, } as any) +const ApiSandboxFilePersistenceRoute = + ApiSandboxFilePersistenceRouteImport.update({ + id: '/api/sandbox-file-persistence', + path: '/api/sandbox-file-persistence', + getParentRoute: () => rootRouteImport, + } as any) const ApiSandboxDurabilityRoute = ApiSandboxDurabilityRouteImport.update({ id: '/api/sandbox-durability', path: '/api/sandbox-durability', @@ -478,6 +485,7 @@ export interface FileRoutesByFullPath { '/api/otel-usage': typeof ApiOtelUsageRoute '/api/persistence-durability': typeof ApiPersistenceDurabilityRoute '/api/sandbox-durability': typeof ApiSandboxDurabilityRoute + '/api/sandbox-file-persistence': typeof ApiSandboxFilePersistenceRoute '/api/sandbox-tool-history': typeof ApiSandboxToolHistoryRoute '/api/summarize': typeof ApiSummarizeRoute '/api/tool-call-lifecycle-wire': typeof ApiToolCallLifecycleWireRoute @@ -547,6 +555,7 @@ export interface FileRoutesByTo { '/api/otel-usage': typeof ApiOtelUsageRoute '/api/persistence-durability': typeof ApiPersistenceDurabilityRoute '/api/sandbox-durability': typeof ApiSandboxDurabilityRoute + '/api/sandbox-file-persistence': typeof ApiSandboxFilePersistenceRoute '/api/sandbox-tool-history': typeof ApiSandboxToolHistoryRoute '/api/summarize': typeof ApiSummarizeRoute '/api/tool-call-lifecycle-wire': typeof ApiToolCallLifecycleWireRoute @@ -617,6 +626,7 @@ export interface FileRoutesById { '/api/otel-usage': typeof ApiOtelUsageRoute '/api/persistence-durability': typeof ApiPersistenceDurabilityRoute '/api/sandbox-durability': typeof ApiSandboxDurabilityRoute + '/api/sandbox-file-persistence': typeof ApiSandboxFilePersistenceRoute '/api/sandbox-tool-history': typeof ApiSandboxToolHistoryRoute '/api/summarize': typeof ApiSummarizeRoute '/api/tool-call-lifecycle-wire': typeof ApiToolCallLifecycleWireRoute @@ -688,6 +698,7 @@ export interface FileRouteTypes { | '/api/otel-usage' | '/api/persistence-durability' | '/api/sandbox-durability' + | '/api/sandbox-file-persistence' | '/api/sandbox-tool-history' | '/api/summarize' | '/api/tool-call-lifecycle-wire' @@ -757,6 +768,7 @@ export interface FileRouteTypes { | '/api/otel-usage' | '/api/persistence-durability' | '/api/sandbox-durability' + | '/api/sandbox-file-persistence' | '/api/sandbox-tool-history' | '/api/summarize' | '/api/tool-call-lifecycle-wire' @@ -826,6 +838,7 @@ export interface FileRouteTypes { | '/api/otel-usage' | '/api/persistence-durability' | '/api/sandbox-durability' + | '/api/sandbox-file-persistence' | '/api/sandbox-tool-history' | '/api/summarize' | '/api/tool-call-lifecycle-wire' @@ -896,6 +909,7 @@ export interface RootRouteChildren { ApiOtelUsageRoute: typeof ApiOtelUsageRoute ApiPersistenceDurabilityRoute: typeof ApiPersistenceDurabilityRoute ApiSandboxDurabilityRoute: typeof ApiSandboxDurabilityRoute + ApiSandboxFilePersistenceRoute: typeof ApiSandboxFilePersistenceRoute ApiSandboxToolHistoryRoute: typeof ApiSandboxToolHistoryRoute ApiSummarizeRoute: typeof ApiSummarizeRoute ApiToolCallLifecycleWireRoute: typeof ApiToolCallLifecycleWireRoute @@ -1083,6 +1097,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiSandboxToolHistoryRouteImport parentRoute: typeof rootRouteImport } + '/api/sandbox-file-persistence': { + id: '/api/sandbox-file-persistence' + path: '/api/sandbox-file-persistence' + fullPath: '/api/sandbox-file-persistence' + preLoaderRoute: typeof ApiSandboxFilePersistenceRouteImport + parentRoute: typeof rootRouteImport + } '/api/sandbox-durability': { id: '/api/sandbox-durability' path: '/api/sandbox-durability' @@ -1493,6 +1514,7 @@ const rootRouteChildren: RootRouteChildren = { ApiOtelUsageRoute: ApiOtelUsageRoute, ApiPersistenceDurabilityRoute: ApiPersistenceDurabilityRoute, ApiSandboxDurabilityRoute: ApiSandboxDurabilityRoute, + ApiSandboxFilePersistenceRoute: ApiSandboxFilePersistenceRoute, ApiSandboxToolHistoryRoute: ApiSandboxToolHistoryRoute, ApiSummarizeRoute: ApiSummarizeRoute, ApiToolCallLifecycleWireRoute: ApiToolCallLifecycleWireRoute, diff --git a/testing/e2e/src/routes/api.sandbox-file-persistence.ts b/testing/e2e/src/routes/api.sandbox-file-persistence.ts new file mode 100644 index 0000000000..22bd9f3ce3 --- /dev/null +++ b/testing/e2e/src/routes/api.sandbox-file-persistence.ts @@ -0,0 +1,463 @@ +import { createFileRoute } from '@tanstack/react-router' +import { EventType, chat } from '@tanstack/ai' +import { + InMemorySandboxInstanceStore, + defineSandbox, + defineWorkspace, + forkFromSandboxSnapshot, + memorySandboxSnapshots, + resolveSnapshotArtifact, + saveNamedSandboxSnapshot, + withSandbox, +} from '@tanstack/ai-sandbox' +import { withPersistence } from '@tanstack/ai-persistence' +import type { + DefaultMessageMetadataByModality, + StreamChunk, + TextAdapter, + TextOptions, +} from '@tanstack/ai' +import type { + SandboxCapabilities, + SandboxCheckpoint, + SandboxFsStat, + SandboxHandle, + SandboxProvider, +} from '@tanstack/ai-sandbox' + +const capabilities: SandboxCapabilities = { + fs: true, + exec: true, + env: true, + ports: false, + backgroundProcesses: false, + writableStdin: false, + killableProcesses: false, + snapshots: false, + networkPolicy: false, + durableFilesystem: false, + fork: false, +} + +type WorkspaceEntry = { kind: 'dir' } | { kind: 'file'; content: string } + +type SnapshotValue = SandboxCheckpoint | null + +function sameValue(left: SnapshotValue, right: SnapshotValue): boolean { + if (!left || !right) return left === right + return ( + left.id === right.id && + left.threadId === right.threadId && + left.parentCheckpointId === right.parentCheckpointId && + left.createdAt === right.createdAt && + left.reason === right.reason && + left.label === right.label && + left.sourceRunId === right.sourceRunId && + left.files.length === right.files.length && + left.files.every((entry, index) => { + const other = right.files[index] + return ( + other !== undefined && + entry.path === other.path && + entry.kind === other.kind && + (entry.kind !== 'file' || + (other.kind === 'file' && + entry.blobKey === other.blobKey && + entry.size === other.size)) + ) + }) && + left.conversation.length === right.conversation.length && + left.conversation.every( + (message, index) => + message.role === right.conversation[index]?.role && + message.content === right.conversation[index]?.content, + ) && + left.artifacts.length === right.artifacts.length && + left.artifacts.every((artifact, index) => { + const other = right.artifacts[index] + return ( + other !== undefined && + artifact.artifactId === other.artifactId && + artifact.name === other.name && + artifact.mimeType === other.mimeType && + artifact.blobKey === other.blobKey && + artifact.size === other.size && + artifact.createdAt === other.createdAt + ) + }) + ) +} + +function sameValues( + left: ReadonlyArray, + right: ReadonlyArray, +): boolean { + return ( + left.length === right.length && + left.every((value, index) => sameValue(value, right[index])) + ) +} + +export function sameSerializableValue(left: unknown, right: unknown): boolean { + if (Object.is(left, right)) return true + if (typeof left !== typeof right || left === null || right === null) + return false + if (left instanceof Date || right instanceof Date) { + return ( + left instanceof Date && + right instanceof Date && + left.getTime() === right.getTime() + ) + } + if (Array.isArray(left) || Array.isArray(right)) { + if (!Array.isArray(left) || !Array.isArray(right)) return false + return ( + left.length === right.length && + left.every((value, index) => sameSerializableValue(value, right[index])) + ) + } + if (typeof left !== 'object' || typeof right !== 'object') return false + const leftEntries = Object.entries(left).sort(([first], [second]) => + first.localeCompare(second), + ) + const rightEntries = Object.entries(right).sort(([first], [second]) => + first.localeCompare(second), + ) + return ( + leftEntries.length === rightEntries.length && + leftEntries.every(([key, value], index) => { + const other = rightEntries[index] + return ( + other !== undefined && + key === other[0] && + sameSerializableValue(value, other[1]) + ) + }) + ) +} + +export async function forkWithSourceMessageSnapshot({ + loadSourceMessages, + fork, +}: { + loadSourceMessages: () => Promise + fork: () => Promise +}) { + const sourceMessagesBeforeFork = structuredClone(await loadSourceMessages()) + const result = await fork() + const sourceMessagesAfterFork = await loadSourceMessages() + return { + result, + sourceMessagesUnchanged: sameSerializableValue( + sourceMessagesBeforeFork, + sourceMessagesAfterFork, + ), + } +} + +function fileContent(entry: WorkspaceEntry | undefined): string { + return entry?.kind === 'file' ? entry.content : '' +} + +function fakeHandle( + id: string, + files: Map, +): SandboxHandle { + const lstat = (path: string): Promise => { + const entry = files.get(path) + if (!entry) return Promise.resolve(undefined) + return Promise.resolve( + entry.kind === 'dir' + ? { type: 'dir', mode: 0o755 } + : { + type: 'file', + mode: 0o644, + size: new TextEncoder().encode(entry.content).byteLength, + }, + ) + } + return { + id, + provider: 'fake', + capabilities, + fs: { + read: async (path) => fileContent(files.get(path)), + readBytes: async (path) => + new TextEncoder().encode(fileContent(files.get(path))), + write: async (path, value) => { + files.set(path, { + kind: 'file', + content: + typeof value === 'string' ? value : new TextDecoder().decode(value), + }) + }, + list: async (path) => { + const prefix = `${path}/` + return [...files].flatMap(([entryPath, entry]) => { + const name = entryPath.startsWith(prefix) + ? entryPath.slice(prefix.length) + : '' + return name && !name.includes('/') + ? [{ name, path: entryPath, type: entry.kind }] + : [] + }) + }, + lstat, + mkdir: async (path) => { + files.set(path, { kind: 'dir' }) + }, + remove: async (path) => { + files.delete(path) + }, + rename: async () => {}, + exists: async (path) => files.has(path), + }, + git: { + clone: async () => {}, + status: async () => '', + add: async () => {}, + commit: async () => {}, + push: async () => {}, + pull: async () => {}, + branch: async () => 'main', + }, + process: { + exec: async () => ({ stdout: '', stderr: '', exitCode: 0 }), + spawn: async () => { + throw new Error('not supported') + }, + }, + ports: { + connect: async () => { + throw new Error('not supported') + }, + }, + env: { set: async () => {} }, + destroy: async () => {}, + } +} + +function fixedRun(threadId: string, runId: string): AsyncIterable { + return (async function* () { + yield { type: EventType.RUN_STARTED, threadId, runId, timestamp: 1 } + yield { + type: EventType.TEXT_MESSAGE_START, + threadId, + runId, + timestamp: 1, + messageId: 'automatic-message', + role: 'assistant', + } + yield { + type: EventType.TEXT_MESSAGE_CONTENT, + threadId, + runId, + timestamp: 1, + messageId: 'automatic-message', + delta: 'automatic conversation', + } + yield { + type: EventType.TEXT_MESSAGE_END, + threadId, + runId, + timestamp: 1, + messageId: 'automatic-message', + } + yield { + type: EventType.RUN_FINISHED, + threadId, + runId, + finishReason: 'stop', + timestamp: 1, + } + })() +} + +type FixedAdapter = TextAdapter< + 'test-model', + Record, + readonly ['text'], + DefaultMessageMetadataByModality, + readonly [], + unknown, + unknown +> + +const adapter: FixedAdapter = { + kind: 'text', + name: 'fixed', + model: 'test-model', + '~types': { + providerOptions: {}, + inputModalities: ['text'], + messageMetadataByModality: { + text: undefined, + image: undefined, + audio: undefined, + video: undefined, + document: undefined, + }, + toolCapabilities: [], + toolCallMetadata: undefined, + systemPromptMetadata: undefined, + }, + chatStream: (options: TextOptions>) => + fixedRun( + options.threadId ?? 'missing-thread', + options.runId ?? 'missing-run', + ), + structuredOutput: () => Promise.resolve({ data: {}, rawText: '{}' }), +} + +export const Route = createFileRoute('/api/sandbox-file-persistence')({ + server: { + handlers: { + POST: async () => { + const threadId = `snapshot-source-${crypto.randomUUID()}` + const destinationThreadId = `snapshot-fork-${crypto.randomUUID()}` + const sourceFiles = new Map([ + ['/workspace', { kind: 'dir' }], + ['/workspace/notes.txt', { kind: 'file', content: 'saved file' }], + ['/workspace/empty', { kind: 'dir' }], + ['/workspace/.env', { kind: 'file', content: 'secret' }], + ['/workspace/.git', { kind: 'dir' }], + ['/workspace/.git/config', { kind: 'file', content: 'private' }], + ]) + const restoredFiles = new Map([ + ['/workspace', { kind: 'dir' }], + ]) + let originalExists = true + const provider: SandboxProvider = { + name: 'fake', + capabilities: () => capabilities, + create: async () => fakeHandle('created', restoredFiles), + resume: async () => + originalExists ? fakeHandle('original', sourceFiles) : null, + destroy: async () => {}, + } + const snapshots = await memorySandboxSnapshots() + const instances = new InMemorySandboxInstanceStore() + const sandbox = defineSandbox({ + id: 'file-persistence', + provider, + workspace: defineWorkspace({ source: { type: 'none' } }), + fileEvents: false, + }) + const key = sandbox.key({ threadId, runId: 'save', store: instances }) + await instances.upsert({ + key, + provider: 'fake', + providerSandboxId: 'original', + threadId, + updatedAt: 1, + }) + await snapshots.persistence.stores.messages.saveThread(threadId, [ + { role: 'user', content: 'saved conversation' }, + ]) + await snapshots.persistence.stores.blobs.put( + 'artifact-source', + 'artifact data', + ) + await snapshots.persistence.stores.artifacts.save({ + artifactId: 'artifact-1', + threadId, + runId: 'save', + name: 'artifact.txt', + mimeType: 'text/plain', + blobKey: 'artifact-source', + size: 13, + createdAt: 1, + }) + + const saved = await saveNamedSandboxSnapshot({ + definition: sandbox, + threadId, + runId: 'save', + instances, + snapshots, + label: 'release-1', + }) + originalExists = false + const recovery = chat({ + adapter, + messages: [{ role: 'user', content: 'recover' }], + runId: 'recover', + threadId, + middleware: [ + withPersistence(snapshots.persistence), + withSandbox(sandbox, { instances, snapshots }), + ], + }) + for await (const _ of recovery) void _ + const automaticHeadId = await snapshots.checkpoints.getHead(threadId) + if (!automaticHeadId || automaticHeadId === saved.id) + throw new Error('Expected a newer automatic checkpoint') + const automaticCheckpoint = + await snapshots.checkpoints.get(automaticHeadId) + if (!automaticCheckpoint) + throw new Error('Expected the automatic checkpoint') + const artifact = await resolveSnapshotArtifact({ + threadId, + checkpointId: saved.id, + artifactId: 'artifact-1', + snapshots, + }) + const sourceBeforeFork = await snapshots.checkpoints.list(threadId) + const { result: fork, sourceMessagesUnchanged } = + await forkWithSourceMessageSnapshot({ + loadSourceMessages: () => + snapshots.persistence.stores.messages.loadThread(threadId), + fork: () => + forkFromSandboxSnapshot({ + sourceThreadId: threadId, + sourceCheckpointId: saved.id, + destinationThreadId, + snapshots, + destinationCheckpointId: 'fork-root', + createdAt: 2, + }), + }) + const sourceAfterFork = await snapshots.checkpoints.list(threadId) + const sourceHeadAfterFork = + await snapshots.checkpoints.getHead(threadId) + const sourceCheckpointAfterFork = + await snapshots.checkpoints.get(automaticHeadId) + const readFile = restoredFiles.get('/workspace/notes.txt') + return Response.json({ + namedSave: saved.label, + recoveredFiles: readFile?.kind === 'file' ? ['notes.txt'] : [], + recoveredFileBytes: + readFile?.kind === 'file' + ? Array.from(new TextEncoder().encode(readFile.content)) + : [], + recoveredEmptyDirectories: + restoredFiles.get('/workspace/empty')?.kind === 'dir' + ? ['empty'] + : [], + artifactText: new TextDecoder().decode(artifact.bytes), + conversation: saved.conversation, + excluded: ['.env', '.git'].filter( + (name) => + !saved.files.some( + (entry) => + entry.path === name || entry.path.startsWith(`${name}/`), + ), + ), + fork: { + selectedCheckpointIsHead: saved.id === automaticHeadId, + files: fork.files + .filter((entry) => entry.kind === 'file') + .map((entry) => entry.path), + sourceThreadUnchanged: + sameValues(sourceBeforeFork, sourceAfterFork) && + automaticHeadId === sourceHeadAfterFork && + sameValue(automaticCheckpoint, sourceCheckpointAfterFork), + sourceMessagesUnchanged, + conversation: fork.conversation, + }, + automaticConversation: automaticCheckpoint.conversation, + }) + }, + }, + }, +}) diff --git a/testing/e2e/tests/sandbox-file-persistence.spec.ts b/testing/e2e/tests/sandbox-file-persistence.spec.ts new file mode 100644 index 0000000000..eeadccf0bf --- /dev/null +++ b/testing/e2e/tests/sandbox-file-persistence.spec.ts @@ -0,0 +1,60 @@ +import { expect, test } from '@playwright/test' +import { + forkWithSourceMessageSnapshot, + sameSerializableValue, +} from '../src/routes/api.sandbox-file-persistence' + +test.describe('sandbox portable file snapshots', () => { + test('captures source messages before the fork and compares dates', async () => { + const messages = [{ id: 'message-1', nested: { value: 'before' } }] + const result = await forkWithSourceMessageSnapshot({ + loadSourceMessages: () => Promise.resolve(messages), + fork: () => { + const message = messages[0] + if (!message) throw new Error('Expected a message') + message.nested.value = 'after' + return Promise.resolve() + }, + }) + + expect(result.sourceMessagesUnchanged).toBe(false) + expect( + sameSerializableValue( + { createdAt: new Date('2026-01-01T00:00:00.000Z') }, + { createdAt: new Date('2026-01-02T00:00:00.000Z') }, + ), + ).toBe(false) + }) + + test('saves, recovers, reads, and forks an immutable snapshot', async ({ + request, + }) => { + const response = await request.post('/api/sandbox-file-persistence') + + expect(response.ok()).toBe(true) + await expect(response.json()).resolves.toEqual({ + namedSave: 'release-1', + recoveredFiles: ['notes.txt'], + recoveredFileBytes: [115, 97, 118, 101, 100, 32, 102, 105, 108, 101], + recoveredEmptyDirectories: ['empty'], + artifactText: 'artifact data', + conversation: [{ content: 'saved conversation', role: 'user' }], + automaticConversation: [ + { content: 'recover', role: 'user' }, + { + content: 'automatic conversation', + id: 'automatic-message', + role: 'assistant', + }, + ], + excluded: ['.env', '.git'], + fork: { + selectedCheckpointIsHead: false, + files: ['notes.txt'], + sourceThreadUnchanged: true, + sourceMessagesUnchanged: true, + conversation: [{ content: 'saved conversation', role: 'user' }], + }, + }) + }) +}) From 32771f9383f88a3f06a5a8f1d4c59f6aecd667c3 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Fri, 14 Aug 2026 14:38:25 +0200 Subject: [PATCH 02/12] fix(ai): harden sandbox snapshot CI and public policy export Remove the accidental CapabilityRegistry export. Type SQLite checkpoint tests, reject non-string ids, and keep the original error when ROLLBACK fails. Export defaultSandboxSnapshotPolicy so a custom redact can keep the default exclusions. Assign the capture task before capture work starts. --- .changeset/quiet-sandbox-middleware.md | 5 - docs/config.json | 14 +- docs/sandbox/portable-snapshots.md | 38 +++++- .../src/lib/sqlite-persistence.test.ts | 103 +++++++++++++- .../src/lib/sqlite-persistence.ts | 38 ++++-- .../tests/persistence-completion.test.ts | 5 +- .../ai-sandbox/skills/ai-sandbox/SKILL.md | 45 +++--- packages/ai-sandbox/src/index.ts | 2 +- packages/ai-sandbox/src/middleware.ts | 128 +++++++++--------- .../tests/ai-middleware-subpath.test.ts | 6 - packages/ai-sandbox/tests/fakes.ts | 28 +++- .../tests/snapshot-operations.test.ts | 36 +++++ .../tests/snapshot-policy-export.test-d.ts | 4 + packages/ai/src/middlewares/index.ts | 2 - packages/ai/tests/middlewares/index.test.ts | 17 --- 15 files changed, 333 insertions(+), 138 deletions(-) delete mode 100644 .changeset/quiet-sandbox-middleware.md delete mode 100644 packages/ai-sandbox/tests/ai-middleware-subpath.test.ts delete mode 100644 packages/ai/tests/middlewares/index.test.ts diff --git a/.changeset/quiet-sandbox-middleware.md b/.changeset/quiet-sandbox-middleware.md deleted file mode 100644 index 73f6c75138..0000000000 --- a/.changeset/quiet-sandbox-middleware.md +++ /dev/null @@ -1,5 +0,0 @@ ---- -'@tanstack/ai': patch ---- - -Export `CapabilityRegistry` from `@tanstack/ai/middlewares`. diff --git a/docs/config.json b/docs/config.json index 9eb6d68775..7d7243472b 100644 --- a/docs/config.json +++ b/docs/config.json @@ -299,7 +299,7 @@ "label": "Build a Generation Adapter", "to": "persistence/build-your-own-generation-adapter", "addedAt": "2026-08-04", - "updatedAt": "2026-08-13" + "updatedAt": "2026-08-14" }, { "label": "Build a Sandbox Adapter", @@ -310,7 +310,7 @@ "label": "Store Reference", "to": "persistence/store-reference", "addedAt": "2026-08-04", - "updatedAt": "2026-08-13" + "updatedAt": "2026-08-14" }, { "label": "How Persistence Works", @@ -511,7 +511,7 @@ "label": "Overview", "to": "sandbox/overview", "addedAt": "2026-06-16", - "updatedAt": "2026-08-13" + "updatedAt": "2026-08-14" }, { "label": "Quick Start", @@ -523,7 +523,7 @@ "label": "Providers", "to": "sandbox/providers", "addedAt": "2026-06-29", - "updatedAt": "2026-08-13" + "updatedAt": "2026-08-14" }, { "label": "Harnesses", @@ -553,18 +553,18 @@ "label": "Lifecycle & Snapshots", "to": "sandbox/lifecycle", "addedAt": "2026-06-29", - "updatedAt": "2026-08-13" + "updatedAt": "2026-08-14" }, { "label": "Portable Sandbox Snapshots", "to": "sandbox/portable-snapshots", - "addedAt": "2026-08-13" + "addedAt": "2026-08-14" }, { "label": "Instance Durability", "to": "sandbox/durability", "addedAt": "2026-08-04", - "updatedAt": "2026-08-13" + "updatedAt": "2026-08-14" }, { "label": "Durable Runs", diff --git a/docs/sandbox/portable-snapshots.md b/docs/sandbox/portable-snapshots.md index bcf62879de..365aa18bc0 100644 --- a/docs/sandbox/portable-snapshots.md +++ b/docs/sandbox/portable-snapshots.md @@ -5,8 +5,8 @@ order: 10 description: "Store a completed sandbox workspace as durable files, artifacts, and conversation data, then rebuild it in a new sandbox." --- -An agent can finish work in a sandbox, then the sandbox can disappear. You can -need the same files when the page reloads or when a later run starts. Portable +An agent can finish work in a sandbox, then the sandbox can disappear. You need +the same files when the page reloads or when a later run starts. Portable sandbox snapshots save the completed workspace in your persistence stores. This feature saves a checkpoint after each successful terminal run. A later @@ -92,6 +92,7 @@ that the owner can access the thread before you call the helper. ```ts import { saveNamedSandboxSnapshot } from '@tanstack/ai-sandbox' +import { requireSession } from './auth' import { sandbox, snapshots, instances } from './sandbox-server' export async function POST(request: Request) { @@ -153,6 +154,7 @@ reference counts. It must reject a non-empty destination thread. ```ts import { forkFromSandboxSnapshot } from '@tanstack/ai-sandbox' +import { requireSession } from './auth' import { snapshots } from './sandbox-server' export async function POST(request: Request) { @@ -217,6 +219,7 @@ and `Uint8Array` bytes. It does not create an HTTP response. ```ts import { resolveSnapshotArtifact } from '@tanstack/ai-sandbox' +import { requireSession } from './auth' import { snapshots } from './sandbox-server' export async function GET(request: Request) { @@ -239,7 +242,7 @@ export async function GET(request: Request) { artifactId, snapshots, }) - return new Response(bytes, { + return new Response(bytes.slice(), { headers: { 'content-type': artifact.mimeType, 'content-length': String(artifact.size), @@ -302,12 +305,37 @@ Portable snapshots support regular files and directories only. A capture or restore fails safely when it finds a symlink, an executable file, or a special filesystem entry. -The default policy excludes these paths at every depth: +The default policy excludes these path segments at every depth: - `.git` - `node_modules` - `.env*` -- The workspace projection marker, `.tanstack-projected-` + +It also excludes these exact paths: + +- The projection marker, `.tanstack-projected-`, at the + workspace root only. +- `CLAUDE.md` and `GEMINI.md` at the workspace root. +- Direct `.claude/skills/`, `.codex/skills/`, and + `.grok/skills/` paths. + +These exclusions use paths even when an entry is a regular file or a copied +file. A custom policy replaces the default exclusions. If you pass only +`redact` or `include`, `.env`, `.git`, and `node_modules` are captured unless +you copy the default policy first: + +```ts +import { defaultSandboxSnapshotPolicy } from '@tanstack/ai-sandbox' + +const policy = { + ...defaultSandboxSnapshotPolicy(), + redact({ bytes }: { bytes: Uint8Array }) { + return bytes + }, +} +``` + +The exact projection marker remains protected for the workspace. Resolved secret values are replaced with zero bytes before their content is hashed or stored. A custom policy cannot capture or restore the exact projection diff --git a/examples/ts-react-chat/src/lib/sqlite-persistence.test.ts b/examples/ts-react-chat/src/lib/sqlite-persistence.test.ts index 299f767bfb..f4e30a188d 100644 --- a/examples/ts-react-chat/src/lib/sqlite-persistence.test.ts +++ b/examples/ts-react-chat/src/lib/sqlite-persistence.test.ts @@ -242,7 +242,7 @@ describe('sqliteSandboxSnapshots fork transaction', () => { } }) - it.each([ + it.each<['files' | 'artifacts', unknown]>([ ['files', null], ['files', 'not-an-array'], ['files', [null]], @@ -278,6 +278,54 @@ describe('sqliteSandboxSnapshots fork transaction', () => { } }) + it('rejects empty and non-string identifiers without changing state', async () => { + const snapshots = sqliteSandboxSnapshots({ url: ':memory:', migrate: true }) + try { + const writer = await snapshots.checkpoints.acquireWriter('thread') + const before = { + head: await snapshots.checkpoints.getHead('thread'), + list: await snapshots.checkpoints.list('thread'), + references: await snapshots.checkpoints.listBlobReferences(), + } + await expect( + Reflect.apply( + snapshots.checkpoints.acquireWriter, + snapshots.checkpoints, + [''], + ), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ID' }) + await expect( + Reflect.apply(snapshots.checkpoints.list, snapshots.checkpoints, [{}]), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ID' }) + await expect( + Reflect.apply(snapshots.checkpoints.append, snapshots.checkpoints, [ + { + checkpoint: { + id: 'invalid', + threadId: 7, + parentCheckpointId: null, + createdAt: 1, + reason: 'named', + files: [], + conversation: [], + artifacts: [], + }, + expectedHeadId: null, + writer, + }, + ]), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_INVALID_ID' }) + expect(await snapshots.checkpoints.getHead('thread')).toBe(before.head) + expect(await snapshots.checkpoints.list('thread')).toEqual(before.list) + expect(await snapshots.checkpoints.listBlobReferences()).toEqual( + before.references, + ) + await writer.release() + } finally { + snapshots.close() + } + }) + it('rolls back the destination when checkpoint storage fails after transcript staging', async () => { const dir = mkdtempSync(join(tmpdir(), 'tanstack-sqlite-fork-')) const file = join(dir, 'snapshots.db') @@ -446,3 +494,56 @@ describe('sqlitePersistence migrate — an existing pre-durability database', () } }) }) + +describe('sqliteSandboxSnapshots migrate — an existing artifacts table', () => { + it('adds the artifact thread-order index and uses it for listForThread', () => { + const dir = mkdtempSync(join(tmpdir(), 'tanstack-sqlite-artifacts-')) + const file = join(dir, 'old.db') + try { + const old = new DatabaseSync(file) + old.exec(` + CREATE TABLE artifacts ( + artifact_id text PRIMARY KEY NOT NULL, + run_id text NOT NULL, + thread_id text NOT NULL, + blob_key text, + name text NOT NULL, + mime_type text NOT NULL, + size integer NOT NULL, + source_url text, + created_at integer NOT NULL + ); + CREATE INDEX artifacts_run_order + ON artifacts (run_id, created_at ASC, artifact_id ASC); + `) + old.close() + + const snapshots = sqliteSandboxSnapshots({ url: file, migrate: true }) + snapshots.close() + + const inspector = new DatabaseSync(file) + try { + const indexes = inspector + .prepare("SELECT name FROM pragma_index_list('artifacts')") + .all() + expect(indexes).toContainEqual( + expect.objectContaining({ name: 'artifacts_thread_order' }), + ) + + const plan = inspector + .prepare( + 'EXPLAIN QUERY PLAN SELECT * FROM artifacts WHERE thread_id = ? ORDER BY created_at ASC, artifact_id ASC', + ) + .all('thread') + expect(JSON.stringify(plan)).toContain( + 'SEARCH artifacts USING INDEX artifacts_thread_order (thread_id=?)', + ) + expect(JSON.stringify(plan)).not.toContain('USE TEMP B-TREE') + } finally { + inspector.close() + } + } finally { + rmSync(dir, { recursive: true, force: true }) + } + }) +}) diff --git a/examples/ts-react-chat/src/lib/sqlite-persistence.ts b/examples/ts-react-chat/src/lib/sqlite-persistence.ts index 6d184d75d9..425ee0bf04 100644 --- a/examples/ts-react-chat/src/lib/sqlite-persistence.ts +++ b/examples/ts-react-chat/src/lib/sqlite-persistence.ts @@ -154,6 +154,8 @@ CREATE TABLE IF NOT EXISTS artifacts ( ); CREATE INDEX IF NOT EXISTS artifacts_run_order ON artifacts (run_id, created_at ASC, artifact_id ASC); +CREATE INDEX IF NOT EXISTS artifacts_thread_order + ON artifacts (thread_id, created_at ASC, artifact_id ASC); -- The bytes themselves. \`body\` is a BLOB column, so this file IS the object -- store; a production adapter would keep metadata here and put bytes in S3/R2. CREATE TABLE IF NOT EXISTS blobs ( @@ -1176,14 +1178,30 @@ function hasUnpairedSurrogate(value: string): boolean { return false } -function assertCheckpointId(value: string, label: string): void { - if (!value || value.includes('\0') || hasUnpairedSurrogate(value)) { +function assertCheckpointId( + value: unknown, + label: string, +): asserts value is string { + if ( + typeof value !== 'string' || + value.length === 0 || + value.includes('\0') || + hasUnpairedSurrogate(value) + ) { throw new SandboxCheckpointInvalidIdError( `${label} must be a non-empty well-formed Unicode string`, ) } } +function rollbackIfActive(db: DatabaseSync): void { + try { + db.exec('ROLLBACK') + } catch { + // SQLite can already abort the transaction. Keep the original error. + } +} + function assertCheckpoint(checkpoint: SandboxCheckpoint): void { assertCheckpointId(checkpoint.id, 'Checkpoint id') assertCheckpointId(checkpoint.threadId, 'Checkpoint thread id') @@ -1201,9 +1219,11 @@ function assertCheckpoint(checkpoint: SandboxCheckpoint): void { throw new SandboxCheckpointInvalidEntryError( 'Checkpoint artifacts must be an array', ) + const files: ReadonlyArray = checkpoint.files + const artifacts: ReadonlyArray = checkpoint.artifacts const paths = new Set() const kinds = new Map() - for (const entry of checkpoint.files) { + for (const entry of files) { if (entry === null || typeof entry !== 'object') throw new SandboxCheckpointInvalidEntryError( 'Checkpoint entry must be an object', @@ -1263,7 +1283,7 @@ function assertCheckpoint(checkpoint: SandboxCheckpoint): void { } kinds.set(entry.path, entry.kind) } - for (const artifact of checkpoint.artifacts) { + for (const artifact of artifacts) { if (artifact === null || typeof artifact !== 'object') throw new SandboxCheckpointInvalidEntryError( 'Checkpoint artifact must be an object', @@ -1506,7 +1526,7 @@ function createCheckpointStore( incrementReferences(checkpoint) db.exec('COMMIT') } catch (error) { - db.exec('ROLLBACK') + rollbackIfActive(db) throw error } return { headId: checkpoint.id } @@ -1560,7 +1580,7 @@ function createCheckpointStore( ) db.exec('COMMIT') } catch (error) { - db.exec('ROLLBACK') + rollbackIfActive(db) throw error } }, @@ -1597,7 +1617,7 @@ function createCheckpointStore( ).run(threadId, lease.ownerToken, fence, lease.expiresAt) db.exec('COMMIT') } catch (error) { - db.exec('ROLLBACK') + rollbackIfActive(db) throw error } return { @@ -1623,7 +1643,7 @@ function createCheckpointStore( db.exec('COMMIT') return { expiresAt } } catch (error) { - db.exec('ROLLBACK') + rollbackIfActive(db) throw error } }, @@ -1737,7 +1757,7 @@ function createCheckpointStore( db.exec('COMMIT') return { checkpoint: cloneCheckpoint(checkpoint) } } catch (error) { - db.exec('ROLLBACK') + rollbackIfActive(db) throw error } }, diff --git a/packages/ai-persistence/tests/persistence-completion.test.ts b/packages/ai-persistence/tests/persistence-completion.test.ts index 4b9d1eaa3d..f15279b59b 100644 --- a/packages/ai-persistence/tests/persistence-completion.test.ts +++ b/packages/ai-persistence/tests/persistence-completion.test.ts @@ -258,10 +258,13 @@ describe('PersistenceCompletionCapability', () => { const persistence = memoryPersistence() const failure = new Error('final save failed') let saves = 0 + const originalSave = persistence.stores.messages.saveThread.bind( + persistence.stores.messages, + ) persistence.stores.messages.saveThread = async (...args) => { saves += 1 if (saves === 2) throw failure - await memoryPersistence().stores.messages.saveThread(...args) + await originalSave(...args) } const { completion, setupReady, run } = createCompletionRun(persistence, { adapter: mockAdapter(finishedChunks()), diff --git a/packages/ai-sandbox/skills/ai-sandbox/SKILL.md b/packages/ai-sandbox/skills/ai-sandbox/SKILL.md index 00946b91d6..f69b9d11b9 100644 --- a/packages/ai-sandbox/skills/ai-sandbox/SKILL.md +++ b/packages/ai-sandbox/skills/ai-sandbox/SKILL.md @@ -217,12 +217,18 @@ const middleware = [ Each successful terminal run saves regular files, empty directories, durable conversation data, and persisted thread artifacts. A later run restores the latest checkpoint only into a new private sandbox. A live resumed sandbox is -never overwritten. The default policy excludes `.git`, `node_modules`, `.env*`, -and the workspace projection marker. Resolved secrets are redacted before the -data is stored. Symlinks, executables, and special filesystem entries fail the -capture or restore. Each thread has one writer lease. Pause and detach release -the lease without a partial checkpoint. Blob retention is manual because there -is no automatic garbage collection yet. +never overwritten. The default policy excludes `.git`, `node_modules`, and +`.env*` path segments at every depth. It excludes the exact projection marker +only at the workspace root. It also excludes root `CLAUDE.md` and `GEMINI.md`, +plus direct `.claude/skills/`, `.codex/skills/`, and +`.grok/skills/` paths. These exclusions use paths even for regular files +or copies. A custom policy replaces them, except for exact projection-marker +protection. If you only pass `redact` or `include`, copy +`defaultSandboxSnapshotPolicy()` first or `.env` files are captured. Resolved secrets are redacted before the data is stored. Symlinks, +executables, and special filesystem entries fail the capture or restore. Each +thread has one writer lease. Pause and detach release the lease without a +partial checkpoint. Blob retention is manual because there is no automatic +garbage collection yet. Read `docs/sandbox/portable-snapshots.md` for the full server-only setup and the restore safety rules. @@ -248,10 +254,14 @@ including its copied conversation. A partial transaction breaks snapshot consistency. Snapshot capture supports regular files and empty directories only. It excludes -`.git`, `node_modules`, `.env*`, and the workspace projection marker. It rejects -symlinks, executable files, and special filesystem entries. Restore verifies -the manifest and blobs before it changes a new private sandbox. It never writes -into a live resumed sandbox. +`.git`, `node_modules`, and `.env*` path segments at every depth. It excludes +the exact projection marker only at the workspace root. It also excludes root +`CLAUDE.md` and `GEMINI.md`, plus direct `.claude/skills/`, +`.codex/skills/`, and `.grok/skills/` paths. These exclusions use +paths even for regular files or copies. A custom policy replaces them, except +for exact projection-marker protection. It rejects symlinks, executable files, +and special filesystem entries. Restore verifies the manifest and blobs before +it changes a new private sandbox. It never writes into a live resumed sandbox. ## Providers @@ -1051,14 +1061,15 @@ including the client `joinRun` side, is in `docs/sandbox/takeover.md`. - **Harness adapters require a sandbox.** Always include `withSandbox(...)` in `middleware` — without it `chat()` throws a missing-capability error. -- **Secrets** (`workspace.secrets`) are injected into the sandbox env and never - persisted (no snapshots, no sandbox store, no event log). Always create them - with `createSecrets(...)` so the values stay hidden behind `SecretRef` tokens. - The agent binary (`claude`) must exist in the sandbox image (install it in - `setup` or bake it into the image). +- **Secrets** (`workspace.secrets`) are injected into the sandbox env. Their + raw values are never persisted in snapshots, the sandbox store, or the event + log. Always create them with `createSecrets(...)` so the values stay hidden + behind `SecretRef` tokens. The agent binary (`claude`) must exist in the + sandbox image (install it in `setup` or bake it into the image). - **Secret-bearing projected files** (e.g. MCP config with resolved header - values) are re-written on every projection call so rotated secrets re-apply; - they are never included in a snapshot. + values) can be included by default capture. Capture replaces resolved secret + bytes with zero bytes before it hashes or writes snapshot blobs. Restore runs + before projection, so projection writes current secret values after restore. - **chat()-provided `tools` are bridged** into the in-sandbox agent over a host-side MCP tool-proxy: the agent calls them as `mcp__tanstack__` and each call is proxied back to the host where the tool's `execute()` runs (with diff --git a/packages/ai-sandbox/src/index.ts b/packages/ai-sandbox/src/index.ts index 79050723d0..b0f3ecd7c2 100644 --- a/packages/ai-sandbox/src/index.ts +++ b/packages/ai-sandbox/src/index.ts @@ -60,7 +60,7 @@ export type { } from './checkpoint-store' // File snapshot policy used by provider snapshot create/restore inputs. -export { SandboxSnapshotError } from './snapshots' +export { SandboxSnapshotError, defaultSandboxSnapshotPolicy } from './snapshots' export type { SandboxSnapshotErrorCode, SandboxSnapshotPolicy, diff --git a/packages/ai-sandbox/src/middleware.ts b/packages/ai-sandbox/src/middleware.ts index ee0a0c28d5..4495d3f090 100644 --- a/packages/ai-sandbox/src/middleware.ts +++ b/packages/ai-sandbox/src/middleware.ts @@ -1017,76 +1017,78 @@ export function withSandbox( let primaryError: unknown try { - const snapshotCaptureTask = (async (): Promise => { - 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 + const snapshotCaptureTask = Promise.resolve().then( + async (): Promise => { + 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 + 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 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 ?? DEFAULT_WORKSPACE_ROOT, - }, - state.snapshotPolicy, - definition.workspace?.secrets !== undefined - ? resolveAllSecrets(definition.workspace.secrets) - : {}, - ) - if (!canPublishPortableSnapshot(state, lease)) return + const files = await captureSandboxFiles( + handle, + { + blobs: config.persistence.stores.blobs, + workspaceRoot: + definition.workspace?.root ?? DEFAULT_WORKSPACE_ROOT, + }, + state.snapshotPolicy, + definition.workspace?.secrets !== undefined + ? 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 !== undefined - ? 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 !== undefined + ? resolveAllSecrets(definition.workspace.secrets) + : {}, + ) + if (!canPublishPortableSnapshot(state, lease)) return - const parentCheckpointId = await config.checkpoints.getHead( - ctx.threadId, - ) - 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) { + 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 - throw error - } - canPublishPortableSnapshot(state, lease) - })() + }, + ) state.snapshotCaptureTask = snapshotCaptureTask try { await snapshotCaptureTask diff --git a/packages/ai-sandbox/tests/ai-middleware-subpath.test.ts b/packages/ai-sandbox/tests/ai-middleware-subpath.test.ts deleted file mode 100644 index 7e1b45eaf5..0000000000 --- a/packages/ai-sandbox/tests/ai-middleware-subpath.test.ts +++ /dev/null @@ -1,6 +0,0 @@ -import { expect, it } from 'vitest' -import { CapabilityRegistry } from '@tanstack/ai/middlewares' - -it('exports a working CapabilityRegistry from the built middlewares subpath', () => { - expect(new CapabilityRegistry()).toBeInstanceOf(CapabilityRegistry) -}) diff --git a/packages/ai-sandbox/tests/fakes.ts b/packages/ai-sandbox/tests/fakes.ts index 249fec95b2..8bdaa9d383 100644 --- a/packages/ai-sandbox/tests/fakes.ts +++ b/packages/ai-sandbox/tests/fakes.ts @@ -1,5 +1,4 @@ import { resolveDebugOption } from '@tanstack/ai/adapter-internals' -import { CapabilityRegistry } from '@tanstack/ai/middlewares' import { makeFakeShellSpawn } from '../src/testkit/shell-spawn' import type { InternalLogger } from '@tanstack/ai/adapter-internals' import type { @@ -217,9 +216,28 @@ export function captureLogger(): { * asserting on a provided capability is not exercising a stub that always * answers the same way. * - * The capabilities field uses the public `CapabilityRegistry` export, so this - * fake exercises the same registry implementation as production code. + * The capabilities field uses this test-local registry. It has the same + * behavior that capability accessors need, without making internal middleware + * bookkeeping a public API. */ +class TestCapabilityRegistry { + private readonly provided = new Set() + private onDuplicate?: (name: string) => void + + setOnDuplicate(callback: (name: string) => void): void { + this.onDuplicate = callback + } + + markProvided(handle: { capabilityName: string }): void { + if (this.provided.has(handle)) this.onDuplicate?.(handle.capabilityName) + this.provided.add(handle) + } + + has(handle: object): boolean { + return this.provided.has(handle) + } +} + export function makeMiddlewareCtx(input: { threadId: string runId: string @@ -254,7 +272,9 @@ export function makeMiddlewareCtx(input: { messages: [], createId: (prefix: string) => `${prefix}-${Math.random().toString(36).slice(2)}`, - capabilities: new CapabilityRegistry(), + capabilities: + // @ts-expect-error This test-only registry has the required methods, but the production class has private state and is nominally typed. + new TestCapabilityRegistry() as ChatMiddlewareContext['capabilities'], get: (capability) => capability[0](ctx), getOptional: (capability) => capability[0](ctx, { optional: true }), provide: (capability, value) => capability[1](ctx, value), diff --git a/packages/ai-sandbox/tests/snapshot-operations.test.ts b/packages/ai-sandbox/tests/snapshot-operations.test.ts index 09db241f71..9e16920e70 100644 --- a/packages/ai-sandbox/tests/snapshot-operations.test.ts +++ b/packages/ai-sandbox/tests/snapshot-operations.test.ts @@ -999,6 +999,42 @@ describe('public sandbox snapshot operations', () => { expect(listCalls).not.toContain(`/workspace/${marker}`) }) + it('uses the computed workspace hash instead of a stale caller hash', async () => { + const workspace: WorkspaceDefinition = { + source: { type: 'none' }, + root: '/custom', + } + const workspaceHash = computeWorkspaceHash(workspace) + const computedMarker = `.tanstack-projected-${workspaceHash}` + const staleMarker = '.tanstack-projected-caller-hash' + const listCalls: Array = [] + const fixture = await namedFixture({ + workspace, + policy: { workspaceHash: 'caller-hash', exclude: () => false }, + }) + resumeWithHandle( + fixture, + workspaceHandle({ + root: workspace.root ?? '/workspace', + listCalls, + files: [ + { path: 'kept.txt', content: 'kept' }, + { path: `${computedMarker}/private.txt`, content: 'private' }, + { path: `${staleMarker}/captured.txt`, content: 'captured' }, + ], + }), + ) + + const saved = await namedSave(fixture) + + expect(saved.files.map((entry) => entry.path).sort()).toEqual([ + `${staleMarker}/captured.txt`, + 'kept.txt', + ]) + expect(listCalls).not.toContain(`/custom/${computedMarker}`) + expect(listCalls).toContain(`/custom/${staleMarker}`) + }) + it('preserves custom redaction and passes resolved workspace secrets', async () => { const seenSecrets: Array = [] const workspace: WorkspaceDefinition = { diff --git a/packages/ai-sandbox/tests/snapshot-policy-export.test-d.ts b/packages/ai-sandbox/tests/snapshot-policy-export.test-d.ts index d2086b6cfa..f2223b0426 100644 --- a/packages/ai-sandbox/tests/snapshot-policy-export.test-d.ts +++ b/packages/ai-sandbox/tests/snapshot-policy-export.test-d.ts @@ -1,4 +1,5 @@ import { expectTypeOf } from 'vitest' +import { defaultSandboxSnapshotPolicy } from '../src' import type { SandboxCheckpointStoreOptions, SandboxSnapshotPolicy, @@ -26,3 +27,6 @@ const policy: SandboxSnapshotPolicy = { } expectTypeOf(policy).toMatchTypeOf() +expectTypeOf( + defaultSandboxSnapshotPolicy, +).returns.toMatchTypeOf() diff --git a/packages/ai/src/middlewares/index.ts b/packages/ai/src/middlewares/index.ts index d3fe58ecc9..1470cdcff1 100644 --- a/packages/ai/src/middlewares/index.ts +++ b/packages/ai/src/middlewares/index.ts @@ -12,8 +12,6 @@ export { type ContentFilteredInfo, } from './content-guard' -export { CapabilityRegistry } from '../activities/chat/middleware/capabilities' - // otelMiddleware is exported from the dedicated subpath // `@tanstack/ai/middlewares/otel` so that importing the main middlewares barrel // does not eagerly require `@opentelemetry/api` (which is an optional peer diff --git a/packages/ai/tests/middlewares/index.test.ts b/packages/ai/tests/middlewares/index.test.ts deleted file mode 100644 index 367bdf0e34..0000000000 --- a/packages/ai/tests/middlewares/index.test.ts +++ /dev/null @@ -1,17 +0,0 @@ -import { describe, expect, it } from 'vitest' -import { CapabilityRegistry } from '../../src/middlewares' -import { createCapability } from '../../src/activities/chat/middleware/capabilities' - -describe('middlewares public exports', () => { - it('exports a working CapabilityRegistry', () => { - const registry = new CapabilityRegistry() - const capability = createCapability()('value') - const [getValue, provideValue] = capability - const context = { capabilities: registry } - - expect(registry.has(capability)).toBe(false) - provideValue(context, 42) - expect(registry.has(capability)).toBe(true) - expect(getValue(context)).toBe(42) - }) -}) From b62ae4689a2f09fa9a5205678ee1805b51763a5f Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Fri, 14 Aug 2026 14:58:15 +0200 Subject: [PATCH 03/12] fix(e2e): accept createdAt on snapshot conversation messages After merging message-metadata preservation, automatic assistant turns keep createdAt. The portable snapshot spec compared the JSON body exactly and failed. --- testing/e2e/tests/sandbox-file-persistence.spec.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/testing/e2e/tests/sandbox-file-persistence.spec.ts b/testing/e2e/tests/sandbox-file-persistence.spec.ts index eeadccf0bf..6828bf8446 100644 --- a/testing/e2e/tests/sandbox-file-persistence.spec.ts +++ b/testing/e2e/tests/sandbox-file-persistence.spec.ts @@ -43,6 +43,7 @@ test.describe('sandbox portable file snapshots', () => { { content: 'recover', role: 'user' }, { content: 'automatic conversation', + createdAt: expect.any(String), id: 'automatic-message', role: 'assistant', }, From 2d81a48aa510ce0f4c4025bcf4d43938c66b9b10 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Fri, 14 Aug 2026 16:05:31 +0200 Subject: [PATCH 04/12] feat(ai-sandbox): fuse snapshot helpers into one object createSandboxSnapshots and memorySandboxSnapshots now return save, fork, and readArtifact on the same object. --- .changeset/fuzzy-snapshots-build.md | 2 +- docs/config.json | 3 +- docs/sandbox/portable-snapshots.md | 81 ++++--- .../src/lib/sqlite-persistence.ts | 10 +- packages/ai-sandbox/README.md | 9 +- .../ai-sandbox/skills/ai-sandbox/SKILL.md | 34 ++- packages/ai-sandbox/src/index.ts | 18 +- packages/ai-sandbox/src/memory-snapshots.ts | 37 +++- .../ai-sandbox/src/snapshot-operations.ts | 175 +++++++++++++-- packages/ai-sandbox/src/snapshots.ts | 3 + .../memory-snapshots-declaration.test-d.ts | 3 + .../tests/memory-snapshots-import.test.ts | 3 + .../tests/snapshot-operations.test-d.ts | 6 + .../tests/snapshot-operations.test.ts | 199 ++++++++++++------ .../routes/api.sandbox-file-persistence.ts | 23 +- 15 files changed, 426 insertions(+), 180 deletions(-) diff --git a/.changeset/fuzzy-snapshots-build.md b/.changeset/fuzzy-snapshots-build.md index b8c6c25f13..293776afbd 100644 --- a/.changeset/fuzzy-snapshots-build.md +++ b/.changeset/fuzzy-snapshots-build.md @@ -2,4 +2,4 @@ '@tanstack/ai-sandbox': minor --- -Add portable sandbox checkpoints, named snapshot saves, selected-checkpoint forks, and checkpoint artifact reads. +Add portable sandbox checkpoints. `createSandboxSnapshots` and `memorySandboxSnapshots` return one object with `save`, `fork`, and `readArtifact`. diff --git a/docs/config.json b/docs/config.json index 7d7243472b..25a17b8d2e 100644 --- a/docs/config.json +++ b/docs/config.json @@ -558,7 +558,8 @@ { "label": "Portable Sandbox Snapshots", "to": "sandbox/portable-snapshots", - "addedAt": "2026-08-14" + "addedAt": "2026-08-14", + "updatedAt": "2026-08-14" }, { "label": "Instance Durability", diff --git a/docs/sandbox/portable-snapshots.md b/docs/sandbox/portable-snapshots.md index 365aa18bc0..14eeefda3a 100644 --- a/docs/sandbox/portable-snapshots.md +++ b/docs/sandbox/portable-snapshots.md @@ -35,7 +35,6 @@ import { } from '@tanstack/ai-sandbox' import { dockerSandbox } from '@tanstack/ai-sandbox-docker' -const snapshots = await memorySandboxSnapshots() const instances = new InMemorySandboxInstanceStore() const userId = 'user-123' // Read this from the server session. @@ -46,6 +45,11 @@ const sandbox = defineSandbox({ lifecycle: { reuse: 'thread' }, }) +const snapshots = await memorySandboxSnapshots({ + sandbox, + instances, +}) + const result = chat({ threadId: 'app-thread', context: { userId }, @@ -55,10 +59,7 @@ const result = chat({ withPersistence(snapshots.persistence), withSandbox(sandbox, { instances, - snapshots: { - persistence: snapshots.persistence, - checkpoints: snapshots.checkpoints, - }, + snapshots, }), ], }) @@ -66,12 +67,31 @@ const result = chat({ void result ``` -`memorySandboxSnapshots()` creates an in-memory persistence object directly. It -does not load `@tanstack/ai-persistence` at runtime. +`memorySandboxSnapshots()` creates the in-memory persistence, checkpoint store, +and snapshot methods as one object. It does not load `@tanstack/ai-persistence` +at runtime. + +You can bind `sandbox`, `instances`, `tenant`, and `locks` at create time. A +named save can override those values. + +If you already have a persistence object, pass that same object to +`createSandboxSnapshots`. Use that same object in `withPersistence`. + +```ts +import { createSandboxSnapshots } from '@tanstack/ai-sandbox' +import { checkpoints, instances, persistence, sandbox } from './sandbox-server' + +const snapshots = createSandboxSnapshots({ + persistence, + checkpoints, + sandbox, + instances, +}) +``` Keep `instances` in the same server module as this middleware. A named save must use this same store. Pass the session `userId` in `context` for every run. -Pass that same user id as `tenant.userId` to `saveNamedSandboxSnapshot`. +Pass that same user id as `tenant.userId` on `snapshots.save`. Use a durable persistence implementation and checkpoint store in production. The memory factory is useful for local development and examples only. @@ -83,7 +103,7 @@ default policy. ## Save a named checkpoint Automatic saves protect each completed run. Use a named save when a user marks -one workspace state, such as a release candidate. The helper requires a live, +one workspace state, such as a release candidate. The method requires a live, reusable sandbox for the thread. A lifecycle with `reuse: 'none'` cannot create a named checkpoint. @@ -91,9 +111,8 @@ Keep this route on the server. Derive the owner from the session. Then make sure that the owner can access the thread before you call the helper. ```ts -import { saveNamedSandboxSnapshot } from '@tanstack/ai-sandbox' import { requireSession } from './auth' -import { sandbox, snapshots, instances } from './sandbox-server' +import { snapshots } from './sandbox-server' export async function POST(request: Request) { const session = await requireSession(request) @@ -110,12 +129,9 @@ export async function POST(request: Request) { return new Response('Not found', { status: 404 }) } - const checkpoint = await saveNamedSandboxSnapshot({ - definition: sandbox, + const checkpoint = await snapshots.save({ threadId, runId, - instances, - snapshots, label, tenant: { userId: session.userId }, }) @@ -153,34 +169,31 @@ atomically copy the selected checkpoint, the conversation, the head, and blob reference counts. It must reject a non-empty destination thread. ```ts -import { forkFromSandboxSnapshot } from '@tanstack/ai-sandbox' import { requireSession } from './auth' import { snapshots } from './sandbox-server' export async function POST(request: Request) { const session = await requireSession(request) - const { sourceThreadId, sourceCheckpointId, destinationThreadId } = - await request.json() + const { threadId, checkpointId, destinationThreadId } = await request.json() if ( - typeof sourceThreadId !== 'string' || - typeof sourceCheckpointId !== 'string' || + typeof threadId !== 'string' || + typeof checkpointId !== 'string' || typeof destinationThreadId !== 'string' ) { return new Response('Invalid request', { status: 400 }) } - if (!(await session.canAccessThread(sourceThreadId))) { + if (!(await session.canAccessThread(threadId))) { return new Response('Not found', { status: 404 }) } if (!(await session.canCreateThread(destinationThreadId))) { return new Response('Not found', { status: 404 }) } - const checkpoint = await forkFromSandboxSnapshot({ - sourceThreadId, - sourceCheckpointId, + const checkpoint = await snapshots.fork({ + threadId, + checkpointId, destinationThreadId, - snapshots, }) return Response.json({ checkpointId: checkpoint.id }) @@ -192,16 +205,16 @@ threads. Do not accept a client-selected checkpoint as proof of access. ```ts export async function forkCheckpoint( - sourceThreadId: string, - sourceCheckpointId: string, + threadId: string, + checkpointId: string, destinationThreadId: string, ) { const response = await fetch('/api/snapshots/fork', { method: 'POST', headers: { 'content-type': 'application/json' }, body: JSON.stringify({ - sourceThreadId, - sourceCheckpointId, + threadId, + checkpointId, destinationThreadId, }), }) @@ -212,13 +225,12 @@ export async function forkCheckpoint( ## Read a snapshot artifact -`resolveSnapshotArtifact` reads copied artifact bytes from one checkpoint. It +`snapshots.readArtifact` reads copied artifact bytes from one checkpoint. It checks that the checkpoint belongs to the supplied thread. Your route must still -authorize that thread before it calls the helper. The helper returns metadata +authorize that thread before it calls the method. The method returns metadata and `Uint8Array` bytes. It does not create an HTTP response. ```ts -import { resolveSnapshotArtifact } from '@tanstack/ai-sandbox' import { requireSession } from './auth' import { snapshots } from './sandbox-server' @@ -236,11 +248,10 @@ export async function GET(request: Request) { return new Response('Not found', { status: 404 }) } - const { artifact, bytes } = await resolveSnapshotArtifact({ + const { artifact, bytes } = await snapshots.readArtifact({ threadId, checkpointId, artifactId, - snapshots, }) return new Response(bytes.slice(), { headers: { @@ -252,7 +263,7 @@ export async function GET(request: Request) { ``` The client can use the authorized route as an artifact URL. It must not read the -blob store or call `resolveSnapshotArtifact` in the browser. +blob store or call `snapshots.readArtifact` in the browser. ```ts export function snapshotArtifactUrl( diff --git a/examples/ts-react-chat/src/lib/sqlite-persistence.ts b/examples/ts-react-chat/src/lib/sqlite-persistence.ts index 425ee0bf04..fc8958f435 100644 --- a/examples/ts-react-chat/src/lib/sqlite-persistence.ts +++ b/examples/ts-react-chat/src/lib/sqlite-persistence.ts @@ -60,6 +60,7 @@ import type { RunStore, } from '@tanstack/ai-persistence' import { + createSandboxSnapshots, SandboxCheckpointConflictError, SandboxCheckpointDuplicateIdError, SandboxCheckpointError, @@ -77,6 +78,7 @@ import type { SandboxCheckpointWriter, SandboxSnapshotArtifact, SandboxSnapshotEntry, + SandboxSnapshots, } from '@tanstack/ai-sandbox' // --------------------------------------------------------------------------- @@ -1836,9 +1838,7 @@ export function sqlitePersistence( /** Build the seven persistence stores and a durable SQLite checkpoint store. */ export function sqliteSandboxSnapshots( options: SqlitePersistenceOptions & SandboxCheckpointStoreOptions, -): { - persistence: SqliteAIPersistence - checkpoints: ForkCapableSandboxCheckpointStore +): SandboxSnapshots & { close: () => void } { const filename = normalizeSqliteUrl(options.url) @@ -1862,10 +1862,10 @@ export function sqliteSandboxSnapshots( }, }) const checkpoints = createCheckpointStore(db, options) + const snapshots = createSandboxSnapshots({ persistence, checkpoints }) let closed = false return { - persistence, - checkpoints, + ...snapshots, close() { if (closed) return db.close() diff --git a/packages/ai-sandbox/README.md b/packages/ai-sandbox/README.md index 48e8e5d257..f5ff37f9ac 100644 --- a/packages/ai-sandbox/README.md +++ b/packages/ai-sandbox/README.md @@ -147,16 +147,11 @@ persistence value and a checkpoint store to `withSandbox`. import { withPersistence } from '@tanstack/ai-persistence' import { memorySandboxSnapshots, withSandbox } from '@tanstack/ai-sandbox' -const snapshots = await memorySandboxSnapshots() +const snapshots = await memorySandboxSnapshots({ sandbox, instances }) const middleware = [ withPersistence(snapshots.persistence), - withSandbox(sandbox, { - snapshots: { - persistence: snapshots.persistence, - checkpoints: snapshots.checkpoints, - }, - }), + withSandbox(sandbox, { instances, snapshots }), ] ``` diff --git a/packages/ai-sandbox/skills/ai-sandbox/SKILL.md b/packages/ai-sandbox/skills/ai-sandbox/SKILL.md index f69b9d11b9..8a2dd5cd16 100644 --- a/packages/ai-sandbox/skills/ai-sandbox/SKILL.md +++ b/packages/ai-sandbox/skills/ai-sandbox/SKILL.md @@ -11,9 +11,8 @@ description: > snapshotMaxAge TTL. It also covers portable snapshots after a successful terminal run with withPersistence before withSandbox and memorySandboxSnapshots for local examples. It covers named saves with - saveNamedSandboxSnapshot, selected-checkpoint forks with - forkFromSandboxSnapshot, and authorized artifact reads with - resolveSnapshotArtifact. It covers defineWorkspace + snapshots.save, selected-checkpoint forks with snapshots.fork, and + authorized artifact reads with snapshots.readArtifact. It covers defineWorkspace (git/setup/scripts/skills/secrets/ instructions/plugins), defineSandboxPolicy (allow/ask/deny), lifecycle/resume, the SandboxHandle (fs/git/process/ports), capability tokens, defineSandbox @@ -201,16 +200,11 @@ middleware in this order, with the same persistence value in both places: import { withPersistence } from '@tanstack/ai-persistence' import { memorySandboxSnapshots, withSandbox } from '@tanstack/ai-sandbox' -const snapshots = await memorySandboxSnapshots() +const snapshots = await memorySandboxSnapshots({ sandbox, instances }) const middleware = [ withPersistence(snapshots.persistence), - withSandbox(sandbox, { - snapshots: { - persistence: snapshots.persistence, - checkpoints: snapshots.checkpoints, - }, - }), + withSandbox(sandbox, { instances, snapshots }), ] ``` @@ -233,18 +227,18 @@ garbage collection yet. Read `docs/sandbox/portable-snapshots.md` for the full server-only setup and the restore safety rules. -For a user-marked workspace state, call `saveNamedSandboxSnapshot` on the -server. It needs `definition`, `threadId`, `runId`, `instances`, `snapshots`, -and a label. It requires a live reusable sandbox. `reuse: 'none'` cannot save a -named checkpoint. +For a user-marked workspace state, call `snapshots.save` on the server. Bind +`sandbox` and `instances` at create time, or pass them on `save`. The call +needs `threadId`, `runId`, and a label. It requires a live reusable sandbox. +`reuse: 'none'` cannot save a named checkpoint. -To branch from a selected checkpoint, call `forkFromSandboxSnapshot` with the -source thread id, source checkpoint id, destination thread id, and `snapshots`. -The store must implement atomic `forkFromCheckpoint`. The destination thread -must be empty. A fork copies the selected snapshot, not the latest snapshot. +To branch from a selected checkpoint, call `snapshots.fork` with the thread id, +checkpoint id, and destination thread id. The store must implement atomic +`forkFromCheckpoint`. The destination thread must be empty. A fork copies the +selected snapshot, not the latest snapshot. -To send a checkpoint artifact, call `resolveSnapshotArtifact` on the server. -First authorize the caller for the supplied thread. The helper checks that the +To send a checkpoint artifact, call `snapshots.readArtifact` on the server. +First authorize the caller for the supplied thread. The method checks that the checkpoint belongs to that thread, then returns its metadata and bytes. It does not authorize a caller or create an HTTP response. diff --git a/packages/ai-sandbox/src/index.ts b/packages/ai-sandbox/src/index.ts index b0f3ecd7c2..e4b9724b3f 100644 --- a/packages/ai-sandbox/src/index.ts +++ b/packages/ai-sandbox/src/index.ts @@ -66,13 +66,19 @@ export type { SandboxSnapshotPolicy, } from './snapshots' export { memorySandboxSnapshots } from './memory-snapshots' -export type { MemorySandboxSnapshots } from './memory-snapshots' -export { - forkFromSandboxSnapshot, - resolveSnapshotArtifact, - saveNamedSandboxSnapshot, +export type { + MemorySandboxSnapshots, + MemorySandboxSnapshotsOptions, +} from './memory-snapshots' +export { createSandboxSnapshots } from './snapshot-operations' +export type { + CreateSandboxSnapshotsInput, + ForkSandboxSnapshotInput, + ReadSandboxSnapshotArtifactInput, + SandboxSnapshots, + SaveSandboxSnapshotInput, + SnapshotPersistence, } from './snapshot-operations' -export type { SandboxSnapshots } from './snapshot-operations' // Workspace projection capability (provided by withSandbox, consumed by harness adapters) export { diff --git a/packages/ai-sandbox/src/memory-snapshots.ts b/packages/ai-sandbox/src/memory-snapshots.ts index aaf3a7553a..2d13214c7b 100644 --- a/packages/ai-sandbox/src/memory-snapshots.ts +++ b/packages/ai-sandbox/src/memory-snapshots.ts @@ -26,6 +26,11 @@ import type { SandboxCheckpointWriterLease, SandboxCheckpointForkInput, } from './checkpoint-store' +import { createSandboxSnapshots } from './snapshot-operations' +import type { + CreateSandboxSnapshotsInput, + SandboxSnapshots, +} from './snapshot-operations' type BlobGetOptions = { range?: { offset: number; length?: number } } @@ -55,10 +60,18 @@ function resolveBlobRange( } } -export interface MemorySandboxSnapshots { - persistence: MemorySnapshotPersistence - checkpoints: ForkCapableSandboxCheckpointStore -} +export type MemorySandboxSnapshots = SandboxSnapshots< + MemorySnapshotPersistence, + ForkCapableSandboxCheckpointStore +> + +export type MemorySandboxSnapshotsOptions = Omit< + CreateSandboxSnapshotsInput< + MemorySnapshotPersistence, + ForkCapableSandboxCheckpointStore + >, + 'persistence' | 'checkpoints' +> const encoder = new TextEncoder() const compare = (a: string, b: string) => { @@ -617,11 +630,21 @@ async function bodyBytes(body: BlobBody): Promise { throw new TypeError('Unsupported blob body.') } -export async function memorySandboxSnapshots(): Promise { - return createMemorySandboxSnapshots() +export async function memorySandboxSnapshots( + options: MemorySandboxSnapshotsOptions = {}, +): Promise { + const { persistence, checkpoints } = await createMemorySandboxSnapshots() + return createSandboxSnapshots({ + persistence, + checkpoints, + ...options, + }) } -async function createMemorySandboxSnapshots(): Promise { +async function createMemorySandboxSnapshots(): Promise<{ + persistence: MemorySnapshotPersistence + checkpoints: ForkCapableSandboxCheckpointStore +}> { const messages = new Map>() const runs = new Map() const generations = new Map() diff --git a/packages/ai-sandbox/src/snapshot-operations.ts b/packages/ai-sandbox/src/snapshot-operations.ts index 156e853eee..770c34e6f4 100644 --- a/packages/ai-sandbox/src/snapshot-operations.ts +++ b/packages/ai-sandbox/src/snapshot-operations.ts @@ -19,7 +19,7 @@ import type { SandboxDefinition } from './sandbox' import type { SandboxSnapshotBundle, SandboxSnapshotPolicy } from './snapshots' import type { WorkspaceDefinition } from './workspace' -type SnapshotPersistence = { +export interface SnapshotPersistence { stores: { messages: { loadThread: (threadId: string) => Promise> @@ -29,10 +29,58 @@ type SnapshotPersistence = { } } -export interface SandboxSnapshots { - persistence: SnapshotPersistence - checkpoints: SandboxCheckpointStore +export interface CreateSandboxSnapshotsInput< + TPersistence extends SnapshotPersistence = SnapshotPersistence, + TCheckpoints extends SandboxCheckpointStore = SandboxCheckpointStore, +> { + persistence: TPersistence + checkpoints: TCheckpoints policy?: SandboxSnapshotPolicy + sandbox?: SandboxDefinition + instances?: SandboxInstanceStore + tenant?: { userId?: string; orgId?: string } + locks?: LockStore +} + +export interface SaveSandboxSnapshotInput { + threadId: string + runId: string + label: string + sandbox?: SandboxDefinition + instances?: SandboxInstanceStore + tenant?: { userId?: string; orgId?: string } + locks?: LockStore + signal?: AbortSignal + adapterName?: string +} + +export interface ForkSandboxSnapshotInput { + threadId: string + checkpointId: string + destinationThreadId: string + destinationCheckpointId?: string + createdAt?: number +} + +export interface ReadSandboxSnapshotArtifactInput { + threadId: string + checkpointId: string + artifactId: string +} + +export interface SandboxSnapshots< + TPersistence extends SnapshotPersistence = SnapshotPersistence, + TCheckpoints extends SandboxCheckpointStore = SandboxCheckpointStore, +> { + persistence: TPersistence + checkpoints: TCheckpoints + policy?: SandboxSnapshotPolicy + save: (input: SaveSandboxSnapshotInput) => Promise + fork: (input: ForkSandboxSnapshotInput) => Promise + readArtifact: (input: ReadSandboxSnapshotArtifactInput) => Promise<{ + artifact: SandboxCheckpoint['artifacts'][number] + bytes: Uint8Array + }> } type Failure = { error: unknown } @@ -161,12 +209,95 @@ function stageLockStore(locks: LockStore | undefined): LockStore | undefined { return { withLock } } -export async function saveNamedSandboxSnapshot(input: { +function requireSnapshotPersistence( + persistence: TPersistence, +): TPersistence { + const stores = persistence.stores + if (!stores?.messages || !stores.artifacts || !stores.blobs) { + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_MISSING_PERSISTENCE_STORES', + 'Sandbox snapshots require persistence stores.messages, stores.artifacts, and stores.blobs', + ) + } + return persistence +} + +export function createSandboxSnapshots< + TPersistence extends SnapshotPersistence, + TCheckpoints extends SandboxCheckpointStore, +>( + input: CreateSandboxSnapshotsInput, +): SandboxSnapshots { + const persistence = requireSnapshotPersistence(input.persistence) + const checkpoints = input.checkpoints + const policy = input.policy + const boundSandbox = input.sandbox + const boundInstances = input.instances + const boundTenant = input.tenant + const boundLocks = input.locks + + return { + persistence, + checkpoints, + ...(policy === undefined ? {} : { policy }), + async save(saveInput) { + const sandbox = saveInput.sandbox ?? boundSandbox + const instances = saveInput.instances ?? boundInstances + if (sandbox === undefined) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_MISSING_SANDBOX', + 'Named snapshots require a sandbox at create time or on save', + ) + if (instances === undefined) + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_MISSING_INSTANCES', + 'Named snapshots require instances at create time or on save', + ) + return saveNamedSandboxSnapshot({ + definition: sandbox, + threadId: saveInput.threadId, + runId: saveInput.runId, + instances, + persistence, + checkpoints, + policy, + label: saveInput.label, + tenant: saveInput.tenant ?? boundTenant, + locks: saveInput.locks ?? boundLocks, + signal: saveInput.signal, + adapterName: saveInput.adapterName, + }) + }, + fork(forkInput) { + return forkFromSandboxSnapshot({ + threadId: forkInput.threadId, + checkpointId: forkInput.checkpointId, + destinationThreadId: forkInput.destinationThreadId, + checkpoints, + destinationCheckpointId: forkInput.destinationCheckpointId, + createdAt: forkInput.createdAt, + }) + }, + readArtifact(readInput) { + return resolveSnapshotArtifact({ + threadId: readInput.threadId, + checkpointId: readInput.checkpointId, + artifactId: readInput.artifactId, + persistence, + checkpoints, + }) + }, + } +} + +async function saveNamedSandboxSnapshot(input: { definition: SandboxDefinition threadId: string runId: string instances: SandboxInstanceStore - snapshots: SandboxSnapshots + persistence: SnapshotPersistence + checkpoints: SandboxCheckpointStore + policy?: SandboxSnapshotPolicy label: string tenant?: { userId?: string; orgId?: string } locks?: LockStore @@ -177,7 +308,6 @@ export async function saveNamedSandboxSnapshot(input: { const threadId = input.threadId const runId = input.runId const instances = stageInstanceStore(input.instances) - const snapshots = input.snapshots const label = input.label const suppliedTenant = input.tenant const tenantUserId = suppliedTenant?.userId @@ -200,7 +330,7 @@ export async function saveNamedSandboxSnapshot(input: { const providerName = provider.name const resume = provider.resume.bind(provider) const ensureExisting = stageEnsureExistingSandbox(definition) - const persistence = snapshots.persistence + const persistence = input.persistence const stores = persistence.stores const messages = stores.messages const loadThread = messages.loadThread.bind(messages) @@ -215,12 +345,12 @@ export async function saveNamedSandboxSnapshot(input: { head: headBlob, put: putBlob, } - const checkpoints = snapshots.checkpoints + const checkpoints = input.checkpoints const acquireWriter = checkpoints.acquireWriter.bind(checkpoints) const getHead = checkpoints.getHead.bind(checkpoints) const append = checkpoints.append.bind(checkpoints) const policy = effectivePolicy( - snapshots.policy, + input.policy, workspace === undefined ? undefined : computeWorkspaceHash(workspace), ) const workspaceSecrets = workspace?.secrets @@ -306,24 +436,23 @@ export async function saveNamedSandboxSnapshot(input: { ) } -export async function forkFromSandboxSnapshot(input: { - sourceThreadId: string - sourceCheckpointId: string +async function forkFromSandboxSnapshot(input: { + threadId: string + checkpointId: string destinationThreadId: string - snapshots: SandboxSnapshots + checkpoints: SandboxCheckpointStore destinationCheckpointId?: string createdAt?: number }): Promise { - const sourceThreadId = input.sourceThreadId - const sourceCheckpointId = input.sourceCheckpointId + const sourceThreadId = input.threadId + const sourceCheckpointId = input.checkpointId const destinationThreadId = input.destinationThreadId - const snapshots = input.snapshots const suppliedDestinationCheckpointId = input.destinationCheckpointId const suppliedCreatedAt = input.createdAt const destinationCheckpointId = suppliedDestinationCheckpointId ?? crypto.randomUUID() const createdAt = suppliedCreatedAt ?? Date.now() - const checkpoints = snapshots.checkpoints + const checkpoints = input.checkpoints const acquireWriter = checkpoints.acquireWriter.bind(checkpoints) const forkFromCheckpoint = checkpoints.forkFromCheckpoint?.bind(checkpoints) @@ -356,11 +485,12 @@ async function sha256(bytes: Uint8Array): Promise { ).join('') } -export async function resolveSnapshotArtifact(input: { +async function resolveSnapshotArtifact(input: { threadId: string checkpointId: string artifactId: string - snapshots: SandboxSnapshots + persistence: SnapshotPersistence + checkpoints: SandboxCheckpointStore }): Promise<{ artifact: SandboxCheckpoint['artifacts'][number] bytes: Uint8Array @@ -368,10 +498,9 @@ export async function resolveSnapshotArtifact(input: { const threadId = input.threadId const checkpointId = input.checkpointId const artifactId = input.artifactId - const snapshots = input.snapshots - const checkpoints = snapshots.checkpoints + const checkpoints = input.checkpoints const getCheckpoint = checkpoints.get.bind(checkpoints) - const persistence = snapshots.persistence + const persistence = input.persistence const stores = persistence.stores const blobs = stores.blobs const getBlob = blobs.get.bind(blobs) diff --git a/packages/ai-sandbox/src/snapshots.ts b/packages/ai-sandbox/src/snapshots.ts index d9a81c678d..cdd4fb88fa 100644 --- a/packages/ai-sandbox/src/snapshots.ts +++ b/packages/ai-sandbox/src/snapshots.ts @@ -37,6 +37,9 @@ export interface SandboxSnapshotBundle { } export type SandboxSnapshotErrorCode = + | 'SANDBOX_SNAPSHOT_MISSING_SANDBOX' + | 'SANDBOX_SNAPSHOT_MISSING_INSTANCES' + | 'SANDBOX_SNAPSHOT_MISSING_PERSISTENCE_STORES' | 'SANDBOX_SNAPSHOT_MISSING_REUSABLE_SANDBOX' | 'SANDBOX_SNAPSHOT_REUSE_NONE' | 'SANDBOX_SNAPSHOT_MISSING_CHECKPOINT_ARTIFACT' diff --git a/packages/ai-sandbox/tests/memory-snapshots-declaration.test-d.ts b/packages/ai-sandbox/tests/memory-snapshots-declaration.test-d.ts index 8511956dd4..852e34a9e9 100644 --- a/packages/ai-sandbox/tests/memory-snapshots-declaration.test-d.ts +++ b/packages/ai-sandbox/tests/memory-snapshots-declaration.test-d.ts @@ -25,5 +25,8 @@ declare const snapshots: MemorySandboxSnapshots expectTypeOf(snapshots.persistence).toMatchTypeOf< AIPersistence >() +expectTypeOf(snapshots.save).toBeFunction() +expectTypeOf(snapshots.fork).toBeFunction() +expectTypeOf(snapshots.readArtifact).toBeFunction() // @ts-expect-error immutable identity fields are not patchable snapshots.persistence.stores.generationRuns.update('run', { threadId: 'other' }) diff --git a/packages/ai-sandbox/tests/memory-snapshots-import.test.ts b/packages/ai-sandbox/tests/memory-snapshots-import.test.ts index 2bbab650de..576fe922d4 100644 --- a/packages/ai-sandbox/tests/memory-snapshots-import.test.ts +++ b/packages/ai-sandbox/tests/memory-snapshots-import.test.ts @@ -11,6 +11,9 @@ describe('memorySandboxSnapshots runtime dependencies', () => { await expect(memorySandboxSnapshots()).resolves.toMatchObject({ persistence: expect.any(Object), checkpoints: expect.any(Object), + save: expect.any(Function), + fork: expect.any(Function), + readArtifact: expect.any(Function), }) }) }) diff --git a/packages/ai-sandbox/tests/snapshot-operations.test-d.ts b/packages/ai-sandbox/tests/snapshot-operations.test-d.ts index dc00aeb8b1..f95f04010e 100644 --- a/packages/ai-sandbox/tests/snapshot-operations.test-d.ts +++ b/packages/ai-sandbox/tests/snapshot-operations.test-d.ts @@ -3,6 +3,9 @@ import { memorySandboxSnapshots, SandboxSnapshotError } from '../src' import type { SandboxSnapshotErrorCode, SandboxSnapshots } from '../src' type ExpectedSandboxSnapshotErrorCode = + | 'SANDBOX_SNAPSHOT_MISSING_SANDBOX' + | 'SANDBOX_SNAPSHOT_MISSING_INSTANCES' + | 'SANDBOX_SNAPSHOT_MISSING_PERSISTENCE_STORES' | 'SANDBOX_SNAPSHOT_MISSING_REUSABLE_SANDBOX' | 'SANDBOX_SNAPSHOT_REUSE_NONE' | 'SANDBOX_SNAPSHOT_MISSING_CHECKPOINT_ARTIFACT' @@ -28,6 +31,9 @@ expectTypeOf(sourceError.code).toEqualTypeOf() async function assignActualMemorySnapshots(): Promise { const snapshots = await memorySandboxSnapshots() const structuralSnapshots: SandboxSnapshots = snapshots + expectTypeOf(snapshots.save).toBeFunction() + expectTypeOf(snapshots.fork).toBeFunction() + expectTypeOf(snapshots.readArtifact).toBeFunction() void structuralSnapshots } void assignActualMemorySnapshots diff --git a/packages/ai-sandbox/tests/snapshot-operations.test.ts b/packages/ai-sandbox/tests/snapshot-operations.test.ts index 9e16920e70..ac6ed4ddec 100644 --- a/packages/ai-sandbox/tests/snapshot-operations.test.ts +++ b/packages/ai-sandbox/tests/snapshot-operations.test.ts @@ -3,15 +3,13 @@ import { describe, expect, it, vi } from 'vitest' import { computeSandboxKey, computeWorkspaceHash, + createSandboxSnapshots, createSecrets, defineSandbox, - forkFromSandboxSnapshot, InMemorySandboxCheckpointStore, InMemorySandboxInstanceStore, memorySandboxSnapshots, - resolveSnapshotArtifact, SandboxSnapshotError, - saveNamedSandboxSnapshot, } from '../src' import { makeFakeHandle, makeFakeProvider } from './fakes' import type { @@ -188,37 +186,36 @@ async function namedFixture( updatedAt: Date.now(), }) } + const locks = new InMemoryLockStore() return { definition, instances, provider, memory, - snapshots: { + snapshots: createSandboxSnapshots({ persistence: memory.persistence, checkpoints: options.checkpoints ?? memory.checkpoints, + sandbox: definition, + instances, + locks, ...(options.policy === undefined ? {} : { policy: options.policy }), - }, - locks: new InMemoryLockStore(), + }), + locks, } } -type NamedSaveInput = Parameters[0] +type NamedSaveInput = Parameters[0] function namedSave( fixture: NamedFixture, changes: Partial = {}, ) { - const input: NamedSaveInput = { - definition: fixture.definition, + return fixture.snapshots.save({ threadId: THREAD, runId: RUN, - instances: fixture.instances, - snapshots: fixture.snapshots, label: LABEL, - locks: fixture.locks, ...changes, - } - return saveNamedSandboxSnapshot(input) + }) } function checkpoint( @@ -242,7 +239,7 @@ function checkpoint( } async function appendCheckpoint( - snapshots: SandboxSnapshots, + snapshots: { checkpoints: SandboxCheckpointStore }, value: SandboxCheckpoint, ): Promise { const writer = await snapshots.checkpoints.acquireWriter(value.threadId) @@ -299,17 +296,23 @@ async function artifactSnapshot(input: { function resolveArtifact( snapshots: SandboxSnapshots, - changes: Partial[0]> = {}, + changes: Partial[0]> = {}, ) { - return resolveSnapshotArtifact({ + return snapshots.readArtifact({ threadId: THREAD, checkpointId: 'checkpoint', artifactId: 'artifact', - snapshots, ...changes, }) } +function operationsFor( + persistence: SandboxSnapshots['persistence'], + checkpoints: SandboxCheckpointStore, +): SandboxSnapshots { + return createSandboxSnapshots({ persistence, checkpoints }) +} + function workspaceHandle(input: { files: ReadonlyArray<{ path: string; content: string }> listCalls?: Array @@ -492,7 +495,7 @@ describe('public sandbox snapshot operations', () => { } await expect( - namedSave(fixture, { definition: structuralDefinition }), + namedSave(fixture, { sandbox: structuralDefinition }), ).resolves.toMatchObject({ reason: 'named' }) expect(ensureReads).toBe(1) }) @@ -614,11 +617,10 @@ describe('public sandbox snapshot operations', () => { }, } - const fork = forkFromSandboxSnapshot({ - sourceThreadId: 'source', - sourceCheckpointId: 'source-checkpoint', + const fork = operationsFor(snapshots.persistence, checkpoints).fork({ + threadId: 'source', + checkpointId: 'source-checkpoint', destinationThreadId: 'destination', - snapshots: { persistence: snapshots.persistence, checkpoints }, }) await acquisitionStarted.promise acquisitionGate.resolve() @@ -646,11 +648,10 @@ describe('public sandbox snapshot operations', () => { } await expect( - forkFromSandboxSnapshot({ - sourceThreadId: 'source', - sourceCheckpointId: 'checkpoint', + operationsFor(snapshots.persistence, checkpoints).fork({ + threadId: 'source', + checkpointId: 'checkpoint', destinationThreadId: 'destination', - snapshots: { persistence: snapshots.persistence, checkpoints }, }), ).rejects.toThrow('fork staging failed') expect(acquisitionCalls).toBe(0) @@ -674,7 +675,7 @@ describe('public sandbox snapshot operations', () => { }) let threadReads = 0 let checkpointPending = false - const input: Parameters[0] = { + const input: Parameters[0] = { get threadId() { threadReads++ if (checkpointPending) throw new Error('late threadId read') @@ -682,10 +683,9 @@ describe('public sandbox snapshot operations', () => { }, checkpointId: 'checkpoint', artifactId: 'artifact', - snapshots, } - const resolved = resolveSnapshotArtifact(input) + const resolved = snapshots.readArtifact(input) await checkpointStarted.promise checkpointPending = true checkpointGate.resolve(stored) @@ -1167,13 +1167,12 @@ describe('public sandbox snapshot operations', () => { const sourceConversationBefore = await snapshots.persistence.stores.messages.loadThread('source') - const result = await forkFromSandboxSnapshot({ - sourceThreadId: 'source', - sourceCheckpointId: 'selected', + const result = await snapshots.fork({ + threadId: 'source', + checkpointId: 'selected', destinationThreadId: 'destination', destinationCheckpointId: 'fork', createdAt: 3, - snapshots, }) expect(result).toMatchObject({ @@ -1198,21 +1197,17 @@ describe('public sandbox snapshot operations', () => { it('releases once and does not renew a successful fork', async () => { const snapshots = await memorySandboxSnapshots() const checkpoints = new ForkProbeStore(snapshots.checkpoints) - const bundle: SandboxSnapshots = { - persistence: snapshots.persistence, - checkpoints, - } + const bundle = operationsFor(snapshots.persistence, checkpoints) await appendCheckpoint( bundle, checkpoint({ id: 'source-checkpoint', threadId: 'source' }), ) await expect( - forkFromSandboxSnapshot({ - sourceThreadId: 'source', - sourceCheckpointId: 'source-checkpoint', + bundle.fork({ + threadId: 'source', + checkpointId: 'source-checkpoint', destinationThreadId: 'destination', - snapshots: bundle, }), ).resolves.toMatchObject({ reason: 'fork-root' }) expect(checkpoints.releases).toBe(2) @@ -1225,14 +1220,10 @@ describe('public sandbox snapshot operations', () => { const snapshots = await memorySandboxSnapshots() await expect( - forkFromSandboxSnapshot({ - sourceThreadId: 'source', - sourceCheckpointId: 'source-checkpoint', + operationsFor(snapshots.persistence, withoutFork(checkpoints)).fork({ + threadId: 'source', + checkpointId: 'source-checkpoint', destinationThreadId: 'destination', - snapshots: { - persistence: snapshots.persistence, - checkpoints: withoutFork(checkpoints), - }, }), ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_FORK_UNAVAILABLE' }) expect(checkpoints.releases).toBe(1) @@ -1245,11 +1236,10 @@ describe('public sandbox snapshot operations', () => { checkpoints.releaseError = new Error('release failed') await expect( - forkFromSandboxSnapshot({ - sourceThreadId: 'source', - sourceCheckpointId: 'source-checkpoint', + operationsFor(snapshots.persistence, checkpoints).fork({ + threadId: 'source', + checkpointId: 'source-checkpoint', destinationThreadId: 'destination', - snapshots: { persistence: snapshots.persistence, checkpoints }, }), ).rejects.toThrow('fork failed') expect(checkpoints.releases).toBe(1) @@ -1258,10 +1248,7 @@ describe('public sandbox snapshot operations', () => { it('reports release failure after a successful fork publication', async () => { const snapshots = await memorySandboxSnapshots() const checkpoints = new ForkProbeStore(snapshots.checkpoints) - const bundle: SandboxSnapshots = { - persistence: snapshots.persistence, - checkpoints, - } + const bundle = operationsFor(snapshots.persistence, checkpoints) await appendCheckpoint( bundle, checkpoint({ id: 'source-checkpoint', threadId: 'source' }), @@ -1269,11 +1256,10 @@ describe('public sandbox snapshot operations', () => { checkpoints.releaseError = new Error('release failed') await expect( - forkFromSandboxSnapshot({ - sourceThreadId: 'source', - sourceCheckpointId: 'source-checkpoint', + bundle.fork({ + threadId: 'source', + checkpointId: 'source-checkpoint', destinationThreadId: 'destination', - snapshots: bundle, }), ).rejects.toThrow('release failed') expect(await checkpoints.getHead('destination')).not.toBeNull() @@ -1371,4 +1357,95 @@ describe('public sandbox snapshot operations', () => { ).rejects.toBeInstanceOf(SandboxSnapshotError) }) }) + + describe('create and bind', () => { + it('returns save, fork, and readArtifact from memorySandboxSnapshots', async () => { + const snapshots = await memorySandboxSnapshots() + expect(snapshots.save).toEqual(expect.any(Function)) + expect(snapshots.fork).toEqual(expect.any(Function)) + expect(snapshots.readArtifact).toEqual(expect.any(Function)) + }) + + it('rejects create when persistence stores are missing', async () => { + const snapshots = await memorySandboxSnapshots() + Reflect.deleteProperty(snapshots.persistence.stores, 'messages') + expect(() => + createSandboxSnapshots({ + persistence: snapshots.persistence, + checkpoints: snapshots.checkpoints, + }), + ).toThrow( + expect.objectContaining({ + code: 'SANDBOX_SNAPSHOT_MISSING_PERSISTENCE_STORES', + }), + ) + }) + + it('rejects save when sandbox and instances are missing', async () => { + const snapshots = await memorySandboxSnapshots() + await expect( + snapshots.save({ threadId: THREAD, runId: RUN, label: LABEL }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_MISSING_SANDBOX' }) + }) + + it('rejects save when only instances are missing', async () => { + const fixture = await namedFixture() + const snapshots = createSandboxSnapshots({ + persistence: fixture.memory.persistence, + checkpoints: fixture.memory.checkpoints, + sandbox: fixture.definition, + }) + await expect( + snapshots.save({ threadId: THREAD, runId: RUN, label: LABEL }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_MISSING_INSTANCES' }) + }) + + it('uses save overrides for sandbox and instances', async () => { + const fixture = await namedFixture() + const snapshots = await memorySandboxSnapshots() + await snapshots.persistence.stores.messages.saveThread(THREAD, [ + { role: 'user', content: 'saved' }, + ]) + const saved = await snapshots.save({ + threadId: THREAD, + runId: RUN, + label: 'override', + sandbox: fixture.definition, + instances: fixture.instances, + locks: fixture.locks, + }) + expect(saved).toMatchObject({ + reason: 'named', + label: 'override', + conversation: [{ role: 'user', content: 'saved' }], + }) + }) + + it('binds sandbox and instances on the memory factory', async () => { + const instances = new InMemorySandboxInstanceStore() + const provider = makeFakeProvider() + const sandbox = defineSandbox({ id: 'sandbox', provider }) + await instances.upsert({ + key: sandbox.key({ threadId: THREAD, runId: 'old' }), + provider: provider.name, + providerSandboxId: 'existing', + threadId: THREAD, + updatedAt: Date.now(), + }) + const snapshots = await memorySandboxSnapshots({ + sandbox, + instances, + locks: new InMemoryLockStore(), + }) + await snapshots.persistence.stores.messages.saveThread(THREAD, [ + { role: 'user', content: 'bound' }, + ]) + await expect( + snapshots.save({ threadId: THREAD, runId: RUN, label: LABEL }), + ).resolves.toMatchObject({ + reason: 'named', + conversation: [{ role: 'user', content: 'bound' }], + }) + }) + }) }) diff --git a/testing/e2e/src/routes/api.sandbox-file-persistence.ts b/testing/e2e/src/routes/api.sandbox-file-persistence.ts index 22bd9f3ce3..fba283a2fe 100644 --- a/testing/e2e/src/routes/api.sandbox-file-persistence.ts +++ b/testing/e2e/src/routes/api.sandbox-file-persistence.ts @@ -4,10 +4,7 @@ import { InMemorySandboxInstanceStore, defineSandbox, defineWorkspace, - forkFromSandboxSnapshot, memorySandboxSnapshots, - resolveSnapshotArtifact, - saveNamedSandboxSnapshot, withSandbox, } from '@tanstack/ai-sandbox' import { withPersistence } from '@tanstack/ai-persistence' @@ -335,7 +332,6 @@ export const Route = createFileRoute('/api/sandbox-file-persistence')({ originalExists ? fakeHandle('original', sourceFiles) : null, destroy: async () => {}, } - const snapshots = await memorySandboxSnapshots() const instances = new InMemorySandboxInstanceStore() const sandbox = defineSandbox({ id: 'file-persistence', @@ -343,6 +339,10 @@ export const Route = createFileRoute('/api/sandbox-file-persistence')({ workspace: defineWorkspace({ source: { type: 'none' } }), fileEvents: false, }) + const snapshots = await memorySandboxSnapshots({ + sandbox, + instances, + }) const key = sandbox.key({ threadId, runId: 'save', store: instances }) await instances.upsert({ key, @@ -369,12 +369,9 @@ export const Route = createFileRoute('/api/sandbox-file-persistence')({ createdAt: 1, }) - const saved = await saveNamedSandboxSnapshot({ - definition: sandbox, + const saved = await snapshots.save({ threadId, runId: 'save', - instances, - snapshots, label: 'release-1', }) originalExists = false @@ -396,11 +393,10 @@ export const Route = createFileRoute('/api/sandbox-file-persistence')({ await snapshots.checkpoints.get(automaticHeadId) if (!automaticCheckpoint) throw new Error('Expected the automatic checkpoint') - const artifact = await resolveSnapshotArtifact({ + const artifact = await snapshots.readArtifact({ threadId, checkpointId: saved.id, artifactId: 'artifact-1', - snapshots, }) const sourceBeforeFork = await snapshots.checkpoints.list(threadId) const { result: fork, sourceMessagesUnchanged } = @@ -408,11 +404,10 @@ export const Route = createFileRoute('/api/sandbox-file-persistence')({ loadSourceMessages: () => snapshots.persistence.stores.messages.loadThread(threadId), fork: () => - forkFromSandboxSnapshot({ - sourceThreadId: threadId, - sourceCheckpointId: saved.id, + snapshots.fork({ + threadId, + checkpointId: saved.id, destinationThreadId, - snapshots, destinationCheckpointId: 'fork-root', createdAt: 2, }), From 7fbd218af8d9375665f276a53a341a80125e1537 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Fri, 14 Aug 2026 16:16:07 +0200 Subject: [PATCH 05/12] docs(sandbox): split portable snapshots into journey pages One page per reader goal: reload, named save, fork, artifact download, and policy. --- docs/config.json | 33 +- docs/persistence/build-a-sandbox-adapter.md | 4 + docs/persistence/overview.md | 4 + docs/sandbox/durability.md | 4 +- docs/sandbox/lifecycle.md | 8 +- docs/sandbox/overview.md | 4 +- docs/sandbox/portable-snapshots-artifacts.md | 68 +++ docs/sandbox/portable-snapshots-configure.md | 158 +++++++ docs/sandbox/portable-snapshots-fork.md | 88 ++++ docs/sandbox/portable-snapshots-safety.md | 68 +++ docs/sandbox/portable-snapshots-save.md | 73 ++++ docs/sandbox/portable-snapshots.md | 404 ++---------------- docs/sandbox/providers.md | 2 +- packages/ai-sandbox/README.md | 10 +- .../ai-sandbox/skills/ai-sandbox/SKILL.md | 19 +- 15 files changed, 559 insertions(+), 388 deletions(-) create mode 100644 docs/sandbox/portable-snapshots-artifacts.md create mode 100644 docs/sandbox/portable-snapshots-configure.md create mode 100644 docs/sandbox/portable-snapshots-fork.md create mode 100644 docs/sandbox/portable-snapshots-safety.md create mode 100644 docs/sandbox/portable-snapshots-save.md diff --git a/docs/config.json b/docs/config.json index 25a17b8d2e..8cf461935c 100644 --- a/docs/config.json +++ b/docs/config.json @@ -240,7 +240,8 @@ { "label": "Overview", "to": "persistence/overview", - "addedAt": "2026-08-04" + "addedAt": "2026-08-04", + "updatedAt": "2026-08-14" }, { "label": "Chat Persistence", @@ -304,7 +305,8 @@ { "label": "Build a Sandbox Adapter", "to": "persistence/build-a-sandbox-adapter", - "addedAt": "2026-08-04" + "addedAt": "2026-08-04", + "updatedAt": "2026-08-14" }, { "label": "Store Reference", @@ -556,11 +558,36 @@ "updatedAt": "2026-08-14" }, { - "label": "Portable Sandbox Snapshots", + "label": "Portable Snapshots", "to": "sandbox/portable-snapshots", "addedAt": "2026-08-14", "updatedAt": "2026-08-14" }, + { + "label": "Keep Files After Reload", + "to": "sandbox/portable-snapshots-configure", + "addedAt": "2026-08-14" + }, + { + "label": "Save a Named Version", + "to": "sandbox/portable-snapshots-save", + "addedAt": "2026-08-14" + }, + { + "label": "Branch From a Version", + "to": "sandbox/portable-snapshots-fork", + "addedAt": "2026-08-14" + }, + { + "label": "Send a Frozen File", + "to": "sandbox/portable-snapshots-artifacts", + "addedAt": "2026-08-14" + }, + { + "label": "What a Snapshot Stores", + "to": "sandbox/portable-snapshots-safety", + "addedAt": "2026-08-14" + }, { "label": "Instance Durability", "to": "sandbox/durability", diff --git a/docs/persistence/build-a-sandbox-adapter.md b/docs/persistence/build-a-sandbox-adapter.md index 540463f310..8f351a2eae 100644 --- a/docs/persistence/build-a-sandbox-adapter.md +++ b/docs/persistence/build-a-sandbox-adapter.md @@ -27,6 +27,10 @@ place the two meet. It is the third of the adapter walkthroughs, next to [generation](./build-your-own-generation-adapter), and it needs neither of their store contracts. +To rebuild completed workspace files after the provider sandbox is gone, use +[Keep Files After Reload](../sandbox/portable-snapshots-configure). Pass the same +persistence object that `withPersistence` uses. + ## Decide what you store | You keep | Wire | You get | You give up | diff --git a/docs/persistence/overview.md b/docs/persistence/overview.md index cfdaa019e4..b9568fcf79 100644 --- a/docs/persistence/overview.md +++ b/docs/persistence/overview.md @@ -23,6 +23,10 @@ streaming. That is [Resumable Streams](../resumable-streams/overview), a differe layer you can add on its own. Step 3 below combines them, which is what most apps end up wanting. +When the provider sandbox is gone, the workspace files can also disappear. Use +the same persistence object with +[Keep Files After Reload](../sandbox/portable-snapshots-configure). + ## Install ```bash diff --git a/docs/sandbox/durability.md b/docs/sandbox/durability.md index e1a7e3720a..c5ce0a8506 100644 --- a/docs/sandbox/durability.md +++ b/docs/sandbox/durability.md @@ -21,8 +21,8 @@ owned by `@tanstack/ai-sandbox`, independent of `@tanstack/ai-persistence` you compose a separate middleware. Instance durability does not copy the workspace into your application storage. -Use [Portable Sandbox Snapshots](./portable-snapshots) when a new sandbox must -rebuild completed files, artifacts, and the saved conversation. +Use [Keep Files After Reload](./portable-snapshots-configure) when a new +sandbox must rebuild completed files, artifacts, and the saved conversation. It is also not the agent's *output*. This page keeps a sandbox findable across processes; keeping the run's event stream readable across processes is diff --git a/docs/sandbox/lifecycle.md b/docs/sandbox/lifecycle.md index c92cfd90a2..477dae8ea3 100644 --- a/docs/sandbox/lifecycle.md +++ b/docs/sandbox/lifecycle.md @@ -11,9 +11,9 @@ cost once and reuse the result: keep one sandbox per thread, snapshot it after setup, and resume instead of re-bootstrapping on the next run. When you must also recover files after a sandbox is gone, configure -[Portable Sandbox Snapshots](./portable-snapshots). Provider-native snapshots -make bootstrap faster. Portable snapshots save the completed workspace as -durable application data. +[Keep Files After Reload](./portable-snapshots-configure). Provider-native +snapshots make bootstrap faster. Portable snapshots save the completed +workspace as durable application data. ```ts import { defineSandbox, defineWorkspace, githubRepo } from '@tanstack/ai-sandbox' @@ -131,7 +131,7 @@ full bootstrap. Portable sandbox snapshots run after this lifecycle work. They restore a saved workspace only into a newly created private sandbox. They never overwrite a -live resumed sandbox. See [Portable Sandbox Snapshots](./portable-snapshots). +live resumed sandbox. See [Portable Snapshots](./portable-snapshots). > Which providers support durable disk, snapshots, and resume-by-id is listed on > [Providers](./providers). diff --git a/docs/sandbox/overview.md b/docs/sandbox/overview.md index 8d1898ea07..2984c1563a 100644 --- a/docs/sandbox/overview.md +++ b/docs/sandbox/overview.md @@ -116,8 +116,8 @@ After that, pick the piece you need: - [Tools](./tools): bridge your app's own tools into the in-sandbox agent. - [Policy](./policy): allow, ask or deny guardrails on what the agent may run. - [Lifecycle & Snapshots](./lifecycle): reuse a sandbox, snapshot after setup, resume. -- [Portable Sandbox Snapshots](./portable-snapshots): save completed files and - artifacts, then rebuild them in a new sandbox. +- [Portable Snapshots](./portable-snapshots): keep completed files after the + sandbox is gone. Start with [Keep Files After Reload](./portable-snapshots-configure). - [Instance Durability](./durability): reuse it across replicas too. - [Durable Runs](./durable-runs): let a run outlive the tab, and turn it on. - [Events](./events): stream the agent's edits and tool calls to a UI, and choose what diff --git a/docs/sandbox/portable-snapshots-artifacts.md b/docs/sandbox/portable-snapshots-artifacts.md new file mode 100644 index 0000000000..bb39c1c6d0 --- /dev/null +++ b/docs/sandbox/portable-snapshots-artifacts.md @@ -0,0 +1,68 @@ +--- +title: Send a Frozen File +id: portable-snapshots-artifacts +order: 14 +description: "Serve copied artifact bytes from one checkpoint through an authorized server route." +--- + +A user wants to download a generated file from a saved version. The blob store +is not a public URL. `snapshots.readArtifact` reads the copied bytes from one +checkpoint. Your route then returns an HTTP response. + +The method makes sure that the checkpoint belongs to the supplied thread. Your +route must still authorize that thread first. The method returns metadata and +`Uint8Array` bytes. It does not create an HTTP response. + +Keep this route on the server. If the session cannot access the thread, return +404. Then call `snapshots.readArtifact`. + +```ts +import { requireSession } from './auth' +import { snapshots } from './sandbox-server' + +export async function GET(request: Request) { + const session = await requireSession(request) + const url = new URL(request.url) + const threadId = url.searchParams.get('threadId') + const checkpointId = url.searchParams.get('checkpointId') + const artifactId = url.searchParams.get('artifactId') + + if (!threadId || !checkpointId || !artifactId) { + return new Response('Not found', { status: 404 }) + } + if (!(await session.canAccessThread(threadId))) { + return new Response('Not found', { status: 404 }) + } + + const { artifact, bytes } = await snapshots.readArtifact({ + threadId, + checkpointId, + artifactId, + }) + return new Response(bytes.slice(), { + headers: { + 'content-type': artifact.mimeType, + 'content-length': String(artifact.size), + }, + }) +} +``` + +The client uses the authorized route as an artifact URL. It must not read the +blob store or call `snapshots.readArtifact` in the browser. + +```ts +export function snapshotArtifactUrl( + threadId: string, + checkpointId: string, + artifactId: string, +) { + const query = new URLSearchParams({ threadId, checkpointId, artifactId }) + return `/api/snapshots/artifact?${query}` +} +``` + +A named save copies thread artifacts into the checkpoint. Automatic restore +does not change those copied bytes. See +[Save a Named Version](./portable-snapshots-save) and +[What a Snapshot Stores](./portable-snapshots-safety). diff --git a/docs/sandbox/portable-snapshots-configure.md b/docs/sandbox/portable-snapshots-configure.md new file mode 100644 index 0000000000..31ff72ad6d --- /dev/null +++ b/docs/sandbox/portable-snapshots-configure.md @@ -0,0 +1,158 @@ +--- +title: Keep Files After Reload +id: portable-snapshots-configure +order: 11 +description: "Wire portable snapshots so a later run rebuilds completed sandbox files from your persistence." +--- + +You have a sandbox chat. The agent writes files. The provider sandbox then goes +away. The next run starts empty. + +Portable snapshots save those files after each successful terminal run. A later +run restores the latest checkpoint into a new private sandbox. By the end of +this page, `chat()` writes and restores those checkpoints. + +Create one persistence value. Pass that exact value to `withPersistence` and to +the snapshots object. Put `withPersistence` before `withSandbox`. + +## Start with the memory factory + +Use `memorySandboxSnapshots` for local development. It creates persistence, a +checkpoint store, and the snapshot methods as one object. It does not load +`@tanstack/ai-persistence` at runtime. + +```ts +import { chat } from '@tanstack/ai' +import { grokBuildText } from '@tanstack/ai-grok-build' +import { withPersistence } from '@tanstack/ai-persistence' +import { + defineSandbox, + defineWorkspace, + InMemorySandboxInstanceStore, + memorySandboxSnapshots, + withSandbox, +} from '@tanstack/ai-sandbox' +import { dockerSandbox } from '@tanstack/ai-sandbox-docker' + +const instances = new InMemorySandboxInstanceStore() +const userId = 'user-123' + +const sandbox = defineSandbox({ + id: 'app-builder', + provider: dockerSandbox({ image: 'node:22' }), + workspace: defineWorkspace({ source: { type: 'none' } }), + lifecycle: { reuse: 'thread' }, +}) + +const snapshots = await memorySandboxSnapshots({ + sandbox, + instances, +}) + +const result = chat({ + threadId: 'app-thread', + context: { userId }, + adapter: grokBuildText('grok-build'), + messages: [{ role: 'user', content: 'Create a landing page.' }], + middleware: [ + withPersistence(snapshots.persistence), + withSandbox(sandbox, { + instances, + snapshots, + }), + ], +}) + +void result +``` + +You can bind `sandbox`, `instances`, `tenant`, and `locks` at create time. A +later `snapshots.save` call can override those values. See +[Save a Named Version](./portable-snapshots-save). + +Keep `instances` in the same server module as this middleware. A named save +must use this same instance store. + +Pass the session `userId` in `context` for every run. Pass that same user id as +`tenant.userId` on `snapshots.save`. + +## Add snapshots to existing persistence + +If `withPersistence` already uses a persistence object, pass that same object +to `createSandboxSnapshots`. Do not create a second message store. + +The persistence object must include `stores.messages`, `stores.artifacts`, and +`stores.blobs`. + +```ts +import { chat } from '@tanstack/ai' +import { grokBuildText } from '@tanstack/ai-grok-build' +import { withPersistence } from '@tanstack/ai-persistence' +import { + createSandboxSnapshots, + InMemorySandboxCheckpointStore, + withSandbox, +} from '@tanstack/ai-sandbox' +import { instances, persistence, sandbox } from './sandbox-server' + +const snapshots = createSandboxSnapshots({ + persistence, + checkpoints: new InMemorySandboxCheckpointStore(), + sandbox, + instances, +}) + +const result = chat({ + threadId: 'app-thread', + adapter: grokBuildText('grok-build'), + messages: [{ role: 'user', content: 'Create a landing page.' }], + middleware: [ + withPersistence(snapshots.persistence), + withSandbox(sandbox, { + instances, + snapshots, + }), + ], +}) + +void result +``` + +## Use durable stores in production + +`memorySandboxSnapshots` is for local development and tests. Production needs +durable message, artifact, and blob stores, plus a durable checkpoint store. + +The React chat example exports `sqliteSandboxSnapshots()` for a Node 22.5+ +server. That function is an example adapter. It is not a package export. + +Use one SQLite transaction for every checkpoint write, head update, and blob +reference count update. Use one transaction for a fork. The fork transaction +must also copy the source conversation. It must reject a destination thread +that already has persisted state. + +## Writer leases + +Each thread has one checkpoint writer lease. + +- A second run for the same thread gets a writer conflict while the lease is + active. +- The middleware renews the lease while the run is active. +- Pause and detach paths release the lease. They do not publish a partial + checkpoint. +- If the writer loses the lease, the middleware does not publish the + checkpoint. A later successful run can create a new checkpoint. + +## What happens on the next run + +A later run restores the latest checkpoint only into a new private sandbox. A +live resumed sandbox keeps its current files. See +[Portable Sandbox Snapshots](./portable-snapshots) for that restore rule. + +Instance durability finds a provider sandbox across server processes. When that +sandbox is gone, portable snapshots rebuild the workspace. See +[Instance Durability](./durability). + +The default policy excludes `.git`, `node_modules`, and `.env*` paths. Read +[What a Snapshot Stores](./portable-snapshots-safety) before you replace that +policy. diff --git a/docs/sandbox/portable-snapshots-fork.md b/docs/sandbox/portable-snapshots-fork.md new file mode 100644 index 0000000000..fc438b0e14 --- /dev/null +++ b/docs/sandbox/portable-snapshots-fork.md @@ -0,0 +1,88 @@ +--- +title: Branch From a Version +id: portable-snapshots-fork +order: 13 +description: "Copy one selected checkpoint into an empty thread without changing the source thread." +--- + +A user wants to try a new idea from an older workspace. The original thread +must stay unchanged. `snapshots.fork` copies one selected checkpoint into an +empty destination thread. + +The method copies the checkpoint that you pass. It does not copy the latest +checkpoint unless that id is the one you pass. + +Your checkpoint store must implement `forkFromCheckpoint`. That operation must +copy these items in one step: + +- The selected checkpoint. +- The conversation. +- The destination head. +- The blob reference counts. + +The store must reject a destination thread that already has persisted state. + +Keep this route on the server. Make sure that the session can access the source +thread. Make sure that the session can create the destination thread. Then call +`snapshots.fork`. + +```ts +import { requireSession } from './auth' +import { snapshots } from './sandbox-server' + +export async function POST(request: Request) { + const session = await requireSession(request) + const { threadId, checkpointId, destinationThreadId } = await request.json() + + if ( + typeof threadId !== 'string' || + typeof checkpointId !== 'string' || + typeof destinationThreadId !== 'string' + ) { + return new Response('Invalid request', { status: 400 }) + } + if (!(await session.canAccessThread(threadId))) { + return new Response('Not found', { status: 404 }) + } + if (!(await session.canCreateThread(destinationThreadId))) { + return new Response('Not found', { status: 404 }) + } + + const checkpoint = await snapshots.fork({ + threadId, + checkpointId, + destinationThreadId, + }) + + return Response.json({ checkpointId: checkpoint.id }) +} +``` + +The client sends its request to this route. It does not call `snapshots.fork`. + +```ts +export async function forkCheckpoint( + threadId: string, + checkpointId: string, + destinationThreadId: string, +) { + const response = await fetch('/api/snapshots/fork', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + threadId, + checkpointId, + destinationThreadId, + }), + }) + if (!response.ok) throw new Error('Could not fork checkpoint') + return response.json() +} +``` + +Use the same authorization rule for both threads. A client-selected checkpoint +id is not proof of access. + +If you use SQLite, put the fork in one transaction. That transaction must copy +the source conversation and reject a destination thread that is not empty. See +[Keep Files After Reload](./portable-snapshots-configure). diff --git a/docs/sandbox/portable-snapshots-safety.md b/docs/sandbox/portable-snapshots-safety.md new file mode 100644 index 0000000000..707457cc0d --- /dev/null +++ b/docs/sandbox/portable-snapshots-safety.md @@ -0,0 +1,68 @@ +--- +title: What a Snapshot Stores +id: portable-snapshots-safety +order: 15 +description: "See which workspace paths a portable snapshot stores, and how to replace the default policy." +--- + +You do not want secrets, git metadata, or install trees in durable storage. The +default snapshot policy excludes those paths. When a capture or restore finds +an unsafe filesystem entry, it stops. + +Portable snapshots store regular files and directories only. When a capture or +restore finds a symlink, an executable file, or a special filesystem entry, it +fails. + +## Default exclusions + +The default policy excludes these path segments at every depth: + +- `.git` +- `node_modules` +- `.env*` + +It also excludes these exact paths: + +- The projection marker, `.tanstack-projected-`, at the + workspace root only. +- `CLAUDE.md` and `GEMINI.md` at the workspace root. +- Direct `.claude/skills/`, `.codex/skills/`, and + `.grok/skills/` paths. + +These exclusions use paths for regular files and copied files too. + +## Replace the default policy + +A custom policy replaces the default exclusions. If you pass only `redact` or +`include`, the capture includes `.env`, `.git`, and `node_modules` unless you +copy the default policy first. + +```ts +import { defaultSandboxSnapshotPolicy } from '@tanstack/ai-sandbox' + +const policy = { + ...defaultSandboxSnapshotPolicy(), + redact({ bytes }: { bytes: Uint8Array }) { + return bytes + }, +} +``` + +When you create the snapshots object, pass `policy`. See +[Keep Files After Reload](./portable-snapshots-configure). + +The exact projection marker for the workspace stays protected. A custom policy +cannot capture or restore that marker. + +## Secrets + +Resolved secret values are replaced with zero bytes before the content is +hashed or stored. + +## Restore safety + +An invalid manifest, a missing blob, or changed blob content stops the restore +before it writes the workspace. The failed private sandbox is then discarded. +Your existing resumed sandbox stays unchanged. + +See [Providers](./providers) for provider-native snapshot and resume support. diff --git a/docs/sandbox/portable-snapshots-save.md b/docs/sandbox/portable-snapshots-save.md new file mode 100644 index 0000000000..67175d7583 --- /dev/null +++ b/docs/sandbox/portable-snapshots-save.md @@ -0,0 +1,73 @@ +--- +title: Save a Named Version +id: portable-snapshots-save +order: 12 +description: "Let a user mark one live sandbox workspace as a named checkpoint." +--- + +Automatic saves protect each completed run. A user can also mark one workspace +state, such as a release candidate. `snapshots.save` captures that live +sandbox into a named checkpoint. + +The method needs a live, reusable sandbox for the thread. A lifecycle with +`reuse: 'none'` cannot create a named checkpoint. + +Keep this route on the server. Read the owner from the session. If the owner +cannot access the thread, return 404. Then call `snapshots.save`. + +When you create the snapshots object, bind `sandbox` and `instances`. You can +also pass them on `save`. Use the same instance store that `withSandbox` uses. +See [Keep Files After Reload](./portable-snapshots-configure). + +```ts +import { requireSession } from './auth' +import { snapshots } from './sandbox-server' + +export async function POST(request: Request) { + const session = await requireSession(request) + const { threadId, runId, label } = await request.json() + + if ( + typeof threadId !== 'string' || + typeof runId !== 'string' || + typeof label !== 'string' + ) { + return new Response('Invalid request', { status: 400 }) + } + if (!(await session.canAccessThread(threadId))) { + return new Response('Not found', { status: 404 }) + } + + const checkpoint = await snapshots.save({ + threadId, + runId, + label, + tenant: { userId: session.userId }, + }) + return Response.json({ checkpointId: checkpoint.id, label: checkpoint.label }) +} +``` + +The client sends its request to this route. It does not call `snapshots.save` +or read the persistence stores. + +```ts +export async function saveCheckpoint( + threadId: string, + runId: string, + label: string, +) { + const response = await fetch('/api/snapshots/save', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ threadId, runId, label }), + }) + if (!response.ok) throw new Error('Could not save checkpoint') + return response.json() +} +``` + +A named checkpoint stays available for a read or a selected fork. Automatic +restore still uses the latest checkpoint. See +[Branch From a Version](./portable-snapshots-fork) to copy one selected +checkpoint into a new thread. diff --git a/docs/sandbox/portable-snapshots.md b/docs/sandbox/portable-snapshots.md index 14eeefda3a..4806664747 100644 --- a/docs/sandbox/portable-snapshots.md +++ b/docs/sandbox/portable-snapshots.md @@ -2,391 +2,65 @@ title: Portable Sandbox Snapshots id: portable-snapshots order: 10 -description: "Store a completed sandbox workspace as durable files, artifacts, and conversation data, then rebuild it in a new sandbox." +description: "Keep completed sandbox files after the provider sandbox is gone, then rebuild them in a new sandbox." --- -An agent can finish work in a sandbox, then the sandbox can disappear. You need -the same files when the page reloads or when a later run starts. Portable -sandbox snapshots save the completed workspace in your persistence stores. +An agent can finish work in a sandbox, then the sandbox can disappear. A reload +or a later run starts with an empty workspace. Portable snapshots store the +finished files in your persistence. A later run restores them into a new +sandbox. -This feature saves a checkpoint after each successful terminal run. A later -run restores that checkpoint into a new private sandbox before application code -can use the sandbox. You can also save a named checkpoint and fork from one -selected checkpoint. +This work runs on the server. Your client calls routes that you own. -The snapshot helpers run on the server. Your client calls routes that you own. -Do not expose a checkpoint, thread, or artifact id as authorization. +## Pick your path -## Configure snapshots +- You want files to return after a reload. Read + [Keep Files After Reload](./portable-snapshots-configure). +- You want a user to mark one version. Read + [Save a Named Version](./portable-snapshots-save). +- You want a user to branch from one version. Read + [Branch From a Version](./portable-snapshots-fork). +- You want a user to download a generated file. Read + [Send a Frozen File](./portable-snapshots-artifacts). +- You want to control which paths are stored. Read + [What a Snapshot Stores](./portable-snapshots-safety). -Create one persistence value and use that exact value in both middleware calls. -Put `withPersistence` before `withSandbox`. - -```ts -import { chat } from '@tanstack/ai' -import { grokBuildText } from '@tanstack/ai-grok-build' -import { withPersistence } from '@tanstack/ai-persistence' -import { - defineSandbox, - defineWorkspace, - InMemorySandboxInstanceStore, - memorySandboxSnapshots, - withSandbox, -} from '@tanstack/ai-sandbox' -import { dockerSandbox } from '@tanstack/ai-sandbox-docker' - -const instances = new InMemorySandboxInstanceStore() -const userId = 'user-123' // Read this from the server session. - -const sandbox = defineSandbox({ - id: 'app-builder', - provider: dockerSandbox({ image: 'node:22' }), - workspace: defineWorkspace({ source: { type: 'none' } }), - lifecycle: { reuse: 'thread' }, -}) - -const snapshots = await memorySandboxSnapshots({ - sandbox, - instances, -}) - -const result = chat({ - threadId: 'app-thread', - context: { userId }, - adapter: grokBuildText('grok-build'), - messages: [{ role: 'user', content: 'Create a landing page.' }], - middleware: [ - withPersistence(snapshots.persistence), - withSandbox(sandbox, { - instances, - snapshots, - }), - ], -}) - -void result -``` - -`memorySandboxSnapshots()` creates the in-memory persistence, checkpoint store, -and snapshot methods as one object. It does not load `@tanstack/ai-persistence` -at runtime. - -You can bind `sandbox`, `instances`, `tenant`, and `locks` at create time. A -named save can override those values. - -If you already have a persistence object, pass that same object to -`createSandboxSnapshots`. Use that same object in `withPersistence`. - -```ts -import { createSandboxSnapshots } from '@tanstack/ai-sandbox' -import { checkpoints, instances, persistence, sandbox } from './sandbox-server' - -const snapshots = createSandboxSnapshots({ - persistence, - checkpoints, - sandbox, - instances, -}) -``` - -Keep `instances` in the same server module as this middleware. A named save -must use this same store. Pass the session `userId` in `context` for every run. -Pass that same user id as `tenant.userId` on `snapshots.save`. - -Use a durable persistence implementation and checkpoint store in production. -The memory factory is useful for local development and examples only. - -The optional `policy` controls which workspace paths become files in a -checkpoint. See [Snapshot safety](#snapshot-safety) before you replace the -default policy. - -## Save a named checkpoint - -Automatic saves protect each completed run. Use a named save when a user marks -one workspace state, such as a release candidate. The method requires a live, -reusable sandbox for the thread. A lifecycle with `reuse: 'none'` cannot create -a named checkpoint. - -Keep this route on the server. Derive the owner from the session. Then make sure -that the owner can access the thread before you call the helper. - -```ts -import { requireSession } from './auth' -import { snapshots } from './sandbox-server' - -export async function POST(request: Request) { - const session = await requireSession(request) - const { threadId, runId, label } = await request.json() - - if ( - typeof threadId !== 'string' || - typeof runId !== 'string' || - typeof label !== 'string' - ) { - return new Response('Invalid request', { status: 400 }) - } - if (!(await session.canAccessThread(threadId))) { - return new Response('Not found', { status: 404 }) - } - - const checkpoint = await snapshots.save({ - threadId, - runId, - label, - tenant: { userId: session.userId }, - }) - return Response.json({ checkpointId: checkpoint.id, label: checkpoint.label }) -} -``` - -The client sends its request to this route. It does not call the helper or use -the persistence stores directly. - -```ts -export async function saveCheckpoint( - threadId: string, - runId: string, - label: string, -) { - const response = await fetch('/api/snapshots/save', { - method: 'POST', - headers: { 'content-type': 'application/json' }, - body: JSON.stringify({ threadId, runId, label }), - }) - if (!response.ok) throw new Error('Could not save checkpoint') - return response.json() -} -``` - -## Fork from a selected checkpoint - -Forking copies the selected checkpoint into an empty destination thread. It does -not fork from the latest checkpoint unless you pass that checkpoint id. The -source thread and its messages remain unchanged. - -Your checkpoint store must implement `forkFromCheckpoint`. The operation must -atomically copy the selected checkpoint, the conversation, the head, and blob -reference counts. It must reject a non-empty destination thread. - -```ts -import { requireSession } from './auth' -import { snapshots } from './sandbox-server' - -export async function POST(request: Request) { - const session = await requireSession(request) - const { threadId, checkpointId, destinationThreadId } = await request.json() - - if ( - typeof threadId !== 'string' || - typeof checkpointId !== 'string' || - typeof destinationThreadId !== 'string' - ) { - return new Response('Invalid request', { status: 400 }) - } - if (!(await session.canAccessThread(threadId))) { - return new Response('Not found', { status: 404 }) - } - if (!(await session.canCreateThread(destinationThreadId))) { - return new Response('Not found', { status: 404 }) - } - - const checkpoint = await snapshots.fork({ - threadId, - checkpointId, - destinationThreadId, - }) - - return Response.json({ checkpointId: checkpoint.id }) -} -``` - -Use the same server authorization boundary for both source and destination -threads. Do not accept a client-selected checkpoint as proof of access. - -```ts -export async function forkCheckpoint( - threadId: string, - checkpointId: string, - destinationThreadId: string, -) { - const response = await fetch('/api/snapshots/fork', { - method: 'POST', - headers: { 'content-type': 'application/json' }, - body: JSON.stringify({ - threadId, - checkpointId, - destinationThreadId, - }), - }) - if (!response.ok) throw new Error('Could not fork checkpoint') - return response.json() -} -``` - -## Read a snapshot artifact - -`snapshots.readArtifact` reads copied artifact bytes from one checkpoint. It -checks that the checkpoint belongs to the supplied thread. Your route must still -authorize that thread before it calls the method. The method returns metadata -and `Uint8Array` bytes. It does not create an HTTP response. - -```ts -import { requireSession } from './auth' -import { snapshots } from './sandbox-server' - -export async function GET(request: Request) { - const session = await requireSession(request) - const url = new URL(request.url) - const threadId = url.searchParams.get('threadId') - const checkpointId = url.searchParams.get('checkpointId') - const artifactId = url.searchParams.get('artifactId') - - if (!threadId || !checkpointId || !artifactId) { - return new Response('Not found', { status: 404 }) - } - if (!(await session.canAccessThread(threadId))) { - return new Response('Not found', { status: 404 }) - } - - const { artifact, bytes } = await snapshots.readArtifact({ - threadId, - checkpointId, - artifactId, - }) - return new Response(bytes.slice(), { - headers: { - 'content-type': artifact.mimeType, - 'content-length': String(artifact.size), - }, - }) -} -``` - -The client can use the authorized route as an artifact URL. It must not read the -blob store or call `snapshots.readArtifact` in the browser. - -```ts -export function snapshotArtifactUrl( - threadId: string, - checkpointId: string, - artifactId: string, -) { - const query = new URLSearchParams({ threadId, checkpointId, artifactId }) - return `/api/snapshots/artifact?${query}` -} -``` - -## What completes and restores +## What a checkpoint holds After a successful terminal run, the middleware waits for persistence to save -the conversation. It then creates one immutable checkpoint for the thread. +the conversation. It then writes one checkpoint for the thread. -The checkpoint contains: +A checkpoint holds: - Regular workspace files. - Empty directories. - Generated artifacts that already belong to the thread. - The saved conversation for the thread. -File data and copied artifact data use separate content-addressed blob -namespaces. Equal file data deduplicates with file data. Equal artifact data -deduplicates with artifact data. The system does not delete unused blobs -automatically yet. - -On a later run, the middleware uses the latest checkpoint only when it has a -new private sandbox. It restores the files after bootstrap and before the -sandbox is exposed to hooks or the harness. - -A live resumed sandbox is never overwritten. Provider-native snapshots can -make bootstrap faster. Portable checkpoints rebuild the durable workspace when -there is no live sandbox to resume. See [Lifecycle & Snapshots](./lifecycle) -for provider-native snapshot behavior. - -Portable snapshots do not restore into a live sandbox. The next private sandbox -gets the latest checkpoint after bootstrap. A named checkpoint remains available -for reading or for a selected fork. It does not change automatic restore. - -The conversation comes from the durable message store, not from the sandbox -journal. The journal remains a run-output log. See [The Run Journal](./journal) -when you need to replay agent output. - -## Snapshot safety - -Portable snapshots support regular files and directories only. A capture or -restore fails safely when it finds a symlink, an executable file, or a special -filesystem entry. - -The default policy excludes these path segments at every depth: - -- `.git` -- `node_modules` -- `.env*` - -It also excludes these exact paths: - -- The projection marker, `.tanstack-projected-`, at the - workspace root only. -- `CLAUDE.md` and `GEMINI.md` at the workspace root. -- Direct `.claude/skills/`, `.codex/skills/`, and - `.grok/skills/` paths. - -These exclusions use paths even when an entry is a regular file or a copied -file. A custom policy replaces the default exclusions. If you pass only -`redact` or `include`, `.env`, `.git`, and `node_modules` are captured unless -you copy the default policy first: - -```ts -import { defaultSandboxSnapshotPolicy } from '@tanstack/ai-sandbox' - -const policy = { - ...defaultSandboxSnapshotPolicy(), - redact({ bytes }: { bytes: Uint8Array }) { - return bytes - }, -} -``` - -The exact projection marker remains protected for the workspace. - -Resolved secret values are replaced with zero bytes before their content is -hashed or stored. A custom policy cannot capture or restore the exact projection -marker for the workspace. - -An invalid manifest, a missing blob, or changed blob content stops the restore -before it writes the workspace. The failed private sandbox is then discarded. -Your existing resumed sandbox remains unchanged. - -## Durable SQLite store - -`memorySandboxSnapshots()` is for tests and local examples. A production store -needs durable message, artifact, and blob stores, plus a durable checkpoint -store. The SQLite example exports `sqliteSandboxSnapshots()` for a Node 22.5+ -server. It is an example adapter, not a package export. +File data and artifact data use separate content-addressed blob namespaces. +Equal file data shares one file blob. Equal artifact data shares one artifact +blob. Unused blobs stay until you delete them. -Use one SQLite transaction for every checkpoint write, head update, and blob -reference count update. Use one transaction for a fork. The fork transaction -must also copy the source conversation and reject a destination thread that has -any persisted state. A partial transaction can create a checkpoint that points -to missing data or a wrong blob reference count. +## When a later run restores files -## Test route +A later run restores the latest checkpoint only into a new private sandbox. The +restore runs after bootstrap and before hooks or the harness see the sandbox. -`testing/e2e/src/routes/api.sandbox-file-persistence.ts` is a test-only route. -It uses in-memory stores and a fake provider. Do not copy this route into an -application. Use your authenticated server routes and durable stores instead. +A live resumed sandbox keeps its current files. Portable snapshots do not write +into that sandbox. Provider-native snapshots make bootstrap faster. See +[Lifecycle & Snapshots](./lifecycle). -## Operations +A named checkpoint stays available for a read or a selected fork. Automatic +restore still uses the latest checkpoint. -Each thread has one checkpoint writer lease. A second run for the same thread -gets a writer conflict while the existing lease is active. The middleware -renews its lease while the run is active. +The conversation comes from the durable message store. The sandbox journal is a +run-output log. When you need to replay agent output, see +[The Run Journal](./journal). -Pause and detach paths release the lease. They do not publish a partial -checkpoint. If the writer loses its lease, the middleware does not publish the -checkpoint. A later successful run can create a new checkpoint. +## Who calls the methods -Portable snapshots work with [Instance Durability](./durability). Instance -durability finds a provider sandbox across server processes. Portable snapshots -rebuild the workspace when that sandbox is unavailable. +`createSandboxSnapshots` and `memorySandboxSnapshots` return one object. That +object has `save`, `fork`, and `readArtifact`. Call those methods on the server +after you make sure that the session can access the thread. -See [Providers](./providers) for provider-native snapshot and resume support. +Do not treat a checkpoint id, a thread id, or an artifact id as proof of access. diff --git a/docs/sandbox/providers.md b/docs/sandbox/providers.md index 3cb342cbaf..14f849996b 100644 --- a/docs/sandbox/providers.md +++ b/docs/sandbox/providers.md @@ -13,7 +13,7 @@ snapshot/resume behaviour you need; the rest of your sandbox definition stays th same. Provider-native snapshots and resume keep or recreate provider state. They can -reduce bootstrap time. [Portable Sandbox Snapshots](./portable-snapshots) store +reduce bootstrap time. [Portable Snapshots](./portable-snapshots) store completed workspace data in your application persistence for reconstruction. > The provider is _where_ the agent runs. For _which_ agent runs (Grok Build, diff --git a/packages/ai-sandbox/README.md b/packages/ai-sandbox/README.md index f5ff37f9ac..4c6f4f7141 100644 --- a/packages/ai-sandbox/README.md +++ b/packages/ai-sandbox/README.md @@ -140,8 +140,8 @@ lifecycle: { ### Portable snapshots Use portable snapshots when a later run must rebuild completed files after the -provider sandbox is gone. Configure `withPersistence` first, then pass the same -persistence value and a checkpoint store to `withSandbox`. +provider sandbox is gone. Create one snapshots object, pass it to +`withPersistence` and `withSandbox`, and put `withPersistence` first. ```typescript import { withPersistence } from '@tanstack/ai-persistence' @@ -157,9 +157,9 @@ const middleware = [ A successful terminal run saves regular files, empty directories, saved conversation data, and thread artifacts. Restore runs only in a new private -sandbox. It never overwrites a live resumed sandbox. Read the -[Portable Sandbox Snapshots guide](https://tanstack.com/ai/latest/docs/sandbox/portable-snapshots) -for storage, safety, and lease details. +sandbox. A live resumed sandbox keeps its current files. Read +[Keep Files After Reload](https://tanstack.com/ai/latest/docs/sandbox/portable-snapshots-configure) +for the full server setup. ### Secrets diff --git a/packages/ai-sandbox/skills/ai-sandbox/SKILL.md b/packages/ai-sandbox/skills/ai-sandbox/SKILL.md index 8a2dd5cd16..dd8af856f5 100644 --- a/packages/ai-sandbox/skills/ai-sandbox/SKILL.md +++ b/packages/ai-sandbox/skills/ai-sandbox/SKILL.md @@ -12,7 +12,8 @@ description: > terminal run with withPersistence before withSandbox and memorySandboxSnapshots for local examples. It covers named saves with snapshots.save, selected-checkpoint forks with snapshots.fork, and - authorized artifact reads with snapshots.readArtifact. It covers defineWorkspace + authorized artifact reads with snapshots.readArtifact. See + docs/sandbox/portable-snapshots.md. It covers defineWorkspace (git/setup/scripts/skills/secrets/ instructions/plugins), defineSandboxPolicy (allow/ask/deny), lifecycle/resume, the SandboxHandle (fs/git/process/ports), capability tokens, defineSandbox @@ -224,8 +225,14 @@ thread has one writer lease. Pause and detach release the lease without a partial checkpoint. Blob retention is manual because there is no automatic garbage collection yet. -Read `docs/sandbox/portable-snapshots.md` for the full server-only setup and -the restore safety rules. +Read these pages for the server-only setup: + +- `docs/sandbox/portable-snapshots.md` +- `docs/sandbox/portable-snapshots-configure.md` +- `docs/sandbox/portable-snapshots-save.md` +- `docs/sandbox/portable-snapshots-fork.md` +- `docs/sandbox/portable-snapshots-artifacts.md` +- `docs/sandbox/portable-snapshots-safety.md` For a user-marked workspace state, call `snapshots.save` on the server. Bind `sandbox` and `instances` at create time, or pass them on `save`. The call @@ -238,9 +245,9 @@ checkpoint id, and destination thread id. The store must implement atomic selected snapshot, not the latest snapshot. To send a checkpoint artifact, call `snapshots.readArtifact` on the server. -First authorize the caller for the supplied thread. The method checks that the -checkpoint belongs to that thread, then returns its metadata and bytes. It does -not authorize a caller or create an HTTP response. +First authorize the caller for the supplied thread. The method makes sure that +the checkpoint belongs to that thread, then returns its metadata and bytes. It +does not authorize a caller or create an HTTP response. For a SQLite checkpoint store, use one transaction for a checkpoint write, its head update, and every blob reference update. Use one transaction for a fork, From e100fe6376fc4c5b51174d6f9adb7386a7279871 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Fri, 14 Aug 2026 16:22:21 +0200 Subject: [PATCH 06/12] docs(persistence): show sandbox snapshot stores in the adapter table Mark messages, artifacts, and blobs as required to rebuild sandbox files. Spell out create vs reuse paths. --- docs/config.json | 9 +++-- docs/persistence/build-your-own-adapter.md | 42 ++++++++++++-------- docs/persistence/controls.md | 3 ++ docs/persistence/overview.md | 1 + docs/persistence/store-reference.md | 4 +- docs/sandbox/portable-snapshots-configure.md | 26 ++++++++++-- docs/sandbox/portable-snapshots.md | 23 ++++++----- 7 files changed, 72 insertions(+), 36 deletions(-) diff --git a/docs/config.json b/docs/config.json index 8cf461935c..e3f805e8c3 100644 --- a/docs/config.json +++ b/docs/config.json @@ -261,12 +261,14 @@ { "label": "Controls", "to": "persistence/controls", - "addedAt": "2026-08-04" + "addedAt": "2026-08-04", + "updatedAt": "2026-08-14" }, { "label": "Build Your Own Adapter", "to": "persistence/build-your-own-adapter", - "addedAt": "2026-08-04" + "addedAt": "2026-08-04", + "updatedAt": "2026-08-14" }, { "label": "Migrations", @@ -566,7 +568,8 @@ { "label": "Keep Files After Reload", "to": "sandbox/portable-snapshots-configure", - "addedAt": "2026-08-14" + "addedAt": "2026-08-14", + "updatedAt": "2026-08-14" }, { "label": "Save a Named Version", diff --git a/docs/persistence/build-your-own-adapter.md b/docs/persistence/build-your-own-adapter.md index 8cbbe457de..77fee806dd 100644 --- a/docs/persistence/build-your-own-adapter.md +++ b/docs/persistence/build-your-own-adapter.md @@ -67,23 +67,29 @@ object inline so you never annotate it by hand. Each store switches on one capability. Find your column and implement the rows marked with a tick: -| Store | Save the transcript | Rejoin a run after reload | Durable approvals | App key/value | Persist generation runs | Keep generated files | -| --- | :-: | :-: | :-: | :-: | :-: | :-: | -| `messages` | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | -| `runs` | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | -| `interrupts` | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | -| `metadata` | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | -| `generationRuns` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | -| `artifacts` | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | -| `blobs` | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | - -- **Columns stack.** Durable approvals and generated files means the union of both. -- **Two pairs cannot be split.** `interrupts` needs `runs`, and `artifacts` needs - `blobs`. -- **The generation stores feed `withGenerationPersistence`** instead, and need none of - the chat stores. See [Generation persistence](./generation-persistence). - -The common production shape is `messages` + `runs` + `interrupts`. +| Store | Save the transcript | Rejoin a run after reload | Durable approvals | App key/value | Persist generation runs | Keep generated files | Rebuild sandbox files | +| --- | :-: | :-: | :-: | :-: | :-: | :-: | :-: | +| `messages` | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | +| `runs` | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | +| `interrupts` | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | +| `metadata` | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | +| `generationRuns` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | +| `artifacts` | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | +| `blobs` | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | + +- **Columns stack.** Chat plus sandbox files means `messages` + `artifacts` + + `blobs`. Durable approvals plus generated files means the union of both + columns. +- **Two pairs cannot be split.** `interrupts` needs `runs`, and `artifacts` + needs `blobs`. +- **Generation runs need none of the chat stores.** See + [Generation persistence](./generation-persistence). +- **Sandbox files need a checkpoint store too.** That store lives on + `@tanstack/ai-sandbox`, not in this table. See + [Keep Files After Reload](../sandbox/portable-snapshots-configure). + +The common production shape is `messages` + `runs` + `interrupts`. When you +keep generated files or rebuild sandbox files, add `artifacts` and `blobs`. You can also own only part of it. Put `messages` and `runs` in your database and fill the rest from somewhere else with `composePersistence`: @@ -158,6 +164,8 @@ for `withPersistence`, and with the generation stores for artifacts and blobs. - [Build a sandbox adapter](./build-a-sandbox-adapter): the sandbox instance store, and what a durable sandboxed run adds to `runs`. Only if you run sandboxes. +- [Keep Files After Reload](../sandbox/portable-snapshots-configure): reuse this + adapter for portable snapshots. You need `messages`, `artifacts`, and `blobs`. - [Store reference](./store-reference): every signature and invariant, and how the records relate. - [Controls](./controls): compose stores from different systems. diff --git a/docs/persistence/controls.md b/docs/persistence/controls.md index e74d9dff7c..d69a62faa6 100644 --- a/docs/persistence/controls.md +++ b/docs/persistence/controls.md @@ -90,6 +90,9 @@ values arrive from untyped JavaScript. - `withPersistence` requires `messages`. - `interrupts` requires `runs`: an interrupt record is scoped to a run. - `withGenerationPersistence` requires `generationRuns`. +- Portable sandbox snapshots require `messages`, `artifacts`, and `blobs` on + the same persistence object. See + [Keep Files After Reload](../sandbox/portable-snapshots-configure). To define a partial backend directly rather than by composing, use `defineAIPersistence({ stores: { ... } })` and pass only the stores you have. diff --git a/docs/persistence/overview.md b/docs/persistence/overview.md index b9568fcf79..8e0cf080ca 100644 --- a/docs/persistence/overview.md +++ b/docs/persistence/overview.md @@ -168,6 +168,7 @@ To make the `POST` resumable too, hand the same adapter to the response: | A reload mid-answer to pick the answer back up | Steps 1, 2 and 3 | | A dropped socket to resume with the page still open | [Resumable Streams](../resumable-streams/overview) alone | | To pause for a human approval and resume it days later | Step 1 with an `interrupts` store | +| Sandbox files to come back after the provider sandbox is gone | [Keep Files After Reload](../sandbox/portable-snapshots-configure) | ## Where to go next diff --git a/docs/persistence/store-reference.md b/docs/persistence/store-reference.md index 41e24e1c81..0c76b69d93 100644 --- a/docs/persistence/store-reference.md +++ b/docs/persistence/store-reference.md @@ -16,8 +16,8 @@ there is no separate enable list. | `interrupts` | Pending, resolved or cancelled human waits. Needs `runs`. | `withPersistence` | | `metadata` | App and integration key/value state. | `withPersistence` | | `generationRuns` | Generation run status and result metadata, keyed by its own `runId`. | `withGenerationPersistence`, required | -| `artifacts` | Generated-file metadata. Needs `blobs`. | `withGenerationPersistence` | -| `blobs` | The generated bytes. Needs `artifacts`. | `withGenerationPersistence` | +| `artifacts` | File metadata. Needs `blobs`. | `withGenerationPersistence`, portable snapshots | +| `blobs` | File bytes. Needs `artifacts`. | `withGenerationPersistence`, portable snapshots | Named groupings of the chat stores (`ChatTranscriptStores`, `ChatPersistenceStores`, `ChatWithInterruptsStores`) are covered in [Controls](./controls). diff --git a/docs/sandbox/portable-snapshots-configure.md b/docs/sandbox/portable-snapshots-configure.md index 31ff72ad6d..31c86ae0ff 100644 --- a/docs/sandbox/portable-snapshots-configure.md +++ b/docs/sandbox/portable-snapshots-configure.md @@ -15,7 +15,13 @@ this page, `chat()` writes and restores those checkpoints. Create one persistence value. Pass that exact value to `withPersistence` and to the snapshots object. Put `withPersistence` before `withSandbox`. -## Start with the memory factory +This page is enough for automatic save and restore. When you need a named +version, a fork, or a download, add +[Save a Named Version](./portable-snapshots-save), +[Branch From a Version](./portable-snapshots-fork), or +[Send a Frozen File](./portable-snapshots-artifacts). + +## Create new persistence Use `memorySandboxSnapshots` for local development. It creates persistence, a checkpoint store, and the snapshot methods as one object. It does not load @@ -76,13 +82,25 @@ must use this same instance store. Pass the session `userId` in `context` for every run. Pass that same user id as `tenant.userId` on `snapshots.save`. -## Add snapshots to existing persistence +## Reuse existing persistence If `withPersistence` already uses a persistence object, pass that same object to `createSandboxSnapshots`. Do not create a second message store. -The persistence object must include `stores.messages`, `stores.artifacts`, and -`stores.blobs`. +The persistence object must include these stores: + +- `messages` +- `artifacts` (with `listForThread`) +- `blobs` + +You also need a checkpoint store. That store is not a persistence store. + +If you already keep generated files, you already have `artifacts` and `blobs`. +Use those same stores. + +If you only have `messages`, add `artifacts` and `blobs` to that same adapter. +See [Which stores do you need?](../persistence/build-your-own-adapter#which-stores-do-you-need) +and [Build a generation adapter](../persistence/build-your-own-generation-adapter). ```ts import { chat } from '@tanstack/ai' diff --git a/docs/sandbox/portable-snapshots.md b/docs/sandbox/portable-snapshots.md index 4806664747..f280ebe3cf 100644 --- a/docs/sandbox/portable-snapshots.md +++ b/docs/sandbox/portable-snapshots.md @@ -14,16 +14,19 @@ This work runs on the server. Your client calls routes that you own. ## Pick your path -- You want files to return after a reload. Read - [Keep Files After Reload](./portable-snapshots-configure). -- You want a user to mark one version. Read - [Save a Named Version](./portable-snapshots-save). -- You want a user to branch from one version. Read - [Branch From a Version](./portable-snapshots-fork). -- You want a user to download a generated file. Read - [Send a Frozen File](./portable-snapshots-artifacts). -- You want to control which paths are stored. Read - [What a Snapshot Stores](./portable-snapshots-safety). +Start with persistence. Then add a product page only when you need that action. + +| You have | You want | Required pages | +| --- | --- | --- | +| No persistence yet | Files come back after a reload | [Keep Files After Reload](./portable-snapshots-configure#create-new-persistence) | +| Chat persistence already | Files come back after a reload | [Keep Files After Reload](./portable-snapshots-configure#reuse-existing-persistence) | +| Snapshots already wired | A user marks one version | [Save a Named Version](./portable-snapshots-save) | +| Snapshots already wired | A user branches from one version | [Branch From a Version](./portable-snapshots-fork) | +| Snapshots already wired | A user downloads a generated file | [Send a Frozen File](./portable-snapshots-artifacts) | +| Snapshots already wired | Custom include or redact rules | [What a Snapshot Stores](./portable-snapshots-safety) | + +Automatic save and restore needs only the configure page. Save, fork, download, +and policy pages are extra. ## What a checkpoint holds From 803bb4154fac1984448d61015b1b547b282013d0 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Fri, 14 Aug 2026 16:57:06 +0200 Subject: [PATCH 07/12] feat(ai-sandbox): add createSnapshotTools bound to the route thread Host tools save, fork, and read this thread. The model cannot pass thread ids. createThreadId mints each destination thread. --- .changeset/fuzzy-snapshots-build.md | 2 +- docs/config.json | 11 +- docs/sandbox/portable-snapshots-fork.md | 3 + docs/sandbox/portable-snapshots-save.md | 3 + docs/sandbox/portable-snapshots-tools.md | 77 ++++++ docs/sandbox/portable-snapshots.md | 1 + .../ai-sandbox/skills/ai-sandbox/SKILL.md | 1 + packages/ai-sandbox/src/index.ts | 2 + packages/ai-sandbox/src/snapshot-tools.ts | 208 ++++++++++++++ packages/ai-sandbox/src/snapshots.ts | 2 + .../tests/snapshot-operations.test-d.ts | 2 + .../ai-sandbox/tests/snapshot-tools.test.ts | 255 ++++++++++++++++++ 12 files changed, 564 insertions(+), 3 deletions(-) create mode 100644 docs/sandbox/portable-snapshots-tools.md create mode 100644 packages/ai-sandbox/src/snapshot-tools.ts create mode 100644 packages/ai-sandbox/tests/snapshot-tools.test.ts diff --git a/.changeset/fuzzy-snapshots-build.md b/.changeset/fuzzy-snapshots-build.md index 293776afbd..f6b0bc9faf 100644 --- a/.changeset/fuzzy-snapshots-build.md +++ b/.changeset/fuzzy-snapshots-build.md @@ -2,4 +2,4 @@ '@tanstack/ai-sandbox': minor --- -Add portable sandbox checkpoints. `createSandboxSnapshots` and `memorySandboxSnapshots` return one object with `save`, `fork`, and `readArtifact`. +Add portable sandbox checkpoints. `createSandboxSnapshots` and `memorySandboxSnapshots` return one object with `save`, `fork`, and `readArtifact`. `createSnapshotTools` turns those methods into host tools bound to the route `threadId`. diff --git a/docs/config.json b/docs/config.json index e3f805e8c3..7026808467 100644 --- a/docs/config.json +++ b/docs/config.json @@ -574,18 +574,25 @@ { "label": "Save a Named Version", "to": "sandbox/portable-snapshots-save", - "addedAt": "2026-08-14" + "addedAt": "2026-08-14", + "updatedAt": "2026-08-14" }, { "label": "Branch From a Version", "to": "sandbox/portable-snapshots-fork", - "addedAt": "2026-08-14" + "addedAt": "2026-08-14", + "updatedAt": "2026-08-14" }, { "label": "Send a Frozen File", "to": "sandbox/portable-snapshots-artifacts", "addedAt": "2026-08-14" }, + { + "label": "Let the Agent Save and Fork", + "to": "sandbox/portable-snapshots-tools", + "addedAt": "2026-08-14" + }, { "label": "What a Snapshot Stores", "to": "sandbox/portable-snapshots-safety", diff --git a/docs/sandbox/portable-snapshots-fork.md b/docs/sandbox/portable-snapshots-fork.md index fc438b0e14..943c15c5cb 100644 --- a/docs/sandbox/portable-snapshots-fork.md +++ b/docs/sandbox/portable-snapshots-fork.md @@ -86,3 +86,6 @@ id is not proof of access. If you use SQLite, put the fork in one transaction. That transaction must copy the source conversation and reject a destination thread that is not empty. See [Keep Files After Reload](./portable-snapshots-configure). + +When the agent must call fork itself, use +[Let the Agent Save and Fork](./portable-snapshots-tools). diff --git a/docs/sandbox/portable-snapshots-save.md b/docs/sandbox/portable-snapshots-save.md index 67175d7583..01f5f7b60c 100644 --- a/docs/sandbox/portable-snapshots-save.md +++ b/docs/sandbox/portable-snapshots-save.md @@ -71,3 +71,6 @@ A named checkpoint stays available for a read or a selected fork. Automatic restore still uses the latest checkpoint. See [Branch From a Version](./portable-snapshots-fork) to copy one selected checkpoint into a new thread. + +When the agent must call save itself, use +[Let the Agent Save and Fork](./portable-snapshots-tools). diff --git a/docs/sandbox/portable-snapshots-tools.md b/docs/sandbox/portable-snapshots-tools.md new file mode 100644 index 0000000000..b8567f2fb6 --- /dev/null +++ b/docs/sandbox/portable-snapshots-tools.md @@ -0,0 +1,77 @@ +--- +title: Let the Agent Save and Fork +id: portable-snapshots-tools +order: 16 +description: "Give a chat() run host tools that save and fork this thread without letting the model pick thread ids." +--- + +You want the agent to mark a version or open a new direction. The route already +knows `threadId` and `runId`. `createSnapshotTools` turns those values into +host tools that you spread into `chat()`. + +The model can pass a label or a checkpoint id on this thread. It cannot pass a +thread id. Your factory mints every new thread id. + +```ts +import { chat } from '@tanstack/ai' +import { grokBuildText } from '@tanstack/ai-grok-build' +import { withPersistence } from '@tanstack/ai-persistence' +import { createSnapshotTools, withSandbox } from '@tanstack/ai-sandbox' +import { instances, sandbox, snapshots } from './sandbox-server' + +export function POST(threadId: string, runId: string) { + return chat({ + threadId, + runId, + adapter: grokBuildText('grok-build'), + messages: [{ role: 'user', content: 'Save this, then try a dark theme.' }], + tools: [ + ...createSnapshotTools(snapshots, { + threadId, + runId, + createThreadId: () => crypto.randomUUID(), + onForked({ destinationThreadId }) { + void destinationThreadId + }, + }), + ], + middleware: [ + withPersistence(snapshots.persistence), + withSandbox(sandbox, { instances, snapshots }), + ], + }) +} +``` + +The tools run on the server. In a sandbox chat they are bridged back to the +host. See [Tools](./tools). + +## What each tool does + +- `save_sandbox_snapshot`: saves the live sandbox for the bound thread. The + model passes `label` only. +- `fork_sandbox_snapshot`: copies one checkpoint into a new empty thread. The + model can pass `checkpointId`. When it omits that id, the tool copies the + latest checkpoint. `createThreadId()` sets the destination thread id. +- `read_sandbox_snapshot_artifact`: returns artifact metadata for a checkpoint + on this thread. It does not return file bytes. Serve bytes from + [Send a Frozen File](./portable-snapshots-artifacts). + +`onForked` runs after a successful fork. If you want the new branch to work +right away, start `chat()` on `destinationThreadId` in that callback. + +## When save and fork run + +`save_sandbox_snapshot` takes the writer lease for this thread. If this +`chat()` already holds that lease through portable snapshots, the save fails +with `SANDBOX_SNAPSHOT_WRITER_CONFLICT`. Save from a planner thread, or after +this run ends. + +`fork_sandbox_snapshot` takes the writer lease on the **new** thread. A fork +during this run can succeed. It copies a saved checkpoint. It does not copy +files the agent is still writing. + +See [Keep Files After Reload](./portable-snapshots-configure) to create the +`snapshots` object. See [Save a Named Version](./portable-snapshots-save) and +[Branch From a Version](./portable-snapshots-fork) for routes that you call +yourself. diff --git a/docs/sandbox/portable-snapshots.md b/docs/sandbox/portable-snapshots.md index f280ebe3cf..35214d889c 100644 --- a/docs/sandbox/portable-snapshots.md +++ b/docs/sandbox/portable-snapshots.md @@ -23,6 +23,7 @@ Start with persistence. Then add a product page only when you need that action. | Snapshots already wired | A user marks one version | [Save a Named Version](./portable-snapshots-save) | | Snapshots already wired | A user branches from one version | [Branch From a Version](./portable-snapshots-fork) | | Snapshots already wired | A user downloads a generated file | [Send a Frozen File](./portable-snapshots-artifacts) | +| Snapshots already wired | The agent saves or forks this thread | [Let the Agent Save and Fork](./portable-snapshots-tools) | | Snapshots already wired | Custom include or redact rules | [What a Snapshot Stores](./portable-snapshots-safety) | Automatic save and restore needs only the configure page. Save, fork, download, diff --git a/packages/ai-sandbox/skills/ai-sandbox/SKILL.md b/packages/ai-sandbox/skills/ai-sandbox/SKILL.md index dd8af856f5..34bd33fed5 100644 --- a/packages/ai-sandbox/skills/ai-sandbox/SKILL.md +++ b/packages/ai-sandbox/skills/ai-sandbox/SKILL.md @@ -232,6 +232,7 @@ Read these pages for the server-only setup: - `docs/sandbox/portable-snapshots-save.md` - `docs/sandbox/portable-snapshots-fork.md` - `docs/sandbox/portable-snapshots-artifacts.md` +- `docs/sandbox/portable-snapshots-tools.md` - `docs/sandbox/portable-snapshots-safety.md` For a user-marked workspace state, call `snapshots.save` on the server. Bind diff --git a/packages/ai-sandbox/src/index.ts b/packages/ai-sandbox/src/index.ts index e4b9724b3f..e265eba978 100644 --- a/packages/ai-sandbox/src/index.ts +++ b/packages/ai-sandbox/src/index.ts @@ -79,6 +79,8 @@ export type { SaveSandboxSnapshotInput, SnapshotPersistence, } from './snapshot-operations' +export { createSnapshotTools } from './snapshot-tools' +export type { CreateSnapshotToolsOptions } from './snapshot-tools' // Workspace projection capability (provided by withSandbox, consumed by harness adapters) export { diff --git a/packages/ai-sandbox/src/snapshot-tools.ts b/packages/ai-sandbox/src/snapshot-tools.ts new file mode 100644 index 0000000000..8725e42abf --- /dev/null +++ b/packages/ai-sandbox/src/snapshot-tools.ts @@ -0,0 +1,208 @@ +import { toolDefinition } from '@tanstack/ai' +import { SandboxSnapshotError } from './snapshots' +import type { SandboxSnapshots } from './snapshot-operations' + +export interface CreateSnapshotToolsOptions { + threadId: string + runId: string + createThreadId: () => string + tenant?: { userId?: string; orgId?: string } + onForked?: (input: { + destinationThreadId: string + checkpointId: string + }) => void | Promise +} + +function field(value: unknown, key: string): unknown { + if (value === null || typeof value !== 'object') return undefined + return Reflect.get(value, key) +} + +function requiredString(value: unknown, key: string): string { + const candidate = field(value, key) + if (typeof candidate !== 'string' || candidate.length === 0) { + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_TOOL_INPUT', + `Snapshot tool requires a non-empty ${key}`, + ) + } + return candidate +} + +function optionalString(value: unknown, key: string): string | undefined { + const candidate = field(value, key) + if (candidate === undefined) return undefined + if (typeof candidate !== 'string' || candidate.length === 0) { + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_TOOL_INPUT', + `Snapshot tool ${key} must be a non-empty string when provided`, + ) + } + return candidate +} + +function requireIdentifier(value: string, label: string): string { + if (value.length === 0) { + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_TOOL_INPUT', + `${label} must be a non-empty string`, + ) + } + return value +} + +export function createSnapshotTools( + snapshots: SandboxSnapshots, + options: CreateSnapshotToolsOptions, +) { + const threadId = requireIdentifier(options.threadId, 'threadId') + const runId = requireIdentifier(options.runId, 'runId') + const createThreadId = options.createThreadId + const tenant = options.tenant + const onForked = options.onForked + if (typeof createThreadId !== 'function') { + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_INVALID_TOOL_INPUT', + 'createSnapshotTools requires createThreadId', + ) + } + + const save = toolDefinition({ + name: 'save_sandbox_snapshot', + description: + 'Save a named checkpoint of the current live sandbox for this thread. Do not pass a thread id.', + inputSchema: { + type: 'object', + properties: { + label: { + type: 'string', + description: 'A short name for this version, such as release-1.', + }, + }, + required: ['label'], + additionalProperties: false, + }, + outputSchema: { + type: 'object', + properties: { + checkpointId: { type: 'string' }, + label: { type: 'string' }, + threadId: { type: 'string' }, + }, + required: ['checkpointId', 'label', 'threadId'], + additionalProperties: false, + }, + }).server(async (input) => { + const label = requiredString(input, 'label') + const checkpoint = await snapshots.save({ + threadId, + runId, + label, + ...(tenant === undefined ? {} : { tenant }), + }) + return { + checkpointId: checkpoint.id, + label: checkpoint.label ?? label, + threadId, + } + }) + + const fork = toolDefinition({ + name: 'fork_sandbox_snapshot', + description: + 'Copy one checkpoint from this thread into a new empty thread. Omit checkpointId to copy the latest checkpoint. Do not pass thread ids.', + inputSchema: { + type: 'object', + properties: { + checkpointId: { + type: 'string', + description: + 'The checkpoint to copy. When omitted, the latest checkpoint is copied.', + }, + }, + additionalProperties: false, + }, + outputSchema: { + type: 'object', + properties: { + checkpointId: { type: 'string' }, + destinationThreadId: { type: 'string' }, + }, + required: ['checkpointId', 'destinationThreadId'], + additionalProperties: false, + }, + }).server(async (input) => { + const suppliedCheckpointId = optionalString(input, 'checkpointId') + const checkpointId = + suppliedCheckpointId ?? (await snapshots.checkpoints.getHead(threadId)) + if (checkpointId === null) { + throw new SandboxSnapshotError( + 'SANDBOX_SNAPSHOT_MISSING_CHECKPOINT', + 'This thread has no checkpoint to fork', + ) + } + const destinationThreadId = requireIdentifier( + createThreadId(), + 'destinationThreadId', + ) + const checkpoint = await snapshots.fork({ + threadId, + checkpointId, + destinationThreadId, + }) + if (onForked !== undefined) { + await onForked({ + destinationThreadId, + checkpointId: checkpoint.id, + }) + } + return { + checkpointId: checkpoint.id, + destinationThreadId, + } + }) + + const readArtifact = toolDefinition({ + name: 'read_sandbox_snapshot_artifact', + description: + 'Read metadata for one artifact on a checkpoint in this thread. Do not pass a thread id.', + inputSchema: { + type: 'object', + properties: { + checkpointId: { type: 'string' }, + artifactId: { type: 'string' }, + }, + required: ['checkpointId', 'artifactId'], + additionalProperties: false, + }, + outputSchema: { + type: 'object', + properties: { + artifactId: { type: 'string' }, + name: { type: 'string' }, + mimeType: { type: 'string' }, + size: { type: 'number' }, + createdAt: { type: 'number' }, + }, + required: ['artifactId', 'name', 'mimeType', 'size', 'createdAt'], + additionalProperties: false, + }, + }).server(async (input) => { + const checkpointId = requiredString(input, 'checkpointId') + const artifactId = requiredString(input, 'artifactId') + const resolved = await snapshots.readArtifact({ + threadId, + checkpointId, + artifactId, + }) + return { + artifactId: resolved.artifact.artifactId, + name: resolved.artifact.name, + mimeType: resolved.artifact.mimeType, + size: resolved.artifact.size, + createdAt: resolved.artifact.createdAt, + } + }) + + return [save, fork, readArtifact] as const +} diff --git a/packages/ai-sandbox/src/snapshots.ts b/packages/ai-sandbox/src/snapshots.ts index cdd4fb88fa..0d1aaf9fb3 100644 --- a/packages/ai-sandbox/src/snapshots.ts +++ b/packages/ai-sandbox/src/snapshots.ts @@ -37,11 +37,13 @@ export interface SandboxSnapshotBundle { } export type SandboxSnapshotErrorCode = + | 'SANDBOX_SNAPSHOT_INVALID_TOOL_INPUT' | 'SANDBOX_SNAPSHOT_MISSING_SANDBOX' | 'SANDBOX_SNAPSHOT_MISSING_INSTANCES' | 'SANDBOX_SNAPSHOT_MISSING_PERSISTENCE_STORES' | 'SANDBOX_SNAPSHOT_MISSING_REUSABLE_SANDBOX' | 'SANDBOX_SNAPSHOT_REUSE_NONE' + | 'SANDBOX_SNAPSHOT_MISSING_CHECKPOINT' | 'SANDBOX_SNAPSHOT_MISSING_CHECKPOINT_ARTIFACT' | 'SANDBOX_SNAPSHOT_FOREIGN_CHECKPOINT_ARTIFACT' | 'SANDBOX_SNAPSHOT_INVALID_ARTIFACT_BYTES' diff --git a/packages/ai-sandbox/tests/snapshot-operations.test-d.ts b/packages/ai-sandbox/tests/snapshot-operations.test-d.ts index f95f04010e..8e5fcb46b8 100644 --- a/packages/ai-sandbox/tests/snapshot-operations.test-d.ts +++ b/packages/ai-sandbox/tests/snapshot-operations.test-d.ts @@ -3,11 +3,13 @@ import { memorySandboxSnapshots, SandboxSnapshotError } from '../src' import type { SandboxSnapshotErrorCode, SandboxSnapshots } from '../src' type ExpectedSandboxSnapshotErrorCode = + | 'SANDBOX_SNAPSHOT_INVALID_TOOL_INPUT' | 'SANDBOX_SNAPSHOT_MISSING_SANDBOX' | 'SANDBOX_SNAPSHOT_MISSING_INSTANCES' | 'SANDBOX_SNAPSHOT_MISSING_PERSISTENCE_STORES' | 'SANDBOX_SNAPSHOT_MISSING_REUSABLE_SANDBOX' | 'SANDBOX_SNAPSHOT_REUSE_NONE' + | 'SANDBOX_SNAPSHOT_MISSING_CHECKPOINT' | 'SANDBOX_SNAPSHOT_MISSING_CHECKPOINT_ARTIFACT' | 'SANDBOX_SNAPSHOT_FOREIGN_CHECKPOINT_ARTIFACT' | 'SANDBOX_SNAPSHOT_INVALID_ARTIFACT_BYTES' diff --git a/packages/ai-sandbox/tests/snapshot-tools.test.ts b/packages/ai-sandbox/tests/snapshot-tools.test.ts new file mode 100644 index 0000000000..2461b0d1ac --- /dev/null +++ b/packages/ai-sandbox/tests/snapshot-tools.test.ts @@ -0,0 +1,255 @@ +import { describe, expect, it, vi } from 'vitest' +import { + createSandboxSnapshots, + createSnapshotTools, + defineSandbox, + InMemorySandboxInstanceStore, + memorySandboxSnapshots, +} from '../src' +import { makeFakeProvider } from './fakes' +import type { SandboxSnapshots } from '../src' + +const THREAD = 'thread' +const RUN = 'run' + +async function liveSnapshots() { + const memory = await memorySandboxSnapshots() + const instances = new InMemorySandboxInstanceStore() + const provider = makeFakeProvider() + const sandbox = defineSandbox({ id: 'sandbox', provider }) + await instances.upsert({ + key: sandbox.key({ threadId: THREAD, runId: 'old' }), + provider: provider.name, + providerSandboxId: 'existing', + threadId: THREAD, + updatedAt: Date.now(), + }) + const snapshots = createSandboxSnapshots({ + persistence: memory.persistence, + checkpoints: memory.checkpoints, + sandbox, + instances, + }) + return { snapshots, provider } +} + +function toolsFor( + snapshots: SandboxSnapshots, + options: { + createThreadId?: () => string + onForked?: (input: { + destinationThreadId: string + checkpointId: string + }) => void | Promise + } = {}, +) { + return createSnapshotTools(snapshots, { + threadId: THREAD, + runId: RUN, + createThreadId: options.createThreadId ?? (() => 'destination'), + ...(options.onForked === undefined ? {} : { onForked: options.onForked }), + }) +} + +async function executeTool( + tools: ReturnType, + name: string, + input: unknown, +) { + const tool = tools.find((candidate) => candidate.name === name) + if (tool === undefined || tool.execute === undefined) { + throw new Error(`Missing tool ${name}`) + } + return tool.execute(input) +} + +describe('createSnapshotTools', () => { + it('rejects empty factory ids', async () => { + const { snapshots } = await liveSnapshots() + expect(() => + createSnapshotTools(snapshots, { + threadId: '', + runId: RUN, + createThreadId: () => 'destination', + }), + ).toThrow( + expect.objectContaining({ code: 'SANDBOX_SNAPSHOT_INVALID_TOOL_INPUT' }), + ) + }) + + it('saves a named checkpoint without a model-supplied thread id', async () => { + const { snapshots, provider } = await liveSnapshots() + const tools = toolsFor(snapshots) + const saveSchema = tools[0].inputSchema + + expect(saveSchema).toMatchObject({ + required: ['label'], + additionalProperties: false, + }) + expect( + saveSchema && + typeof saveSchema === 'object' && + 'properties' in saveSchema && + saveSchema.properties !== null && + typeof saveSchema.properties === 'object', + ).toBe(true) + if ( + saveSchema && + typeof saveSchema === 'object' && + 'properties' in saveSchema && + saveSchema.properties !== null && + typeof saveSchema.properties === 'object' + ) { + expect(Object.hasOwn(saveSchema.properties, 'threadId')).toBe(false) + } + + await expect( + executeTool(tools, 'save_sandbox_snapshot', { label: 'release-1' }), + ).resolves.toEqual({ + checkpointId: expect.any(String), + label: 'release-1', + threadId: THREAD, + }) + expect(provider.calls).toMatchObject({ create: 0, resume: 1 }) + }) + + it('rejects save while this thread already has a writer', async () => { + const { snapshots } = await liveSnapshots() + const writer = await snapshots.checkpoints.acquireWriter(THREAD) + const tools = toolsFor(snapshots) + + await expect( + executeTool(tools, 'save_sandbox_snapshot', { label: 'busy' }), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_WRITER_CONFLICT' }) + await writer.release() + }) + + it('forks the latest checkpoint into a minted destination thread', async () => { + const { snapshots } = await liveSnapshots() + const tools = toolsFor(snapshots, { createThreadId: () => 'branch-1' }) + const saved = await executeTool(tools, 'save_sandbox_snapshot', { + label: 'base', + }) + if ( + saved === null || + typeof saved !== 'object' || + !('checkpointId' in saved) + ) { + throw new Error('save did not return a checkpointId') + } + const sourceId = saved.checkpointId + + const forked = await executeTool(tools, 'fork_sandbox_snapshot', {}) + expect(forked).toEqual({ + checkpointId: expect.any(String), + destinationThreadId: 'branch-1', + }) + expect(await snapshots.checkpoints.getHead(THREAD)).toBe(sourceId) + expect(await snapshots.checkpoints.getHead('branch-1')).not.toBeNull() + }) + + it('forks a selected checkpoint while the source writer is held', async () => { + const { snapshots } = await liveSnapshots() + const tools = toolsFor(snapshots, { createThreadId: () => 'branch-2' }) + const saved = await executeTool(tools, 'save_sandbox_snapshot', { + label: 'base', + }) + if ( + saved === null || + typeof saved !== 'object' || + !('checkpointId' in saved) || + typeof saved.checkpointId !== 'string' + ) { + throw new Error('save did not return a checkpointId') + } + const writer = await snapshots.checkpoints.acquireWriter(THREAD) + + await expect( + executeTool(tools, 'fork_sandbox_snapshot', { + checkpointId: saved.checkpointId, + }), + ).resolves.toMatchObject({ destinationThreadId: 'branch-2' }) + await writer.release() + }) + + it('rejects a fork when the thread has no checkpoint', async () => { + const { snapshots } = await liveSnapshots() + const tools = toolsFor(snapshots) + + await expect( + executeTool(tools, 'fork_sandbox_snapshot', {}), + ).rejects.toMatchObject({ code: 'SANDBOX_SNAPSHOT_MISSING_CHECKPOINT' }) + }) + + it('calls onForked after a successful fork', async () => { + const { snapshots } = await liveSnapshots() + const onForked = vi.fn() + const tools = toolsFor(snapshots, { + createThreadId: () => 'branch-3', + onForked, + }) + await executeTool(tools, 'save_sandbox_snapshot', { label: 'base' }) + const forked = await executeTool(tools, 'fork_sandbox_snapshot', {}) + if ( + forked === null || + typeof forked !== 'object' || + !('checkpointId' in forked) + ) { + throw new Error('fork did not return a checkpointId') + } + expect(onForked).toHaveBeenCalledWith({ + destinationThreadId: 'branch-3', + checkpointId: forked.checkpointId, + }) + }) + + it('reads artifact metadata from the bound thread', async () => { + const { snapshots } = await liveSnapshots() + const tools = toolsFor(snapshots) + const bytes = new TextEncoder().encode('hello') + const digest = await crypto.subtle.digest('SHA-256', new Uint8Array(bytes)) + const blobKey = `sandbox-artifacts/sha256/${Array.from( + new Uint8Array(digest), + (byte) => byte.toString(16).padStart(2, '0'), + ).join('')}` + await snapshots.persistence.stores.blobs.put(blobKey, bytes) + const writer = await snapshots.checkpoints.acquireWriter(THREAD) + await snapshots.checkpoints.append({ + checkpoint: { + id: 'checkpoint', + threadId: THREAD, + parentCheckpointId: null, + createdAt: 1, + reason: 'named', + files: [], + conversation: [], + artifacts: [ + { + artifactId: 'artifact', + name: 'file.txt', + mimeType: 'text/plain', + size: bytes.byteLength, + blobKey, + createdAt: 1, + }, + ], + }, + expectedHeadId: null, + writer, + }) + await writer.release() + + await expect( + executeTool(tools, 'read_sandbox_snapshot_artifact', { + checkpointId: 'checkpoint', + artifactId: 'artifact', + }), + ).resolves.toEqual({ + artifactId: 'artifact', + name: 'file.txt', + mimeType: 'text/plain', + size: bytes.byteLength, + createdAt: 1, + }) + }) +}) From 952adf451aeeaa5a7661cafe074159cf8a0ece9f Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Fri, 14 Aug 2026 17:56:29 +0200 Subject: [PATCH 08/12] feat(examples): add App Studio fork-and-compare demo Add a /app-studio page in ts-react-chat. The agent builds an app in a Docker sandbox and shows a preview. Fork copies the checkpoint. Compare creates two forks so the user can keep one. --- docs/sandbox/portable-snapshots-fork.md | 19 + docs/sandbox/portable-snapshots.md | 1 + examples/ts-react-chat/README.md | 27 + .../ts-react-chat/src/components/Header.tsx | 13 + .../src/lib/app-studio-helpers.ts | 121 +++ .../ts-react-chat/src/lib/app-studio-store.ts | 22 + .../ts-react-chat/src/lib/app-studio.test.ts | 152 ++++ examples/ts-react-chat/src/lib/app-studio.ts | 99 +++ examples/ts-react-chat/src/routeTree.gen.ts | 63 ++ .../src/routes/api.app-studio-fork.ts | 57 ++ .../src/routes/api.app-studio.ts | 87 +++ .../ts-react-chat/src/routes/app-studio.tsx | 727 ++++++++++++++++++ examples/ts-react-chat/src/routes/index.tsx | 8 + 13 files changed, 1396 insertions(+) create mode 100644 examples/ts-react-chat/src/lib/app-studio-helpers.ts create mode 100644 examples/ts-react-chat/src/lib/app-studio-store.ts create mode 100644 examples/ts-react-chat/src/lib/app-studio.test.ts create mode 100644 examples/ts-react-chat/src/lib/app-studio.ts create mode 100644 examples/ts-react-chat/src/routes/api.app-studio-fork.ts create mode 100644 examples/ts-react-chat/src/routes/api.app-studio.ts create mode 100644 examples/ts-react-chat/src/routes/app-studio.tsx diff --git a/docs/sandbox/portable-snapshots-fork.md b/docs/sandbox/portable-snapshots-fork.md index 943c15c5cb..6f5dfbceb8 100644 --- a/docs/sandbox/portable-snapshots-fork.md +++ b/docs/sandbox/portable-snapshots-fork.md @@ -83,6 +83,25 @@ export async function forkCheckpoint( Use the same authorization rule for both threads. A client-selected checkpoint id is not proof of access. +## See it in the example + +The React chat example has an App Studio page at `/app-studio`. That page +starts from one prompt, shows a live preview, then lets you fork the chat or +compare two directions. + +1. Open `examples/ts-react-chat`. +2. Set `XAI_API_KEY` and start Docker. +3. Run `pnpm dev` and open `/app-studio`. +4. Build an app. Then use **Fork chat** or **Compare two directions**. + +When you select **Compare two directions** and submit, the page calls +`/api/app-studio-fork` with `count: 2`. Each fork gets the same product request +and a different visual prompt. You keep one branch. The source thread stays +unchanged. + +See [App Studio](../../examples/ts-react-chat/README.md#app-studio) +in the example README. The page path is `/app-studio`. + If you use SQLite, put the fork in one transaction. That transaction must copy the source conversation and reject a destination thread that is not empty. See [Keep Files After Reload](./portable-snapshots-configure). diff --git a/docs/sandbox/portable-snapshots.md b/docs/sandbox/portable-snapshots.md index 35214d889c..6bbdbf3ff2 100644 --- a/docs/sandbox/portable-snapshots.md +++ b/docs/sandbox/portable-snapshots.md @@ -22,6 +22,7 @@ Start with persistence. Then add a product page only when you need that action. | Chat persistence already | Files come back after a reload | [Keep Files After Reload](./portable-snapshots-configure#reuse-existing-persistence) | | Snapshots already wired | A user marks one version | [Save a Named Version](./portable-snapshots-save) | | Snapshots already wired | A user branches from one version | [Branch From a Version](./portable-snapshots-fork) | +| Snapshots already wired | A product page that forks and compares two directions | [Branch From a Version](./portable-snapshots-fork#see-it-in-the-example) | | Snapshots already wired | A user downloads a generated file | [Send a Frozen File](./portable-snapshots-artifacts) | | Snapshots already wired | The agent saves or forks this thread | [Let the Agent Save and Fork](./portable-snapshots-tools) | | Snapshots already wired | Custom include or redact rules | [What a Snapshot Stores](./portable-snapshots-safety) | diff --git a/examples/ts-react-chat/README.md b/examples/ts-react-chat/README.md index ff61b66566..3cfc96056e 100644 --- a/examples/ts-react-chat/README.md +++ b/examples/ts-react-chat/README.md @@ -454,3 +454,30 @@ cloud runs skip the tools and do a plain triage. In production you wouldn't need ngrok — your orchestrator already has a public URL to advertise. Set keys in `.env.local`, then `pnpm dev` and open `/sandboxes`. + +## App Studio + +A small product page at `/app-studio` on top of portable snapshots. You describe an app. The +agent scaffolds a TanStack Start app in a Docker sandbox, starts the preview +server, and shows the live URL. After the first build: + +- **Fork chat** copies the checkpoint and conversation into a new thread. + Continue from that copy. The source thread stays unchanged. +- **Compare two directions** creates two forks. Each fork gets a different + visual prompt. Pick **Keep variant A** or **Keep variant B** to continue + on that branch. + +Needs: + +- Docker running on the host +- `XAI_API_KEY` in `.env.local` (Grok Build runs inside the sandbox) + +Stay on the page during the first build. The first run installs the CLI, +scaffolds the app, and starts the preview. If you leave, the request aborts. + +The snapshot store is `src/lib/sqlite-persistence.ts` at +`.data/app-studio.db` (gitignored). The server route is `/api/app-studio`. +The fork route is `/api/app-studio-fork`. + +See [Branch From a Version](../../docs/sandbox/portable-snapshots-fork.md) +for the `snapshots.fork` contract this page uses. diff --git a/examples/ts-react-chat/src/components/Header.tsx b/examples/ts-react-chat/src/components/Header.tsx index 9b4d744f98..fa562489d9 100644 --- a/examples/ts-react-chat/src/components/Header.tsx +++ b/examples/ts-react-chat/src/components/Header.tsx @@ -340,6 +340,19 @@ export default function Header() { Persistent Chat + setIsOpen(false)} + className="flex items-center gap-3 p-3 rounded-lg hover:bg-gray-800 transition-colors mb-2" + activeProps={{ + className: + 'flex items-center gap-3 p-3 rounded-lg bg-cyan-600 hover:bg-cyan-700 transition-colors mb-2', + }} + > + + App Studio + + setIsOpen(false)} diff --git a/examples/ts-react-chat/src/lib/app-studio-helpers.ts b/examples/ts-react-chat/src/lib/app-studio-helpers.ts new file mode 100644 index 0000000000..ce1d62bab0 --- /dev/null +++ b/examples/ts-react-chat/src/lib/app-studio-helpers.ts @@ -0,0 +1,121 @@ +import { + SandboxCheckpointWriterConflictError, + SandboxSnapshotError, +} from '@tanstack/ai-sandbox' +import type { + SandboxSnapshots, + SaveSandboxSnapshotInput, +} from '@tanstack/ai-sandbox' + +export function previewUrlFrom(output: unknown): string | null { + let value: unknown = output + if (typeof output === 'string') { + try { + value = JSON.parse(output) + } catch { + return /^https?:\/\//.test(output) ? output : null + } + } + if (value !== null && typeof value === 'object' && 'url' in value) { + const url = value.url + return typeof url === 'string' ? url : null + } + return null +} + +export const DEFAULT_COMPARE_PROMPT = + 'Keep the same product. Change only the visual direction.' + +export function comparePrompt(userText: string): string { + const trimmed = userText.trim() + return trimmed.length > 0 ? trimmed : DEFAULT_COMPARE_PROMPT +} + +export function variantPrompt(userText: string, variant: 'A' | 'B'): string { + const prompt = comparePrompt(userText) + if (variant === 'A') { + return `${prompt}\n\nThis is variant A. Keep the same product. Use a bold, high-contrast, compact visual direction.` + } + return `${prompt}\n\nThis is variant B. Keep the same product. Use a soft, spacious, calm visual direction.` +} + +export function threadIdsFromForkBody(body: unknown): Array { + if (body === null || typeof body !== 'object') return [] + const forks = Reflect.get(body, 'forks') + if (!Array.isArray(forks)) return [] + return forks.flatMap((fork) => { + if (fork === null || typeof fork !== 'object') return [] + const id = Reflect.get(fork, 'threadId') + return typeof id === 'string' && id.length > 0 ? [id] : [] + }) +} + +export function errorMessageFromBody(body: unknown, fallback: string): string { + if (body !== null && typeof body === 'object' && 'error' in body) { + const error = body.error + if (typeof error === 'string' && error.length > 0) return error + } + return fallback +} + +function canUseSavedHead(error: unknown): boolean { + if (error instanceof SandboxSnapshotError) { + return ( + error.code === 'SANDBOX_SNAPSHOT_MISSING_SANDBOX' || + error.code === 'SANDBOX_SNAPSHOT_MISSING_INSTANCES' || + error.code === 'SANDBOX_SNAPSHOT_MISSING_REUSABLE_SANDBOX' || + error.code === 'SANDBOX_SNAPSHOT_REUSE_NONE' + ) + } + return error instanceof SandboxCheckpointWriterConflictError +} + +export async function forkStudioThreads(input: { + snapshots: SandboxSnapshots + threadId: string + runId: string + count: 1 | 2 + label?: string + sandbox?: SaveSandboxSnapshotInput['sandbox'] + instances?: SaveSandboxSnapshotInput['instances'] + locks?: SaveSandboxSnapshotInput['locks'] +}): Promise<{ + sourceCheckpointId: string + forks: Array<{ threadId: string; checkpointId: string }> +}> { + let sourceCheckpointId: string | null = null + try { + const saved = await input.snapshots.save({ + threadId: input.threadId, + runId: input.runId, + label: input.label ?? 'studio-fork', + ...(input.sandbox === undefined ? {} : { sandbox: input.sandbox }), + ...(input.instances === undefined ? {} : { instances: input.instances }), + ...(input.locks === undefined ? {} : { locks: input.locks }), + }) + sourceCheckpointId = saved.id + } catch (error) { + if (!canUseSavedHead(error)) throw error + sourceCheckpointId = await input.snapshots.checkpoints.getHead( + input.threadId, + ) + } + if (sourceCheckpointId === null) { + throw new Error('Build the app first. Then you can fork or compare.') + } + + const forks: Array<{ threadId: string; checkpointId: string }> = [] + for (let index = 0; index < input.count; index++) { + const destinationThreadId = `studio-${crypto.randomUUID()}` + const checkpoint = await input.snapshots.fork({ + threadId: input.threadId, + checkpointId: sourceCheckpointId, + destinationThreadId, + }) + forks.push({ + threadId: destinationThreadId, + checkpointId: checkpoint.id, + }) + } + return { sourceCheckpointId, forks } +} diff --git a/examples/ts-react-chat/src/lib/app-studio-store.ts b/examples/ts-react-chat/src/lib/app-studio-store.ts new file mode 100644 index 0000000000..c238b6939c --- /dev/null +++ b/examples/ts-react-chat/src/lib/app-studio-store.ts @@ -0,0 +1,22 @@ +import { InMemoryLockStore } from '@tanstack/ai/locks' +import { InMemorySandboxInstanceStore } from '@tanstack/ai-sandbox' +import { sqliteSandboxSnapshots } from './sqlite-persistence' + +let snapshots: ReturnType | undefined +const instances = new InMemorySandboxInstanceStore() +const locks = new InMemoryLockStore() + +export function appStudioSnapshots() { + return (snapshots ??= sqliteSandboxSnapshots({ + url: './.data/app-studio.db', + migrate: true, + })) +} + +export function appStudioInstances() { + return instances +} + +export function appStudioLocks() { + return locks +} diff --git a/examples/ts-react-chat/src/lib/app-studio.test.ts b/examples/ts-react-chat/src/lib/app-studio.test.ts new file mode 100644 index 0000000000..bced5bfccf --- /dev/null +++ b/examples/ts-react-chat/src/lib/app-studio.test.ts @@ -0,0 +1,152 @@ +import { describe, expect, it } from 'vitest' +import { memorySandboxSnapshots } from '@tanstack/ai-sandbox' +import { + comparePrompt, + errorMessageFromBody, + forkStudioThreads, + previewUrlFrom, + threadIdsFromForkBody, + variantPrompt, +} from './app-studio-helpers' + +async function seedHead( + snapshots: Awaited>, + threadId: string, + conversation: Array<{ role: 'user'; content: string }> = [], +) { + const writer = await snapshots.checkpoints.acquireWriter(threadId) + await snapshots.checkpoints.append({ + checkpoint: { + id: `${threadId}-head`, + threadId, + parentCheckpointId: null, + createdAt: 1, + reason: 'automatic', + files: [], + conversation, + artifacts: [], + }, + expectedHeadId: null, + writer, + }) + await writer.release() +} + +describe('app studio helpers', () => { + it('reads a preview URL from an exposePreview result', () => { + expect(previewUrlFrom({ url: 'http://127.0.0.1:5173' })).toBe( + 'http://127.0.0.1:5173', + ) + expect(previewUrlFrom('{"url":"http://127.0.0.1:5173"}')).toBe( + 'http://127.0.0.1:5173', + ) + expect(previewUrlFrom('http://127.0.0.1:5173')).toBe( + 'http://127.0.0.1:5173', + ) + expect(previewUrlFrom('not-a-url')).toBeNull() + }) + + it('builds distinct compare prompts', () => { + expect(variantPrompt('Make it warmer', 'A')).toContain('variant A') + expect(variantPrompt('Make it warmer', 'B')).toContain('variant B') + expect(comparePrompt(' ')).toBe( + 'Keep the same product. Change only the visual direction.', + ) + expect(comparePrompt('Make it warmer')).toBe('Make it warmer') + }) + + it('reads fork thread ids and error text from JSON bodies', () => { + expect( + threadIdsFromForkBody({ + forks: [{ threadId: 'a' }, { threadId: 'b' }], + }), + ).toEqual(['a', 'b']) + expect(threadIdsFromForkBody({ forks: [{}] })).toEqual([]) + expect(errorMessageFromBody({ error: 'no checkpoint' }, 'fallback')).toBe( + 'no checkpoint', + ) + expect(errorMessageFromBody({}, 'fallback')).toBe('fallback') + }) +}) + +describe('forkStudioThreads', () => { + it('throws when the thread has no checkpoint', async () => { + const snapshots = await memorySandboxSnapshots() + await expect( + forkStudioThreads({ + snapshots, + threadId: 'empty', + runId: 'run-1', + count: 1, + }), + ).rejects.toThrow('Build the app first') + }) + + it('forks one destination from the saved head', async () => { + const snapshots = await memorySandboxSnapshots() + await seedHead(snapshots, 'source') + const result = await forkStudioThreads({ + snapshots, + threadId: 'source', + runId: 'run-1', + count: 1, + }) + expect(result.sourceCheckpointId).toBe('source-head') + expect(result.forks).toHaveLength(1) + const destination = result.forks[0] + expect(destination?.threadId.startsWith('studio-')).toBe(true) + expect( + await snapshots.checkpoints.getHead(destination?.threadId ?? ''), + ).toBe(destination?.checkpointId) + }) + + it('forks two destinations for a compare', async () => { + const snapshots = await memorySandboxSnapshots() + await seedHead(snapshots, 'source') + const result = await forkStudioThreads({ + snapshots, + threadId: 'source', + runId: 'run-2', + count: 2, + }) + expect(result.forks).toHaveLength(2) + expect(result.forks[0]?.threadId).not.toBe(result.forks[1]?.threadId) + }) + + it('copies the source conversation onto each fork', async () => { + const snapshots = await memorySandboxSnapshots() + await seedHead(snapshots, 'source', [{ role: 'user', content: 'build it' }]) + const result = await forkStudioThreads({ + snapshots, + threadId: 'source', + runId: 'run-3', + count: 1, + }) + const destination = result.forks[0] + const checkpoint = await snapshots.checkpoints.get( + destination?.checkpointId ?? '', + ) + expect(checkpoint?.conversation).toEqual([ + { role: 'user', content: 'build it' }, + ]) + }) + + it('does not hide a save error that is not a missing-sandbox miss', async () => { + const snapshots = await memorySandboxSnapshots() + await seedHead(snapshots, 'source') + const failing = { + ...snapshots, + save: async () => { + throw new Error('disk full') + }, + } + await expect( + forkStudioThreads({ + snapshots: failing, + threadId: 'source', + runId: 'run-4', + count: 1, + }), + ).rejects.toThrow('disk full') + }) +}) diff --git a/examples/ts-react-chat/src/lib/app-studio.ts b/examples/ts-react-chat/src/lib/app-studio.ts new file mode 100644 index 0000000000..08fdf0146c --- /dev/null +++ b/examples/ts-react-chat/src/lib/app-studio.ts @@ -0,0 +1,99 @@ +import { toolDefinition } from '@tanstack/ai' +import { + GROK_CLI_INSTALL_COMMAND, + grokBuildText, +} from '@tanstack/ai-grok-build' +import { + createSecrets, + defineSandbox, + defineWorkspace, +} from '@tanstack/ai-sandbox' +import { dockerSandbox } from '@tanstack/ai-sandbox-docker' +import { z } from 'zod' +import type { AnyTextAdapter } from '@tanstack/ai' +import type { + SandboxDefinition, + SandboxEnsureContext, +} from '@tanstack/ai-sandbox' + +export const PREVIEW_PORT = 5173 + +const SCAFFOLD = + 'Scaffold with the TanStack CLI via npx. Run it exactly like this: `npx --yes @tanstack/cli create my-app --framework react --no-examples --intent -y`. Do not guess other package names.' + +const APP = + 'Turn it into a self-contained interactive app. No external APIs, no env vars, no keys. Keep state in the browser (localStorage). Make it look polished.' + +const RUN = `Add \`server: { host: true, allowedHosts: true }\` to vite.config.ts. Start the dev server on port ${PREVIEW_PORT}: \`pnpm dev --host 0.0.0.0 --port ${PREVIEW_PORT}\`. When it is listening, call exposePreview with { "port": ${PREVIEW_PORT} } and share the URL.` + +export function missingAppStudioEnv(): Array { + return process.env.XAI_API_KEY ? [] : ['XAI_API_KEY'] +} + +export function buildAppStudioAdapter(): AnyTextAdapter { + return grokBuildText('grok-build') +} + +export function buildAppStudioSandbox(): SandboxDefinition { + const key = process.env.XAI_API_KEY + return defineSandbox({ + id: 'app-studio', + provider: dockerSandbox({ + image: process.env.SANDBOX_IMAGE ?? 'node:22', + publishPorts: [PREVIEW_PORT], + }), + workspace: defineWorkspace({ + source: { type: 'none' }, + setup: ({ serial }) => serial(GROK_CLI_INSTALL_COMMAND), + secrets: createSecrets(key ? { XAI_API_KEY: key } : {}), + }), + lifecycle: { reuse: 'thread' }, + }) +} + +export const tanstackStartRecipe = toolDefinition({ + name: 'tanstackStartRecipe', + description: + 'The recipe for a self-contained TanStack Start app in this sandbox. Call this before you scaffold.', + inputSchema: z.object({ + section: z + .enum(['scaffold', 'app', 'run', 'all']) + .describe('Which part of the recipe you need. Use all first.'), + }), +}).server(({ section }) => { + const recipe = { scaffold: SCAFFOLD, app: APP, run: RUN } + return section === 'all' ? recipe : { [section]: recipe[section] } +}) + +export function makeExposePreviewTool( + definition: SandboxDefinition, + threadId: string, + bookkeeping?: Pick, +) { + return toolDefinition({ + name: 'exposePreview', + description: `Expose the sandbox port the dev server is listening on and return a preview URL. Call this after the server is up on port ${PREVIEW_PORT}.`, + inputSchema: z.object({ + port: z.number().int().min(1024).max(65535), + }), + }).server(async ({ port }) => { + const handle = await definition.ensure({ + threadId, + runId: 'expose-preview', + ...bookkeeping, + }) + const channel = await handle.ports.connect(port) + return { url: channel.url } + }) +} + +export const APP_STUDIO_SYSTEM_PROMPT = [ + 'You work in this sandbox.', + 'If the workspace is empty, call tanstackStartRecipe with section all, then scaffold, build the app, start the preview, and call exposePreview.', + 'If the workspace already has an app, do not scaffold again.', + 'Install dependencies if node_modules is missing.', + 'Apply the requested change.', + `Restart the preview on port ${PREVIEW_PORT}.`, + 'Then call exposePreview.', + `The preview port must be ${PREVIEW_PORT}.`, +].join(' ') diff --git a/examples/ts-react-chat/src/routeTree.gen.ts b/examples/ts-react-chat/src/routeTree.gen.ts index 7dfa0251b7..2192f295a6 100644 --- a/examples/ts-react-chat/src/routeTree.gen.ts +++ b/examples/ts-react-chat/src/routeTree.gen.ts @@ -26,6 +26,7 @@ import { Route as ImageToolReproRouteImport } from './routes/image-tool-repro' import { Route as ImageGenRouteImport } from './routes/image-gen' import { Route as GenerationHooksRouteImport } from './routes/generation-hooks' import { Route as CapabilityDemoRouteImport } from './routes/capability-demo' +import { Route as AppStudioRouteImport } from './routes/app-studio' import { Route as IndexRouteImport } from './routes/index' import { Route as GenerationsVideoRouteImport } from './routes/generations.video' import { Route as GenerationsTranscriptionRouteImport } from './routes/generations.transcription' @@ -59,6 +60,8 @@ import { Route as ApiImageToolReproRouteImport } from './routes/api.image-tool-r import { Route as ApiImageGenRouteImport } from './routes/api.image-gen' import { Route as ApiCapabilityDemoRouteImport } from './routes/api.capability-demo' import { Route as ApiArtifactsRouteImport } from './routes/api.artifacts' +import { Route as ApiAppStudioForkRouteImport } from './routes/api.app-studio-fork' +import { Route as ApiAppStudioRouteImport } from './routes/api.app-studio' import { Route as ExampleGuitarsIndexRouteImport } from './routes/example.guitars/index' import { Route as ExampleGuitarsGuitarIdRouteImport } from './routes/example.guitars/$guitarId' import { Route as ApiGenerateVideoRouteImport } from './routes/api.generate.video' @@ -152,6 +155,11 @@ const CapabilityDemoRoute = CapabilityDemoRouteImport.update({ path: '/capability-demo', getParentRoute: () => rootRouteImport, } as any) +const AppStudioRoute = AppStudioRouteImport.update({ + id: '/app-studio', + path: '/app-studio', + getParentRoute: () => rootRouteImport, +} as any) const IndexRoute = IndexRouteImport.update({ id: '/', path: '/', @@ -321,6 +329,16 @@ const ApiArtifactsRoute = ApiArtifactsRouteImport.update({ path: '/api/artifacts', getParentRoute: () => rootRouteImport, } as any) +const ApiAppStudioForkRoute = ApiAppStudioForkRouteImport.update({ + id: '/api/app-studio-fork', + path: '/api/app-studio-fork', + getParentRoute: () => rootRouteImport, +} as any) +const ApiAppStudioRoute = ApiAppStudioRouteImport.update({ + id: '/api/app-studio', + path: '/api/app-studio', + getParentRoute: () => rootRouteImport, +} as any) const ExampleGuitarsIndexRoute = ExampleGuitarsIndexRouteImport.update({ id: '/example/guitars/', path: '/example/guitars/', @@ -360,6 +378,7 @@ const ApiGenerateImageArtifactRoute = export interface FileRoutesByFullPath { '/': typeof IndexRoute + '/app-studio': typeof AppStudioRoute '/capability-demo': typeof CapabilityDemoRoute '/generation-hooks': typeof GenerationHooksRoute '/image-gen': typeof ImageGenRoute @@ -377,6 +396,8 @@ export interface FileRoutesByFullPath { '/server-fn-chat': typeof ServerFnChatRoute '/threads': typeof ThreadsRoute '/typesafe-tools': typeof TypesafeToolsRoute + '/api/app-studio': typeof ApiAppStudioRoute + '/api/app-studio-fork': typeof ApiAppStudioForkRoute '/api/artifacts': typeof ApiArtifactsRoute '/api/capability-demo': typeof ApiCapabilityDemoRoute '/api/image-gen': typeof ApiImageGenRoute @@ -419,6 +440,7 @@ export interface FileRoutesByFullPath { } export interface FileRoutesByTo { '/': typeof IndexRoute + '/app-studio': typeof AppStudioRoute '/capability-demo': typeof CapabilityDemoRoute '/generation-hooks': typeof GenerationHooksRoute '/image-gen': typeof ImageGenRoute @@ -436,6 +458,8 @@ export interface FileRoutesByTo { '/server-fn-chat': typeof ServerFnChatRoute '/threads': typeof ThreadsRoute '/typesafe-tools': typeof TypesafeToolsRoute + '/api/app-studio': typeof ApiAppStudioRoute + '/api/app-studio-fork': typeof ApiAppStudioForkRoute '/api/artifacts': typeof ApiArtifactsRoute '/api/capability-demo': typeof ApiCapabilityDemoRoute '/api/image-gen': typeof ApiImageGenRoute @@ -479,6 +503,7 @@ export interface FileRoutesByTo { export interface FileRoutesById { __root__: typeof rootRouteImport '/': typeof IndexRoute + '/app-studio': typeof AppStudioRoute '/capability-demo': typeof CapabilityDemoRoute '/generation-hooks': typeof GenerationHooksRoute '/image-gen': typeof ImageGenRoute @@ -496,6 +521,8 @@ export interface FileRoutesById { '/server-fn-chat': typeof ServerFnChatRoute '/threads': typeof ThreadsRoute '/typesafe-tools': typeof TypesafeToolsRoute + '/api/app-studio': typeof ApiAppStudioRoute + '/api/app-studio-fork': typeof ApiAppStudioForkRoute '/api/artifacts': typeof ApiArtifactsRoute '/api/capability-demo': typeof ApiCapabilityDemoRoute '/api/image-gen': typeof ApiImageGenRoute @@ -540,6 +567,7 @@ export interface FileRouteTypes { fileRoutesByFullPath: FileRoutesByFullPath fullPaths: | '/' + | '/app-studio' | '/capability-demo' | '/generation-hooks' | '/image-gen' @@ -557,6 +585,8 @@ export interface FileRouteTypes { | '/server-fn-chat' | '/threads' | '/typesafe-tools' + | '/api/app-studio' + | '/api/app-studio-fork' | '/api/artifacts' | '/api/capability-demo' | '/api/image-gen' @@ -599,6 +629,7 @@ export interface FileRouteTypes { fileRoutesByTo: FileRoutesByTo to: | '/' + | '/app-studio' | '/capability-demo' | '/generation-hooks' | '/image-gen' @@ -616,6 +647,8 @@ export interface FileRouteTypes { | '/server-fn-chat' | '/threads' | '/typesafe-tools' + | '/api/app-studio' + | '/api/app-studio-fork' | '/api/artifacts' | '/api/capability-demo' | '/api/image-gen' @@ -658,6 +691,7 @@ export interface FileRouteTypes { id: | '__root__' | '/' + | '/app-studio' | '/capability-demo' | '/generation-hooks' | '/image-gen' @@ -675,6 +709,8 @@ export interface FileRouteTypes { | '/server-fn-chat' | '/threads' | '/typesafe-tools' + | '/api/app-studio' + | '/api/app-studio-fork' | '/api/artifacts' | '/api/capability-demo' | '/api/image-gen' @@ -718,6 +754,7 @@ export interface FileRouteTypes { } export interface RootRouteChildren { IndexRoute: typeof IndexRoute + AppStudioRoute: typeof AppStudioRoute CapabilityDemoRoute: typeof CapabilityDemoRoute GenerationHooksRoute: typeof GenerationHooksRoute ImageGenRoute: typeof ImageGenRoute @@ -735,6 +772,8 @@ export interface RootRouteChildren { ServerFnChatRoute: typeof ServerFnChatRoute ThreadsRoute: typeof ThreadsRoute TypesafeToolsRoute: typeof TypesafeToolsRoute + ApiAppStudioRoute: typeof ApiAppStudioRoute + ApiAppStudioForkRoute: typeof ApiAppStudioForkRoute ApiArtifactsRoute: typeof ApiArtifactsRoute ApiCapabilityDemoRoute: typeof ApiCapabilityDemoRoute ApiImageGenRoute: typeof ApiImageGenRoute @@ -896,6 +935,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof CapabilityDemoRouteImport parentRoute: typeof rootRouteImport } + '/app-studio': { + id: '/app-studio' + path: '/app-studio' + fullPath: '/app-studio' + preLoaderRoute: typeof AppStudioRouteImport + parentRoute: typeof rootRouteImport + } '/': { id: '/' path: '/' @@ -1127,6 +1173,20 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiArtifactsRouteImport parentRoute: typeof rootRouteImport } + '/api/app-studio-fork': { + id: '/api/app-studio-fork' + path: '/api/app-studio-fork' + fullPath: '/api/app-studio-fork' + preLoaderRoute: typeof ApiAppStudioForkRouteImport + parentRoute: typeof rootRouteImport + } + '/api/app-studio': { + id: '/api/app-studio' + path: '/api/app-studio' + fullPath: '/api/app-studio' + preLoaderRoute: typeof ApiAppStudioRouteImport + parentRoute: typeof rootRouteImport + } '/example/guitars/': { id: '/example/guitars/' path: '/example/guitars' @@ -1192,6 +1252,7 @@ const ApiGenerateImageRouteWithChildren = const rootRouteChildren: RootRouteChildren = { IndexRoute: IndexRoute, + AppStudioRoute: AppStudioRoute, CapabilityDemoRoute: CapabilityDemoRoute, GenerationHooksRoute: GenerationHooksRoute, ImageGenRoute: ImageGenRoute, @@ -1209,6 +1270,8 @@ const rootRouteChildren: RootRouteChildren = { ServerFnChatRoute: ServerFnChatRoute, ThreadsRoute: ThreadsRoute, TypesafeToolsRoute: TypesafeToolsRoute, + ApiAppStudioRoute: ApiAppStudioRoute, + ApiAppStudioForkRoute: ApiAppStudioForkRoute, ApiArtifactsRoute: ApiArtifactsRoute, ApiCapabilityDemoRoute: ApiCapabilityDemoRoute, ApiImageGenRoute: ApiImageGenRoute, diff --git a/examples/ts-react-chat/src/routes/api.app-studio-fork.ts b/examples/ts-react-chat/src/routes/api.app-studio-fork.ts new file mode 100644 index 0000000000..dd29fec3c2 --- /dev/null +++ b/examples/ts-react-chat/src/routes/api.app-studio-fork.ts @@ -0,0 +1,57 @@ +import { createFileRoute } from '@tanstack/react-router' +import { buildAppStudioSandbox } from '../lib/app-studio' +import { forkStudioThreads } from '../lib/app-studio-helpers' +import { + appStudioInstances, + appStudioLocks, + appStudioSnapshots, +} from '../lib/app-studio-store' + +function json(status: number, body: unknown): Response { + return new Response(JSON.stringify(body), { + status, + headers: { 'content-type': 'application/json' }, + }) +} + +export const Route = createFileRoute('/api/app-studio-fork')({ + server: { + handlers: { + POST: async ({ request }) => { + let body: unknown + try { + body = await request.json() + } catch { + return json(400, { error: 'invalid JSON body' }) + } + if (body === null || typeof body !== 'object') { + return json(400, { error: 'invalid JSON body' }) + } + const threadId = Reflect.get(body, 'threadId') + const label = Reflect.get(body, 'label') + const countValue = Reflect.get(body, 'count') + if (typeof threadId !== 'string' || threadId.length === 0) { + return json(400, { error: 'threadId is required' }) + } + const count = countValue === 2 ? 2 : 1 + try { + const result = await forkStudioThreads({ + snapshots: appStudioSnapshots(), + threadId, + runId: `studio-fork-${crypto.randomUUID()}`, + count, + sandbox: buildAppStudioSandbox(), + instances: appStudioInstances(), + locks: appStudioLocks(), + ...(typeof label === 'string' && label.length > 0 ? { label } : {}), + }) + return json(200, result) + } catch (error) { + return json(409, { + error: error instanceof Error ? error.message : 'fork failed', + }) + } + }, + }, + }, +}) diff --git a/examples/ts-react-chat/src/routes/api.app-studio.ts b/examples/ts-react-chat/src/routes/api.app-studio.ts new file mode 100644 index 0000000000..18be4acc8a --- /dev/null +++ b/examples/ts-react-chat/src/routes/api.app-studio.ts @@ -0,0 +1,87 @@ +import { createFileRoute } from '@tanstack/react-router' +import { + chat, + chatParamsFromRequestBody, + toServerSentEventsResponse, +} from '@tanstack/ai' +import { reconstructChat, withPersistence } from '@tanstack/ai-persistence' +import { withLocks } from '@tanstack/ai/locks' +import { withSandbox } from '@tanstack/ai-sandbox' +import { + APP_STUDIO_SYSTEM_PROMPT, + buildAppStudioAdapter, + buildAppStudioSandbox, + makeExposePreviewTool, + missingAppStudioEnv, + tanstackStartRecipe, +} from '../lib/app-studio' +import { + appStudioInstances, + appStudioLocks, + appStudioSnapshots, +} from '../lib/app-studio-store' + +function jsonError(status: number, error: string): Response { + return new Response(JSON.stringify({ error }), { + status, + statusText: error.slice(0, 64), + headers: { 'content-type': 'application/json' }, + }) +} + +export const Route = createFileRoute('/api/app-studio')({ + server: { + handlers: { + POST: async ({ request }) => { + const missing = missingAppStudioEnv() + if (missing.length > 0) { + return jsonError( + 500, + `Missing required env: ${missing.join(', ')}. Set it and restart the dev server.`, + ) + } + + let params: Awaited> + try { + params = await chatParamsFromRequestBody(await request.json()) + } catch { + return jsonError(400, 'invalid JSON body') + } + const snapshots = appStudioSnapshots() + const instances = appStudioInstances() + const locks = appStudioLocks() + const sandbox = buildAppStudioSandbox() + const abortController = new AbortController() + request.signal.addEventListener('abort', () => abortController.abort()) + + const stream = chat({ + adapter: buildAppStudioAdapter(), + messages: params.messages, + threadId: params.threadId, + runId: params.runId, + systemPrompts: [APP_STUDIO_SYSTEM_PROMPT], + tools: [ + tanstackStartRecipe, + makeExposePreviewTool(sandbox, params.threadId, { + store: instances, + locks, + }), + ], + middleware: [ + withPersistence(snapshots.persistence), + withLocks(locks), + withSandbox(sandbox, { instances, snapshots }), + ], + abortController, + }) + + return toServerSentEventsResponse(stream, { abortController }) + }, + GET: ({ request }) => { + return reconstructChat(appStudioSnapshots().persistence, request, { + authorize: async (threadId) => threadId.length > 0, + }) + }, + }, + }, +}) diff --git a/examples/ts-react-chat/src/routes/app-studio.tsx b/examples/ts-react-chat/src/routes/app-studio.tsx new file mode 100644 index 0000000000..1f0651a021 --- /dev/null +++ b/examples/ts-react-chat/src/routes/app-studio.tsx @@ -0,0 +1,727 @@ +import { useCallback, useEffect, useMemo, useRef, useState } from 'react' +import { createFileRoute } from '@tanstack/react-router' +import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' +import { GitBranch, Play } from 'lucide-react' +import ReactMarkdown from 'react-markdown' +import rehypeHighlight from 'rehype-highlight' +import rehypeRaw from 'rehype-raw' +import rehypeSanitize from 'rehype-sanitize' +import remarkGfm from 'remark-gfm' +import { + errorMessageFromBody, + previewUrlFrom, + threadIdsFromForkBody, + variantPrompt, +} from '../lib/app-studio-helpers' +import type { UIMessage } from '@tanstack/ai-react' + +export const Route = createFileRoute('/app-studio')({ + component: AppStudioPage, +}) + +const connection = fetchServerSentEvents('/api/app-studio') +const THREADS_KEY = 'app-studio:threads' + +interface StudioThread { + id: string + title: string + parentId: string | null + variant?: 'A' | 'B' + inheritedPreviewUrls?: Array +} + +interface CompareState { + leftId: string + rightId: string + prompt: string + inheritedPreviewUrls: Array +} + +function loadThreads(): Array { + if (typeof window === 'undefined') return [] + try { + const raw = window.localStorage.getItem(THREADS_KEY) + const parsed: unknown = raw ? JSON.parse(raw) : null + if (!Array.isArray(parsed)) return [] + return parsed.flatMap((item) => { + if (item === null || typeof item !== 'object') return [] + const id = Reflect.get(item, 'id') + const title = Reflect.get(item, 'title') + const parentId = Reflect.get(item, 'parentId') + const variant = Reflect.get(item, 'variant') + const inherited = Reflect.get(item, 'inheritedPreviewUrls') + if (typeof id !== 'string' || typeof title !== 'string') return [] + const inheritedPreviewUrls = Array.isArray(inherited) + ? inherited.filter((url) => typeof url === 'string') + : [] + return [ + { + id, + title, + parentId: typeof parentId === 'string' ? parentId : null, + ...(variant === 'A' || variant === 'B' ? { variant } : {}), + ...(inheritedPreviewUrls.length > 0 ? { inheritedPreviewUrls } : {}), + }, + ] + }) + } catch { + return [] + } +} + +function newThread(parentId: string | null = null): StudioThread { + return { + id: `studio-${crypto.randomUUID()}`, + title: parentId ? 'Fork' : 'New app', + parentId, + } +} + +function ThreadNav({ + threads, + activeId, + compare, + onSelect, + horizontal = false, +}: { + threads: Array + activeId: string + compare: CompareState | null + onSelect: (id: string) => void + horizontal?: boolean +}) { + return ( + + ) +} + +function collectPreviewUrls(messages: Array): Set { + const urls = new Set() + for (const message of messages) { + for (const part of message.parts) { + if (part.type !== 'tool-call' || part.name !== 'exposePreview') continue + const url = previewUrlFrom(part.output) + if (url) urls.add(url) + } + } + return urls +} + +function latestPreview( + messages: Array, + skip: ReadonlySet = new Set(), +): string | null { + for (let index = messages.length - 1; index >= 0; index--) { + const message = messages[index] + if (!message) continue + for (const part of message.parts) { + if (part.type !== 'tool-call' || part.name !== 'exposePreview') continue + const url = previewUrlFrom(part.output) + if (url && !skip.has(url)) return url + } + } + return null +} + +function hasUserText(messages: Array, text: string): boolean { + return messages.some( + (message) => + message.role === 'user' && + message.parts.some( + (part) => part.type === 'text' && part.content === text, + ), + ) +} + +function AppStudioPage() { + const [threads, setThreads] = useState>([]) + const [activeId, setActiveId] = useState(null) + const [hydrated, setHydrated] = useState(false) + const [compare, setCompare] = useState(null) + const [wantCompare, setWantCompare] = useState(false) + const [busy, setBusy] = useState(false) + const [error, setError] = useState(null) + + useEffect(() => { + const loaded = loadThreads() + if (loaded.length === 0) { + const first = newThread() + setThreads([first]) + setActiveId(first.id) + } else { + setThreads(loaded) + setActiveId(loaded[0]?.id ?? null) + } + setHydrated(true) + }, []) + + useEffect(() => { + if (!hydrated) return + window.localStorage.setItem(THREADS_KEY, JSON.stringify(threads)) + }, [threads, hydrated]) + + const createRoot = () => { + const thread = newThread() + setCompare(null) + setWantCompare(false) + setError(null) + setThreads((prev) => [thread, ...prev]) + setActiveId(thread.id) + } + + const titleFrom = useCallback((id: string, title: string) => { + setThreads((prev) => + prev.map((thread) => + thread.id === id && + (thread.title === 'New app' || thread.title === 'Fork') + ? { ...thread, title } + : thread, + ), + ) + }, []) + + const forkOne = async (inheritedPreviewUrls: Array) => { + if (!activeId || busy) return + setBusy(true) + setError(null) + try { + const response = await fetch('/api/app-studio-fork', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ threadId: activeId, label: 'continue' }), + }) + const body: unknown = await response.json() + if (!response.ok) { + throw new Error(errorMessageFromBody(body, 'Could not fork this chat')) + } + const [nextId] = threadIdsFromForkBody(body) + if (nextId === undefined) { + throw new Error('Could not fork this chat') + } + const child: StudioThread = { + id: nextId, + title: 'Fork', + parentId: activeId, + ...(inheritedPreviewUrls.length > 0 ? { inheritedPreviewUrls } : {}), + } + setThreads((prev) => [child, ...prev]) + setCompare(null) + setActiveId(nextId) + } catch (cause) { + setError( + cause instanceof Error ? cause.message : 'Could not fork this chat', + ) + } finally { + setBusy(false) + } + } + + const startCompare = async ( + prompt: string, + inheritedPreviewUrls: Array, + ) => { + if (!activeId || busy) return + setBusy(true) + setError(null) + try { + const response = await fetch('/api/app-studio-fork', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + threadId: activeId, + count: 2, + label: 'compare', + }), + }) + const body: unknown = await response.json() + if (!response.ok) { + throw new Error( + errorMessageFromBody(body, 'Could not start the comparison'), + ) + } + const ids = threadIdsFromForkBody(body) + const leftId = ids[0] + const rightId = ids[1] + if (leftId === undefined || rightId === undefined) { + throw new Error('Could not start the comparison') + } + const inherited = + inheritedPreviewUrls.length > 0 ? { inheritedPreviewUrls } : {} + setThreads((prev) => [ + { + id: leftId, + title: 'Variant A', + parentId: activeId, + variant: 'A', + ...inherited, + }, + { + id: rightId, + title: 'Variant B', + parentId: activeId, + variant: 'B', + ...inherited, + }, + ...prev, + ]) + setCompare({ leftId, rightId, prompt, inheritedPreviewUrls }) + setWantCompare(false) + } catch (cause) { + setError( + cause instanceof Error + ? cause.message + : 'Could not start the comparison', + ) + } finally { + setBusy(false) + } + } + + const keepVariant = (id: string) => { + setActiveId(id) + setCompare(null) + } + + const onActiveTitle = useCallback( + (title: string) => { + if (activeId === null) return + titleFrom(activeId, title) + }, + [activeId, titleFrom], + ) + + if (!hydrated || !activeId) { + return ( +
+

Loading chats.

+
+ ) + } + + const selectThread = (id: string) => { + setCompare(null) + setWantCompare(false) + setActiveId(id) + } + + return ( +
+ + +
+
+ + +
+
+

App Studio

+

+ Describe an app. The agent builds it in a sandbox and shows a + preview. Fork the chat to continue, or compare two directions and + keep one. Needs Docker and XAI_API_KEY. +

+
+ + {error ? ( +

+ {error} +

+ ) : null} + + {compare ? ( +
+ keepVariant(compare.leftId)} + /> + keepVariant(compare.rightId)} + /> +
+ ) : ( + thread.id === activeId) + ?.inheritedPreviewUrls ?? [] + } + wantCompare={wantCompare} + setWantCompare={setWantCompare} + busy={busy} + onFork={forkOne} + onCompare={startCompare} + onTitle={onActiveTitle} + /> + )} +
+
+ ) +} + +function StudioPane({ + threadId, + inheritedPreviewUrls, + wantCompare, + setWantCompare, + busy, + onFork, + onCompare, + onTitle, +}: { + threadId: string + inheritedPreviewUrls: Array + wantCompare: boolean + setWantCompare: (value: boolean) => void + busy: boolean + onFork: (inheritedPreviewUrls: Array) => void + onCompare: ( + prompt: string, + inheritedPreviewUrls: Array, + ) => Promise + onTitle: (title: string) => void +}) { + const { + messages, + sendMessage, + isLoading, + error: chatError, + } = useChat({ + threadId, + connection, + persistence: true, + }) + const [input, setInput] = useState('') + const preview = useMemo( + () => latestPreview(messages, new Set(inheritedPreviewUrls)), + [inheritedPreviewUrls, messages], + ) + const hasBuiltApp = messages.some( + (message) => + message.role === 'assistant' && + message.parts.some( + (part) => + (part.type === 'text' && Boolean(part.content)) || + part.type === 'tool-call', + ), + ) + + useEffect(() => { + const firstUser = messages.find((message) => message.role === 'user') + const part = firstUser?.parts.find((item) => item.type === 'text') + if (part && 'content' in part && typeof part.content === 'string') { + onTitle(part.content.slice(0, 40)) + } + }, [messages, onTitle]) + + const send = async () => { + const trimmed = input.trim() + if (isLoading || busy) return + if (wantCompare) { + if (!hasBuiltApp) return + setInput('') + await onCompare(trimmed, [...collectPreviewUrls(messages)]) + return + } + if (!trimmed) return + setInput('') + void sendMessage(trimmed) + } + + return ( + <> +
+ + +
+
{ + event.preventDefault() + void send() + }} + > + + {chatError ? ( +

+ {chatError.message} +

+ ) : null} +