diff --git a/.changeset/sandbox-persistence.md b/.changeset/sandbox-persistence.md new file mode 100644 index 000000000..f9d74feb6 --- /dev/null +++ b/.changeset/sandbox-persistence.md @@ -0,0 +1,18 @@ +--- +'@tanstack/ai': minor +'@tanstack/ai-sandbox': minor +'@tanstack/ai-persistence': minor +'@tanstack/ai-persistence-drizzle': minor +'@tanstack/ai-persistence-prisma': minor +'@tanstack/ai-persistence-cloudflare': minor +--- + +Add durable **sandbox persistence**: cross-process / multi-instance resume for `@tanstack/ai-sandbox`, provided by the same `withPersistence` used for chat. + +`withSandbox` consumes `SandboxStore` (which sandbox to resume) and `LockStore` (mutual exclusion around ensure) as optional capabilities, defaulting to in-memory (single-process). This makes them durable without a sandbox-specific middleware: + +- `withPersistence` now provides the `SandboxStoreCapability` (and the shared `LocksCapability`) whenever its store set includes them. Compose `[withPersistence(persistence), withSandbox(sandbox)]`. +- `AIPersistenceStores` gains an optional `sandbox?: SandboxStore`. The backends carry it out of the box: `sqlitePersistence` / `drizzlePersistence(db)` (new `sandboxes` table in the shipped schema + migration), `prismaPersistence(prisma)` (new `Sandbox` model; the delegate resolves lazily so chat-only clients without it keep working), and `cloudflarePersistence({ d1 })` (D1, delegating to Drizzle). `memoryPersistence()` includes an in-memory sandbox store. +- The Cloudflare Durable-Object lock (`durableObjects`) doubles as the distributed sandbox lock. + +**Shared tokens in core.** `SandboxStore` / `SandboxRecord` / `SandboxStoreCapability` / `InMemorySandboxStore` and the `LockStore` / `LocksCapability` / `InMemoryLockStore` primitives now live in core `@tanstack/ai` (their neutral home). `@tanstack/ai-sandbox` and `@tanstack/ai-persistence` re-export them, so one shared token reference lets a persistence-provided store and lock reach `withSandbox` with no dependency between the two packages. A `SandboxStore` conformance testkit is exported from `@tanstack/ai-sandbox/testkit` (`runSandboxStoreConformance`); every backend runs it. diff --git a/docs/config.json b/docs/config.json index ba6aec4e9..35a2391ba 100644 --- a/docs/config.json +++ b/docs/config.json @@ -463,7 +463,7 @@ "label": "Overview", "to": "sandbox/overview", "addedAt": "2026-06-16", - "updatedAt": "2026-07-03" + "updatedAt": "2026-07-23" }, { "label": "Quick Start", @@ -508,6 +508,11 @@ "addedAt": "2026-06-29", "updatedAt": "2026-07-09" }, + { + "label": "Persistence", + "to": "sandbox/persistence", + "addedAt": "2026-07-23" + }, { "label": "Events", "to": "sandbox/events", diff --git a/docs/persistence/cloudflare.md b/docs/persistence/cloudflare.md index 69f42d4e9..ed626dea3 100644 --- a/docs/persistence/cloudflare.md +++ b/docs/persistence/cloudflare.md @@ -217,3 +217,10 @@ export function createComposedPersistence(env: Env) { D1 continues to own messages and metadata. Locks stay on `withLocks` if you need them. Cross-backend transactions are not added by composition; design retries and consistency explicitly. + +## Sandbox persistence + +The same D1 binding and Durable Object lock back durable **sandbox** resume at +the edge: `createD1SandboxStore(env.AI_STATE)` and +`createDurableObjectLockStore(env.AI_LOCKS)` with `withLocks`. See +[Sandbox Persistence](../sandbox/persistence). diff --git a/docs/persistence/sql-backends.md b/docs/persistence/sql-backends.md index 9b2195163..2ed656300 100644 --- a/docs/persistence/sql-backends.md +++ b/docs/persistence/sql-backends.md @@ -86,3 +86,12 @@ For Cloudflare D1, use the Drizzle SQLite path above (or the thin `cloudflarePersistence({ d1 })` wrapper). Durable Object locks live in [Cloudflare Persistence](./cloudflare). For another SQL library, start with [Custom Stores](./custom-stores). + +## Sandbox persistence + +Both adapters also export a durable `SandboxStore` for +[`@tanstack/ai-sandbox`](../sandbox/persistence) resume — +`createDrizzleSandboxStore(db, sandboxesTable)` and +`createPrismaSandboxStore(prisma)`. The `sandboxes` table/model is **optional** +and stays out of the chat BYO schema contract. See +[Sandbox Persistence](../sandbox/persistence). diff --git a/docs/sandbox/overview.md b/docs/sandbox/overview.md index 52adcf8ba..536a0a5ab 100644 --- a/docs/sandbox/overview.md +++ b/docs/sandbox/overview.md @@ -133,7 +133,8 @@ hands back a live preview URL, see `examples/sandbox-web` — one app with harne (Claude Code / Codex / OpenCode / Grok) and provider (Docker / local / Vercel / Daytona) pickers. -> **Persistence-ready:** the sandbox layer ships with in-memory stores for -> resume bookkeeping. A future persistence package can provide durable -> `SandboxStore` / `LockStore` implementations (and event-log replay) by -> supplying those optional capabilities — no changes to the sandbox layer. +> **Durable resume:** the sandbox layer ships with in-memory stores for resume +> bookkeeping (single-process). For cross-process / multi-instance resume, +> [Sandbox Persistence](./persistence) provides durable `SandboxStore` / +> `LockStore` implementations (SQLite/Drizzle, Prisma, Cloudflare D1 + Durable +> Object lock) via `withSandboxPersistence` — no changes to the sandbox layer. diff --git a/docs/sandbox/persistence.md b/docs/sandbox/persistence.md new file mode 100644 index 000000000..57fd5ac86 --- /dev/null +++ b/docs/sandbox/persistence.md @@ -0,0 +1,178 @@ +--- +title: Sandbox Persistence +id: sandbox-persistence +order: 9 +description: "Make sandbox resume durable across processes and instances with a SandboxStore and a shared lock." +--- + +Your agent runs behind more than one server instance, or at the edge. A run +spins up a sandbox, clones the repo, installs deps, does its work. The next run +for the same thread should pick that sandbox back up. Instead it builds a fresh +one every time and pays the whole cold-start cost again. + +[Lifecycle & Snapshots](./lifecycle) already knows how to resume, but its +bookkeeping is in-memory, so it only holds within one process. The moment a run +lands on a different replica (or a fresh isolate), that instance has never seen +the sandbox and re-creates it. + +Sandbox persistence makes resume durable across instances. Two pieces: + +- **`SandboxStore`**: the record of which provider sandbox (and snapshot) to + resume for a given key. Durable, shared across instances. When present on a + `withPersistence` bag, the shared `sandbox-store` capability is provided for + `withSandbox`. +- **`LockStore`**: mutual exclusion around resume-or-create, so two runs for the + same thread don't both create a sandbox. Provided separately with + `withLocks`. Across instances this has to be a distributed lock. + +Both capability tokens live in core `@tanstack/ai` so persistence and sandbox +share the same references — no package-to-package dependency. + +## Wire it up + +Node / single process (SQLite convenience factory includes a sandbox store): + +```ts +import { chat } from '@tanstack/ai' +import { grokBuildText } from '@tanstack/ai-grok-build' +import { withSandbox } from '@tanstack/ai-sandbox' +import { withLocks, withPersistence } from '@tanstack/ai-persistence' +import { InMemoryLockStore } from '@tanstack/ai-persistence' +import { sqlitePersistence } from '@tanstack/ai-persistence-drizzle/sqlite' +import { sandbox } from './sandbox' +import { messages } from './chat-context' + +const persistence = sqlitePersistence({ + url: 'file:.tanstack-ai/state.sqlite', +}) + +chat({ + adapter: grokBuildText('grok-build'), + messages, + middleware: [ + withPersistence(persistence), + withLocks(new InMemoryLockStore()), + withSandbox(sandbox), + ], +}) +``` + +With `reuse: 'thread'` (the default), the first run creates and records the +sandbox. A later run for the same `threadId` resumes it, even on a different +instance (when the store and lock are distributed). + +Chat-only `drizzlePersistence` / `prismaPersistence` / `cloudflarePersistence` +do **not** force a sandbox store into the chat schema contract. Attach one +when you need it: + +```ts +import { + createDrizzleSandboxStore, + defaultSqliteSandboxes, + drizzlePersistence, +} from '@tanstack/ai-persistence-drizzle' + +const chat = drizzlePersistence(db, { provider: 'sqlite', schema }) +const sandboxStore = createDrizzleSandboxStore(db, defaultSqliteSandboxes) + +withPersistence({ + stores: { ...chat.stores, sandbox: sandboxStore }, +}) +``` + +## Choose a backend + +### SQLite / Drizzle (Node) + +`sqlitePersistence` opens Node's built-in SQLite, bootstraps the stock chat +tables **and** the optional `sandboxes` table, and returns +`ChatAndSandboxPersistence`: + +```ts +import { sqlitePersistence } from '@tanstack/ai-persistence-drizzle/sqlite' + +export const persistence = sqlitePersistence({ + url: 'file:.tanstack-ai/state.sqlite', +}) +// persistence.stores.sandbox is ready for withPersistence +``` + +BYO Drizzle database: migrate the stock (or custom) `sandboxes` table yourself, +then call `createDrizzleSandboxStore(db, sandboxesTable)`. + +### Prisma + +Add the `Sandbox` model from the [Prisma fragment](../persistence/prisma), +migrate, then: + +```ts +import { createPrismaSandboxStore } from '@tanstack/ai-persistence-prisma' +import { prisma } from './prisma' + +export const sandboxStore = createPrismaSandboxStore(prisma) +``` + +A chat-only client that never adds `withSandbox` does not need the `Sandbox` +model; the store is resolved lazily on first use. + +### Cloudflare (edge) + +At the edge every run can hit a different isolate, so the distributed lock earns +its keep. D1 carries the sandbox mapping; Durable Objects provide the lock: + +```ts +import { + createD1SandboxStore, + createDurableObjectLockStore, +} from '@tanstack/ai-persistence-cloudflare' +import { withLocks, withPersistence } from '@tanstack/ai-persistence' +import { withSandbox } from '@tanstack/ai-sandbox' +import { env } from './env' +import { chatPersistence } from './chat-persistence' + +const sandboxStore = createD1SandboxStore(env.DB) + +chat({ + // ... + middleware: [ + withPersistence({ + stores: { ...chatPersistence.stores, sandbox: sandboxStore }, + }), + withLocks(createDurableObjectLockStore(env.AI_LOCKS)), + withSandbox(sandbox), + ], +}) +``` + +Migrate the `sandboxes` table into your D1 journal alongside chat tables (see +the stock definition in `@tanstack/ai-persistence-drizzle`). + +## Custom store + +Implement `SandboxStore` and pass it on the persistence bag: + +```ts +import type { SandboxStore } from '@tanstack/ai' + +export const sandboxStore: SandboxStore = { + async get(key) { + /* load SandboxRecord | null */ + }, + async upsert(record) { + /* insert or overwrite by record.key */ + }, + async delete(key) { + /* remove */ + }, +} +``` + +Pair it with a distributed `LockStore` via `withLocks` in multi-instance +deployments. The in-memory lock is only correct within one process. + +## See also + +- [Persistence overview](../persistence/overview) +- [Custom stores](../persistence/custom-stores) +- [Cloudflare persistence](../persistence/cloudflare) +- [Sandbox lifecycle](./lifecycle) diff --git a/examples/persistent-chat-drizzle/drizzle/0000_loving_snowbird.sql b/examples/persistent-chat-drizzle/drizzle/0000_loving_snowbird.sql index a37beda6e..8f6a8d796 100644 --- a/examples/persistent-chat-drizzle/drizzle/0000_loving_snowbird.sql +++ b/examples/persistent-chat-drizzle/drizzle/0000_loving_snowbird.sql @@ -32,3 +32,13 @@ CREATE TABLE `runs` ( `error` text, `usage_json` text ); +--> statement-breakpoint +CREATE TABLE `sandboxes` ( + `key` text PRIMARY KEY NOT NULL, + `provider` text NOT NULL, + `provider_sandbox_id` text NOT NULL, + `latest_snapshot_id` text, + `thread_id` text NOT NULL, + `latest_run_id` text, + `updated_at` integer NOT NULL +); diff --git a/examples/persistent-chat-drizzle/src/db/tanstack-ai-schema.ts b/examples/persistent-chat-drizzle/src/db/tanstack-ai-schema.ts index 15a3af11c..032ce8fff 100644 --- a/examples/persistent-chat-drizzle/src/db/tanstack-ai-schema.ts +++ b/examples/persistent-chat-drizzle/src/db/tanstack-ai-schema.ts @@ -82,10 +82,26 @@ export const metadata = sqliteTable( (table) => [primaryKey({ columns: [table.scope, table.key] })], ) +/** + * Persisted sandbox resume records (`@tanstack/ai-sandbox`'s `SandboxStore`). + * Keyed by the compound sandbox key; used by `createDrizzleSandboxStore`, + * independent of the chat stores. + */ +export const sandboxes = sqliteTable('sandboxes', { + key: text('key').primaryKey(), + provider: text('provider').notNull(), + providerSandboxId: text('provider_sandbox_id').notNull(), + latestSnapshotId: text('latest_snapshot_id'), + threadId: text('thread_id').notNull(), + latestRunId: text('latest_run_id'), + updatedAt: integer('updated_at').notNull(), +}) + /** The full state schema, for `drizzlePersistence(db, { schema })` and drizzle-kit. */ export const schema = { messages, runs, interrupts, metadata, + sandboxes, } diff --git a/packages/ai-persistence-cloudflare/package.json b/packages/ai-persistence-cloudflare/package.json index b2a2cfcf3..b7c0dd060 100644 --- a/packages/ai-persistence-cloudflare/package.json +++ b/packages/ai-persistence-cloudflare/package.json @@ -52,6 +52,7 @@ "@tanstack/ai": "workspace:*", "@tanstack/ai-persistence": "workspace:*", "@tanstack/ai-persistence-drizzle": "workspace:*", + "@tanstack/ai-sandbox": "workspace:*", "@vitest/coverage-v8": "4.0.14", "drizzle-orm": "^0.45.0", "miniflare": "^4.20260609.0" diff --git a/packages/ai-persistence-cloudflare/src/d1.ts b/packages/ai-persistence-cloudflare/src/d1.ts index 3bd00c55c..67daddc45 100644 --- a/packages/ai-persistence-cloudflare/src/d1.ts +++ b/packages/ai-persistence-cloudflare/src/d1.ts @@ -1,11 +1,14 @@ import { drizzle } from 'drizzle-orm/d1' import { createDefaultSqliteSchema, + createDrizzleSandboxStore, + defaultSqliteSandboxes, drizzlePersistence, } from '@tanstack/ai-persistence-drizzle' +import type { SandboxStore } from '@tanstack/ai' /** - * Create the structured stores over a Cloudflare D1 binding. + * Create the structured chat stores over a Cloudflare D1 binding. * * Thin wrapper: stock SQLite schema + `drizzle-orm/d1` + * {@link drizzlePersistence}. This package does **not** ship or apply DDL — @@ -26,3 +29,17 @@ export function createD1Stores(d1: D1Database) { metadata: persistence.stores.metadata, } } + +/** + * Durable {@link SandboxStore} over a migrated Cloudflare D1 binding (delegates + * to the Drizzle sandbox store). Pair with `createDurableObjectLockStore` for a + * multi-instance-correct sandbox resume on the edge. The `sandboxes` table is + * **not** part of the chat BYO schema — migrate it separately (or use the stock + * definition from `@tanstack/ai-persistence-drizzle`). + */ +export function createD1SandboxStore(d1: D1Database): SandboxStore { + return createDrizzleSandboxStore( + drizzle(d1, { schema: { sandboxes: defaultSqliteSandboxes } }), + defaultSqliteSandboxes, + ) +} diff --git a/packages/ai-persistence-cloudflare/src/index.ts b/packages/ai-persistence-cloudflare/src/index.ts index b90b06ba5..82edf5d37 100644 --- a/packages/ai-persistence-cloudflare/src/index.ts +++ b/packages/ai-persistence-cloudflare/src/index.ts @@ -1,7 +1,7 @@ import { createD1Stores } from './d1' import type { ChatPersistence } from '@tanstack/ai-persistence' -export { createD1Stores } from './d1' +export { createD1Stores, createD1SandboxStore } from './d1' export { CloudflareLockDurableObject, createDurableObjectLockStore, @@ -25,7 +25,8 @@ export type { * * Locks are a separate concern: use {@link createDurableObjectLockStore} with * `withLocks` from `@tanstack/ai-persistence` when you need multi-instance - * coordination. + * coordination. Sandbox resume uses {@link createD1SandboxStore} + the same + * shared lock token. */ export interface CloudflarePersistenceOptions { d1: D1Database @@ -38,7 +39,8 @@ export interface CloudflarePersistenceOptions { * `@tanstack/ai-persistence-drizzle`. Prefer owning that schema in your app * and calling `drizzlePersistence` directly when you need renames or extra * columns. Durable Object locks are separate via - * {@link createDurableObjectLockStore}. + * {@link createDurableObjectLockStore}. Sandbox resume is separate via + * {@link createD1SandboxStore}. */ export function cloudflarePersistence( options: CloudflarePersistenceOptions, diff --git a/packages/ai-persistence-cloudflare/tests/sandbox-store.conformance.test.ts b/packages/ai-persistence-cloudflare/tests/sandbox-store.conformance.test.ts new file mode 100644 index 000000000..9120c9746 --- /dev/null +++ b/packages/ai-persistence-cloudflare/tests/sandbox-store.conformance.test.ts @@ -0,0 +1,38 @@ +/// +import { afterAll, beforeAll } from 'vitest' +import { Miniflare } from 'miniflare' +import { runSandboxStoreConformance } from '@tanstack/ai-sandbox/testkit' +import { createD1SandboxStore } from '../src/index' + +interface RuntimeBindings { + AI_DB: D1Database +} + +let miniflare: Miniflare +let db: D1Database + +const SANDBOXES_DDL = + 'CREATE TABLE IF NOT EXISTS sandboxes (key TEXT PRIMARY KEY NOT NULL, provider TEXT NOT NULL, provider_sandbox_id TEXT NOT NULL, latest_snapshot_id TEXT, thread_id TEXT NOT NULL, latest_run_id TEXT, updated_at INTEGER NOT NULL)' + +beforeAll(async () => { + miniflare = new Miniflare({ + compatibilityDate: '2026-06-24', + d1Databases: ['AI_DB'], + modules: true, + script: 'export default { fetch() { return new Response("ok") } }', + }) + const bindings = await miniflare.getBindings() + db = bindings.AI_DB + await db.prepare(SANDBOXES_DDL).run() +}) + +afterAll(async () => { + await miniflare.dispose() +}) + +// One D1 binding is shared; truncating the table per store keeps each +// conformance case isolated. +runSandboxStoreConformance('cloudflare-d1', async () => { + await db.prepare('DELETE FROM sandboxes').run() + return createD1SandboxStore(db) +}) diff --git a/packages/ai-persistence-drizzle/package.json b/packages/ai-persistence-drizzle/package.json index 06f4ff242..01714e257 100644 --- a/packages/ai-persistence-drizzle/package.json +++ b/packages/ai-persistence-drizzle/package.json @@ -68,6 +68,7 @@ "@electric-sql/pglite": "^0.3.0", "@tanstack/ai": "workspace:*", "@tanstack/ai-persistence": "workspace:*", + "@tanstack/ai-sandbox": "workspace:*", "@vitest/coverage-v8": "4.0.14", "drizzle-orm": "^0.45.0" } diff --git a/packages/ai-persistence-drizzle/src/index.ts b/packages/ai-persistence-drizzle/src/index.ts index 6ace6326c..4139392ad 100644 --- a/packages/ai-persistence-drizzle/src/index.ts +++ b/packages/ai-persistence-drizzle/src/index.ts @@ -45,3 +45,14 @@ export type { TanstackAiSchema, } from './core/schema-contract' export type { DrizzleSqliteDb, DrizzleDb } from './core/stores' + +/** Durable sandbox resume store (optional; not part of the chat schema contract). */ +export { createDrizzleSandboxStore } from './sandbox-store' +export type { + SandboxTable, + SandboxTableShapes, + SqliteSandboxTable, + PgSandboxTable, +} from './sandbox-store' +export { sandboxes as defaultSqliteSandboxes } from './sqlite/default-schema' +export { sandboxes as defaultPgSandboxes } from './pg/default-schema' diff --git a/packages/ai-persistence-drizzle/src/pg/assets/tanstack-ai-schema.ts b/packages/ai-persistence-drizzle/src/pg/assets/tanstack-ai-schema.ts index 41c3b203d..d760940dc 100644 --- a/packages/ai-persistence-drizzle/src/pg/assets/tanstack-ai-schema.ts +++ b/packages/ai-persistence-drizzle/src/pg/assets/tanstack-ai-schema.ts @@ -39,9 +39,7 @@ import type { ModelMessage, TokenUsage } from '@tanstack/ai' /** Thread message history (`MessageStore`). */ export const messages = pgTable('messages', { threadId: text('thread_id').primaryKey(), - messagesJson: jsonb('messages_json') - .$type>() - .notNull(), + messagesJson: jsonb('messages_json').$type>().notNull(), }) /** Run lifecycle records (`RunStore`). */ diff --git a/packages/ai-persistence-drizzle/src/pg/default-schema.ts b/packages/ai-persistence-drizzle/src/pg/default-schema.ts index 86396c2be..3a52b7109 100644 --- a/packages/ai-persistence-drizzle/src/pg/default-schema.ts +++ b/packages/ai-persistence-drizzle/src/pg/default-schema.ts @@ -30,9 +30,7 @@ import type { TanstackAiPgSchema } from '../core/schema-contract' /** Thread message history (`MessageStore`). */ export const messages = pgTable('messages', { threadId: text('thread_id').primaryKey(), - messagesJson: jsonb('messages_json') - .$type>() - .notNull(), + messagesJson: jsonb('messages_json').$type>().notNull(), }) /** Run lifecycle records (`RunStore`). */ @@ -80,6 +78,24 @@ export const metadata = pgTable( (table) => [primaryKey({ columns: [table.scope, table.key] })], ) +/** + * Persisted sandbox resume records (`@tanstack/ai-sandbox`'s `SandboxStore`). + * Keyed by the compound sandbox key; small metadata only (no blobs). + * + * **Not** part of {@link TanstackAiPgSchema} / the chat BYO contract — + * chat-only apps never need this table. Pass it to + * {@link createDrizzleSandboxStore} when you want durable sandbox resume. + */ +export const sandboxes = pgTable('sandboxes', { + key: text('key').primaryKey(), + provider: text('provider').notNull(), + providerSandboxId: text('provider_sandbox_id').notNull(), + latestSnapshotId: text('latest_snapshot_id'), + threadId: text('thread_id').notNull(), + latestRunId: text('latest_run_id'), + updatedAt: bigint('updated_at', { mode: 'number' }).notNull(), +}) + /** * Build a fresh copy of the default SQLite schema tables. * diff --git a/packages/ai-persistence-drizzle/src/sandbox-store.ts b/packages/ai-persistence-drizzle/src/sandbox-store.ts new file mode 100644 index 000000000..d7fa5bdd0 --- /dev/null +++ b/packages/ai-persistence-drizzle/src/sandbox-store.ts @@ -0,0 +1,118 @@ +/** + * Durable {@link SandboxStore} over a Drizzle database. + * + * Backs `@tanstack/ai-sandbox`'s resume-or-create ensure: maps a compound + * sandbox key to the provider sandbox (and latest snapshot) that should be + * resumed. Independent of the chat schema contract — pass the sandboxes table + * explicitly (or use the stock default from {@link createDefaultSqliteSchema} + * / {@link createDefaultPgSchema} when you opt into sandbox columns). + * + * Multi-instance correctness additionally needs a distributed lock (e.g. the + * Cloudflare Durable Object lock); this store only persists the mapping. + */ +import { eq } from 'drizzle-orm' +import type { SandboxRecord, SandboxStore } from '@tanstack/ai' +import type { DrizzleSqliteDb } from './core/stores' +import type { AnySQLiteColumn, SQLiteTable } from 'drizzle-orm/sqlite-core' +import type { AnyPgColumn, PgTable } from 'drizzle-orm/pg-core' + +/** Column data shapes for a sandboxes table. */ +export interface SandboxTableShapes { + key: string + provider: string + providerSandboxId: string + latestSnapshotId: string + threadId: string + latestRunId: string + updatedAt: number +} + +/** SQLite sandboxes table projection. */ +export type SqliteSandboxTable = SQLiteTable & { + [ColumnKey in keyof SandboxTableShapes]: AnySQLiteColumn<{ + data: SandboxTableShapes[ColumnKey] + }> +} + +/** Postgres sandboxes table projection. */ +export type PgSandboxTable = PgTable & { + [ColumnKey in keyof SandboxTableShapes]: AnyPgColumn<{ + data: SandboxTableShapes[ColumnKey] + }> +} + +export type SandboxTable = SqliteSandboxTable | PgSandboxTable + +type SandboxRow = { + key: string + provider: string + providerSandboxId: string + latestSnapshotId: string | null + threadId: string + latestRunId: string | null + updatedAt: number +} + +function mapSandbox(row: SandboxRow): SandboxRecord { + return { + key: row.key, + provider: row.provider, + providerSandboxId: row.providerSandboxId, + ...(row.latestSnapshotId != null + ? { latestSnapshotId: row.latestSnapshotId } + : {}), + threadId: row.threadId, + ...(row.latestRunId != null ? { latestRunId: row.latestRunId } : {}), + updatedAt: row.updatedAt, + } +} + +/** + * Wire a durable {@link SandboxStore} over a sandboxes table. + * + * Chat-only schemas omit this table; only apps that compose sandbox + * persistence need it. + */ +export function createDrizzleSandboxStore( + db: DrizzleSqliteDb, + sandboxes: SandboxTable, +): SandboxStore { + return { + async get(key) { + const rows = (await db + .select() + .from(sandboxes as never) + .where(eq(sandboxes.key as never, key))) as Array + const row = rows[0] + return row ? mapSandbox(row) : null + }, + async upsert(record) { + const values = { + key: record.key, + provider: record.provider, + providerSandboxId: record.providerSandboxId, + latestSnapshotId: record.latestSnapshotId ?? null, + threadId: record.threadId, + latestRunId: record.latestRunId ?? null, + updatedAt: record.updatedAt, + } + await db + .insert(sandboxes as never) + .values(values as never) + .onConflictDoUpdate({ + target: sandboxes.key as never, + set: { + provider: values.provider, + providerSandboxId: values.providerSandboxId, + latestSnapshotId: values.latestSnapshotId, + threadId: values.threadId, + latestRunId: values.latestRunId, + updatedAt: values.updatedAt, + } as never, + }) + }, + async delete(key) { + await db.delete(sandboxes as never).where(eq(sandboxes.key as never, key)) + }, + } +} diff --git a/packages/ai-persistence-drizzle/src/sqlite.ts b/packages/ai-persistence-drizzle/src/sqlite.ts index 30d68a3ef..1d51ab637 100644 --- a/packages/ai-persistence-drizzle/src/sqlite.ts +++ b/packages/ai-persistence-drizzle/src/sqlite.ts @@ -9,3 +9,5 @@ export { sqlitePersistence, } from './sqlite/factory' export type { SqlitePersistenceOptions } from './sqlite/factory' + +export { sandboxes } from './sqlite/default-schema' diff --git a/packages/ai-persistence-drizzle/src/sqlite/default-schema.ts b/packages/ai-persistence-drizzle/src/sqlite/default-schema.ts index 316ae479c..ef6249ccf 100644 --- a/packages/ai-persistence-drizzle/src/sqlite/default-schema.ts +++ b/packages/ai-persistence-drizzle/src/sqlite/default-schema.ts @@ -76,6 +76,24 @@ export const metadata = sqliteTable( (table) => [primaryKey({ columns: [table.scope, table.key] })], ) +/** + * Persisted sandbox resume records (`@tanstack/ai-sandbox`'s `SandboxStore`). + * Keyed by the compound sandbox key; small metadata only (no blobs). + * + * **Not** part of {@link TanstackAiSqliteSchema} / the chat BYO contract — + * chat-only apps never need this table. Pass it to + * {@link createDrizzleSandboxStore} when you want durable sandbox resume. + */ +export const sandboxes = sqliteTable('sandboxes', { + key: text('key').primaryKey(), + provider: text('provider').notNull(), + providerSandboxId: text('provider_sandbox_id').notNull(), + latestSnapshotId: text('latest_snapshot_id'), + threadId: text('thread_id').notNull(), + latestRunId: text('latest_run_id'), + updatedAt: integer('updated_at').notNull(), +}) + /** * Build a fresh copy of the default SQLite schema tables. * diff --git a/packages/ai-persistence-drizzle/src/sqlite/ensure-tables.ts b/packages/ai-persistence-drizzle/src/sqlite/ensure-tables.ts index 9e32765e7..3c0f50d8d 100644 --- a/packages/ai-persistence-drizzle/src/sqlite/ensure-tables.ts +++ b/packages/ai-persistence-drizzle/src/sqlite/ensure-tables.ts @@ -69,3 +69,15 @@ export function ensureSqliteTables( for (const indexSql of createIndexSql(table)) exec(indexSql) } } + +/** + * Bootstrap a single extra table (e.g. the optional sandboxes table) the same + * way {@link ensureSqliteTables} does for the chat schema. + */ +export function ensureSqliteTable( + exec: (sql: string) => void, + table: SQLiteTable, +): void { + exec(createTableSql(table)) + for (const indexSql of createIndexSql(table)) exec(indexSql) +} diff --git a/packages/ai-persistence-drizzle/src/sqlite/factory.ts b/packages/ai-persistence-drizzle/src/sqlite/factory.ts index d5d8596c8..36adb33dc 100644 --- a/packages/ai-persistence-drizzle/src/sqlite/factory.ts +++ b/packages/ai-persistence-drizzle/src/sqlite/factory.ts @@ -17,9 +17,10 @@ import { fileURLToPath } from 'node:url' import { DatabaseSync } from 'node:sqlite' import { drizzle } from 'drizzle-orm/sqlite-proxy' import { drizzlePersistence } from '../core/persistence' -import { createDefaultSqliteSchema } from './default-schema' -import { ensureSqliteTables } from './ensure-tables' -import type { ChatPersistence } from '@tanstack/ai-persistence' +import { createDefaultSqliteSchema, sandboxes } from './default-schema' +import { createDrizzleSandboxStore } from '../sandbox-store' +import { ensureSqliteTables, ensureSqliteTable } from './ensure-tables' +import type { ChatAndSandboxPersistence } from '@tanstack/ai-persistence' import type { TanstackAiSqliteSchema } from '../core/schema-contract' export { createDefaultSqliteSchema } from './default-schema' @@ -66,7 +67,9 @@ export interface SqlitePersistenceOptions { */ export function sqlitePersistence( options: SqlitePersistenceOptions, -): ChatPersistence & { +): ChatAndSandboxPersistence & { + /** Underlying Drizzle database (for optional extra stores). */ + db: ReturnType /** Close the underlying Node SQLite connection. */ close: () => void } { @@ -77,6 +80,9 @@ export function sqlitePersistence( try { if (options.ensureTables !== false) { ensureSqliteTables((sql) => sqlite.exec(sql), schema) + // Sandboxes stay outside the chat schema contract but the Node factory + // bootstraps the stock table so local/dev sandbox resume works zero-config. + ensureSqliteTable((sql) => sqlite.exec(sql), sandboxes) } } catch (error) { sqlite.close() @@ -100,7 +106,11 @@ export function sqlitePersistence( const persistence = drizzlePersistence(db, { provider: 'sqlite', schema }) let closed = false return { - ...persistence, + stores: { + ...persistence.stores, + sandbox: createDrizzleSandboxStore(db, sandboxes), + }, + db, close() { if (closed) return sqlite.close() diff --git a/packages/ai-persistence-drizzle/tests/api-types.test-d.ts b/packages/ai-persistence-drizzle/tests/api-types.test-d.ts index 685b275ac..f11520e68 100644 --- a/packages/ai-persistence-drizzle/tests/api-types.test-d.ts +++ b/packages/ai-persistence-drizzle/tests/api-types.test-d.ts @@ -2,6 +2,7 @@ import { expectTypeOf } from 'vitest' import type { DrizzleD1Database } from 'drizzle-orm/d1' import type { PgliteDatabase } from 'drizzle-orm/pglite' import type { + ChatAndSandboxPersistence, ChatPersistence, MessageStore, RunStore, @@ -50,7 +51,7 @@ drizzlePersistence(pgDatabase, { provider: 'pg', schema: defaultSchema }) drizzlePersistence(d1Database, { schema: defaultSchema }) const sqlite = sqlitePersistence({ url: ':memory:' }) -expectTypeOf(sqlite.stores).toEqualTypeOf() +expectTypeOf(sqlite.stores).toEqualTypeOf() expectTypeOf(sqlite.close).toEqualTypeOf<() => void>() // Default factories, emitted assets, and renamed/extended variant all satisfy diff --git a/packages/ai-persistence-drizzle/tests/sandbox-store.conformance.test.ts b/packages/ai-persistence-drizzle/tests/sandbox-store.conformance.test.ts new file mode 100644 index 000000000..e82e90cff --- /dev/null +++ b/packages/ai-persistence-drizzle/tests/sandbox-store.conformance.test.ts @@ -0,0 +1,9 @@ +import { runSandboxStoreConformance } from '@tanstack/ai-sandbox/testkit' +import { createDrizzleSandboxStore, defaultSqliteSandboxes } from '../src/index' +import { sqlitePersistence } from '../src/sqlite' + +// A fresh in-memory database per store keeps each conformance case isolated. +runSandboxStoreConformance('drizzle-sqlite', () => { + const persistence = sqlitePersistence({ url: ':memory:' }) + return createDrizzleSandboxStore(persistence.db, defaultSqliteSandboxes) +}) diff --git a/packages/ai-persistence-prisma/package.json b/packages/ai-persistence-prisma/package.json index 46e873719..9dced6942 100644 --- a/packages/ai-persistence-prisma/package.json +++ b/packages/ai-persistence-prisma/package.json @@ -85,6 +85,7 @@ "@prisma/client": "^6.19.3", "@tanstack/ai": "workspace:*", "@tanstack/ai-persistence": "workspace:*", + "@tanstack/ai-sandbox": "workspace:*", "@vitest/coverage-v8": "4.0.14", "prisma": "^6.19.3" } diff --git a/packages/ai-persistence-prisma/prisma/tanstack-ai.prisma b/packages/ai-persistence-prisma/prisma/tanstack-ai.prisma index aa716415b..698df5365 100644 --- a/packages/ai-persistence-prisma/prisma/tanstack-ai.prisma +++ b/packages/ai-persistence-prisma/prisma/tanstack-ai.prisma @@ -48,3 +48,16 @@ model Metadata { @@id([scope, key]) @@map("metadata") } + +/// Persisted sandbox resume records (`@tanstack/ai-sandbox`'s `SandboxStore`). +model Sandbox { + key String @id + provider String + providerSandboxId String @map("provider_sandbox_id") + latestSnapshotId String? @map("latest_snapshot_id") + threadId String @map("thread_id") + latestRunId String? @map("latest_run_id") + updatedAt BigInt @map("updated_at") + + @@map("sandboxes") +} diff --git a/packages/ai-persistence-prisma/src/assets/tanstack-ai.prisma b/packages/ai-persistence-prisma/src/assets/tanstack-ai.prisma index aa716415b..698df5365 100644 --- a/packages/ai-persistence-prisma/src/assets/tanstack-ai.prisma +++ b/packages/ai-persistence-prisma/src/assets/tanstack-ai.prisma @@ -48,3 +48,16 @@ model Metadata { @@id([scope, key]) @@map("metadata") } + +/// Persisted sandbox resume records (`@tanstack/ai-sandbox`'s `SandboxStore`). +model Sandbox { + key String @id + provider String + providerSandboxId String @map("provider_sandbox_id") + latestSnapshotId String? @map("latest_snapshot_id") + threadId String @map("thread_id") + latestRunId String? @map("latest_run_id") + updatedAt BigInt @map("updated_at") + + @@map("sandboxes") +} diff --git a/packages/ai-persistence-prisma/src/index.ts b/packages/ai-persistence-prisma/src/index.ts index 65ef10c3a..947df222b 100644 --- a/packages/ai-persistence-prisma/src/index.ts +++ b/packages/ai-persistence-prisma/src/index.ts @@ -18,6 +18,8 @@ import type { PrismaModelMap } from './model-contract' export { prismaModels, prismaModelsFilename } from './models' export { PrismaModelError } from './model-contract' export type { PrismaModelMap } from './model-contract' +export { createPrismaSandboxStore } from './sandbox-store' +export type { PrismaSandboxStoreOptions } from './sandbox-store' /** * Structural stand-in for a generated Prisma client. @@ -49,13 +51,12 @@ export interface PrismaPersistenceOptions { /** * Wire TanStack AI persistence stores over a migrated Prisma client. * - * State stores only — locks are a separate concern. For multi-instance - * coordination use `withLocks` with a distributed `LockStore` (for example - * `createDurableObjectLockStore` from `@tanstack/ai-persistence-cloudflare`). - * Do not invent an in-process lock and pretend it coordinates across instances. - */ -/** * Returns {@link ChatPersistence} (messages + runs + interrupts + metadata). + * Sandbox resume is separate via {@link createPrismaSandboxStore} so chat-only + * clients need not include the `Sandbox` model. + * + * Locks are a separate concern. For multi-instance coordination use `withLocks` + * with a distributed `LockStore`. */ export function prismaPersistence( prisma: PrismaClientLike, diff --git a/packages/ai-persistence-prisma/src/model-contract.ts b/packages/ai-persistence-prisma/src/model-contract.ts index 1096437fa..5a1deaa21 100644 --- a/packages/ai-persistence-prisma/src/model-contract.ts +++ b/packages/ai-persistence-prisma/src/model-contract.ts @@ -50,6 +50,16 @@ export interface MetadataRow { valueJson: string } +export interface SandboxRow { + key: string + provider: string + providerSandboxId: string + latestSnapshotId: string | null + threadId: string + latestRunId: string | null + updatedAt: bigint +} + export interface MessageDelegate { findUnique: (args: { where: { threadId: string } @@ -129,6 +139,31 @@ export interface MetadataDelegate { }) => Promise } +export interface SandboxDelegate { + findUnique: (args: { where: { key: string } }) => Promise + upsert: (args: { + where: { key: string } + create: { + key: string + provider: string + providerSandboxId: string + latestSnapshotId: string | null + threadId: string + latestRunId: string | null + updatedAt: bigint + } + update: { + provider: string + providerSandboxId: string + latestSnapshotId: string | null + threadId: string + latestRunId: string | null + updatedAt: bigint + } + }) => Promise + deleteMany: (args: { where: { key: string } }) => Promise +} + /** The delegate set the stores operate over. */ export interface TanstackAiDelegates { message: MessageDelegate @@ -147,9 +182,12 @@ export interface PrismaModelMap { runs?: string interrupts?: string metadata?: string + sandbox?: string } -const defaultModelNames: Required = { +type ChatPrismaModelMap = Omit + +const defaultModelNames: Required = { messages: 'message', runs: 'run', interrupts: 'interrupt', @@ -157,11 +195,11 @@ const defaultModelNames: Required = { } const storeKeys = Object.keys(defaultModelNames) as Array< - keyof Required + keyof Required > const delegateKeyByStoreKey: { - [K in keyof Required]: keyof TanstackAiDelegates + [K in keyof Required]: keyof TanstackAiDelegates } = { messages: 'message', runs: 'run', @@ -219,3 +257,23 @@ export function resolveDelegates( // which the generated client's delegates satisfy (asserted in type tests). return resolved as TanstackAiDelegates } + +/** + * Resolve the sandbox model delegate off the client. Independent of the four + * chat delegates (sandbox persistence is its own store), so it takes the + * client accessor name directly (`sandbox` by default). + */ +export function resolveSandboxDelegate( + client: object, + modelName = 'sandbox', +): SandboxDelegate { + const candidate: unknown = Reflect.get(client, modelName) + if (!isDelegate(candidate)) { + throw new PrismaModelError([ + `sandbox store maps to \`client.${modelName}\`, which is not a Prisma model delegate.`, + ]) + } + // Safe narrowing: same rationale as resolveDelegates — shape verified here, + // method signatures pinned by SandboxDelegate and asserted in type tests. + return candidate as SandboxDelegate +} diff --git a/packages/ai-persistence-prisma/src/sandbox-store.ts b/packages/ai-persistence-prisma/src/sandbox-store.ts new file mode 100644 index 000000000..fcc42f5e2 --- /dev/null +++ b/packages/ai-persistence-prisma/src/sandbox-store.ts @@ -0,0 +1,73 @@ +/** + * Durable {@link SandboxStore} over a Prisma client. + * + * Backs `@tanstack/ai-sandbox`'s resume-or-create ensure. Epoch-ms `updatedAt` + * is stored as a provider-neutral `BigInt`. Multi-instance correctness + * additionally needs a distributed lock; this store only persists the mapping. + * + * The `Sandbox` delegate is resolved LAZILY, on first use. Chat-only clients + * whose schema omits the `Sandbox` model never trigger resolution (it only + * runs when `withSandbox` uses the store). + */ +import { resolveSandboxDelegate } from './model-contract' +import type { SandboxDelegate, SandboxRow } from './model-contract' +import type { SandboxRecord, SandboxStore } from '@tanstack/ai' + +function mapSandbox(row: SandboxRow): SandboxRecord { + return { + key: row.key, + provider: row.provider, + providerSandboxId: row.providerSandboxId, + ...(row.latestSnapshotId != null + ? { latestSnapshotId: row.latestSnapshotId } + : {}), + threadId: row.threadId, + ...(row.latestRunId != null ? { latestRunId: row.latestRunId } : {}), + updatedAt: Number(row.updatedAt), + } +} + +/** Options for {@link createPrismaSandboxStore}. */ +export interface PrismaSandboxStoreOptions { + /** + * Client accessor name for the sandbox model if you renamed it in your copy + * of the fragment (default `sandbox`, i.e. `model Sandbox`). + */ + model?: string +} + +/** Wire a durable {@link SandboxStore} over a migrated Prisma client. */ +export function createPrismaSandboxStore( + prisma: object, + options?: PrismaSandboxStoreOptions, +): SandboxStore { + // Resolve the delegate on first use so a chat-only client without the + // `Sandbox` model isn't rejected at construction time. + let delegate: SandboxDelegate | undefined + const sandbox = (): SandboxDelegate => + (delegate ??= resolveSandboxDelegate(prisma, options?.model)) + return { + async get(key) { + const row = await sandbox().findUnique({ where: { key } }) + return row ? mapSandbox(row) : null + }, + async upsert(record) { + const fields = { + provider: record.provider, + providerSandboxId: record.providerSandboxId, + latestSnapshotId: record.latestSnapshotId ?? null, + threadId: record.threadId, + latestRunId: record.latestRunId ?? null, + updatedAt: BigInt(record.updatedAt), + } + await sandbox().upsert({ + where: { key: record.key }, + create: { key: record.key, ...fields }, + update: fields, + }) + }, + async delete(key) { + await sandbox().deleteMany({ where: { key } }) + }, + } +} diff --git a/packages/ai-persistence-prisma/tests/sandbox-store.conformance.test.ts b/packages/ai-persistence-prisma/tests/sandbox-store.conformance.test.ts new file mode 100644 index 000000000..6a1fd5a36 --- /dev/null +++ b/packages/ai-persistence-prisma/tests/sandbox-store.conformance.test.ts @@ -0,0 +1,59 @@ +import { rm } from 'node:fs/promises' +import { mkdtempSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { DatabaseSync } from 'node:sqlite' +import { afterAll } from 'vitest' +import { PrismaClient } from '@prisma/client' +import { runSandboxStoreConformance } from '@tanstack/ai-sandbox/testkit' +import { createPrismaSandboxStore } from '../src/index' + +const clients: Array = [] +const temporaryDirectories: Array = [] + +const sandboxTestSchema = ` + CREATE TABLE sandboxes ( + key TEXT NOT NULL PRIMARY KEY, + provider TEXT NOT NULL, + provider_sandbox_id TEXT NOT NULL, + latest_snapshot_id TEXT, + thread_id TEXT NOT NULL, + latest_run_id TEXT, + updated_at BIGINT NOT NULL + ); +` + +function initializeSqliteTestDatabase(path: string): void { + const database = new DatabaseSync(path) + try { + database.exec(sandboxTestSchema) + } finally { + database.close() + } +} + +/** A PrismaClient over a fresh initialized temporary SQLite database. */ +async function makeTestClient(): Promise { + const dir = mkdtempSync(join(tmpdir(), 'tanstack-ai-prisma-sandbox-')) + temporaryDirectories.push(dir) + const dbPath = join(dir, 'state.db').replace(/\\/g, '/') + initializeSqliteTestDatabase(dbPath) + const prisma = new PrismaClient({ + datasources: { db: { url: `file:${dbPath}` } }, + }) + clients.push(prisma) + return prisma +} + +afterAll(async () => { + await Promise.all(clients.map((client) => client.$disconnect())) + await Promise.all( + temporaryDirectories.map((directory) => + rm(directory, { recursive: true, force: true }), + ), + ) +}) + +runSandboxStoreConformance('prisma', async () => + createPrismaSandboxStore(await makeTestClient()), +) diff --git a/packages/ai-persistence/src/capabilities.ts b/packages/ai-persistence/src/capabilities.ts index 0236a8d10..003e87f0d 100644 --- a/packages/ai-persistence/src/capabilities.ts +++ b/packages/ai-persistence/src/capabilities.ts @@ -1,9 +1,10 @@ /** * Persistence capability tokens. * - * `withPersistence` PROVIDES persistence/interrupts so later middleware can - * read durable state. Locks are a separate concern: `withLocks` PROVIDES - * `LocksCapability` (defined in `./locks`). + * `withPersistence` PROVIDES persistence/interrupts (and optionally the shared + * `sandbox-store` token) so later middleware can read durable state. Locks are a + * separate concern: `withLocks` PROVIDES `LocksCapability` (re-exported from + * core via `./locks`). */ import { createCapability } from '@tanstack/ai' import type { AIPersistence, InterruptStore } from './types' @@ -18,5 +19,10 @@ export const InterruptsCapability = createCapability()( export const [getPersistence, providePersistence] = PersistenceCapability export const [getInterrupts, provideInterrupts] = InterruptsCapability -// Locks token lives in ./locks; re-export for a single import surface. +// Shared tokens from core (locks via ./locks so the re-export path is local). export { LocksCapability, getLocks, provideLocks } from './locks' +export { + SandboxStoreCapability, + getSandboxStore, + provideSandboxStore, +} from '@tanstack/ai' diff --git a/packages/ai-persistence/src/index.ts b/packages/ai-persistence/src/index.ts index e06abd68f..9e6885b0b 100644 --- a/packages/ai-persistence/src/index.ts +++ b/packages/ai-persistence/src/index.ts @@ -16,6 +16,8 @@ export type { ChatTranscriptPersistence, ChatPersistence, ChatWithInterruptsPersistence, + ChatAndSandboxPersistenceStores, + ChatAndSandboxPersistence, AIPersistence, AIPersistenceOverrides, ComposedAIPersistenceStores, @@ -44,7 +46,9 @@ export { memoryPersistence } from './memory' export { createInterruptController } from './interrupts' export type { InterruptController } from './interrupts' -// Capabilities +// Capabilities. Locks + sandbox-store tokens are re-exported from core +// `@tanstack/ai` (their neutral home) so a persistence-provided store reaches +// the sandbox layer through the same token reference. export { PersistenceCapability, InterruptsCapability, @@ -55,8 +59,15 @@ export { LocksCapability, getLocks, provideLocks, + SandboxStoreCapability, + getSandboxStore, + provideSandboxStore, } from './capabilities' // Lock primitive (separate from state stores; provide via withLocks) export { InMemoryLockStore } from './locks' export type { LockStore } from './locks' + +// Sandbox-store primitive (re-exported from core) +export { InMemorySandboxStore } from '@tanstack/ai' +export type { SandboxStore, SandboxRecord } from '@tanstack/ai' diff --git a/packages/ai-persistence/src/locks.ts b/packages/ai-persistence/src/locks.ts index dfac8ff85..8c5df8684 100644 --- a/packages/ai-persistence/src/locks.ts +++ b/packages/ai-persistence/src/locks.ts @@ -1,71 +1,18 @@ /** - * Distributed-mutex primitive for multi-instance coordination. + * The `'locks'` capability lives in core `@tanstack/ai` (its neutral home) so + * the durable lock this package provides via {@link withLocks} reaches + * `@tanstack/ai-sandbox`'s `ensure` through the SAME token reference. This + * module re-exports it unchanged so persistence consumers keep importing + * everything lock-related from `@tanstack/ai-persistence`. * * Locks are **not** part of {@link AIPersistenceStores}. State persistence - * (messages/runs/interrupts/metadata) and mutual exclusion are separate - * concerns: wire locks with {@link withLocks} (or pass a `LockStore` into - * consumers that need one), not via `composePersistence`. - * - * ponytail: the `'locks'` capability token is defined LOCALLY here. Capability - * identity is by object reference (see `createCapability`), not by name, so this - * token does NOT interoperate with the identically-named `LocksCapability` that - * `@tanstack/ai-sandbox` owns. That is fine for this batch: the only consumer of - * a provided lock store is `withSandbox`, and sandbox persistence is deferred. - * When it lands, share ONE handle between provider and consumer (the neutral - * home is core `@tanstack/ai`); until then this token has no cross-package - * consumer and just keeps the API self-contained. - */ -import { createCapability } from '@tanstack/ai' - -/** - * Mutual exclusion around a critical section keyed by `key`. A Cloudflare - * Durable Object store is the reference multi-instance backend; the in-memory - * default is correct within a single process only. Lease-backed implementations - * abort `signal` as soon as ownership can no longer be guaranteed; the callback - * must stop externally visible mutations when it aborts. - */ -export interface LockStore { - withLock: ( - key: string, - fn: (signal: AbortSignal) => Promise, - ) => Promise -} - -/** The lock capability. Provided by {@link withLocks}. */ -export const LocksCapability = createCapability()('locks') - -/** Destructured accessors: `getLocks(ctx)` / `provideLocks(ctx, store)`. */ -export const [getLocks, provideLocks] = LocksCapability - -/** - * In-memory {@link LockStore} — a per-key promise chain. Correct within a single - * process; multi-instance correctness needs a distributed lock (e.g. Cloudflare - * Durable Objects via `@tanstack/ai-persistence-cloudflare`). + * (messages/runs/interrupts/metadata/sandbox) and mutual exclusion are separate + * concerns: wire locks with {@link withLocks}, not via `composePersistence`. */ -export class InMemoryLockStore implements LockStore { - private readonly chains = new Map>() - - withLock( - key: string, - fn: (signal: AbortSignal) => Promise, - ): Promise { - const prior = this.chains.get(key) ?? Promise.resolve() - const runCriticalSection = () => fn(new AbortController().signal) - // Chain after the prior holder regardless of how it settled. - const run = prior.then(runCriticalSection, runCriticalSection) - // Swallow rejections so one failure doesn't poison the lock, then drop the - // chain entry once this tail is still the latest — otherwise long-lived - // processes accumulate settled promises for every distinct key forever. - const settled = run.then( - () => undefined, - () => undefined, - ) - this.chains.set(key, settled) - void settled.then(() => { - if (this.chains.get(key) === settled) { - this.chains.delete(key) - } - }) - return run - } -} +export { + LocksCapability, + getLocks, + provideLocks, + InMemoryLockStore, +} from '@tanstack/ai' +export type { LockStore } from '@tanstack/ai' diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index b7da747a0..8c91a2046 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -3,9 +3,11 @@ import { InterruptsCapability, LocksCapability, PersistenceCapability, + SandboxStoreCapability, provideInterrupts, provideLocks, providePersistence, + provideSandboxStore, } from './capabilities' import { validateChatPersistenceStores, @@ -268,13 +270,16 @@ function interruptPayload(interrupt: unknown): Record { interface PersistencePlan { wantsInterrupts: boolean + wantsSandbox: boolean runs: AIPersistence['stores']['runs'] } function resolvePersistencePlan(persistence: AIPersistence): PersistencePlan { + const stores = persistence.stores as AIPersistenceStores return { - wantsInterrupts: persistence.stores.interrupts !== undefined, - runs: persistence.stores.runs, + wantsInterrupts: stores.interrupts !== undefined, + wantsSandbox: stores.sandbox !== undefined, + runs: stores.runs, } } @@ -424,7 +429,7 @@ export function withPersistence( const snapshotStreaming = options.snapshotStreaming ?? false const snapshotIntervalMs = options.snapshotIntervalMs ?? 1000 const plan = resolvePersistencePlan(persistence) - const { wantsInterrupts, runs } = plan + const { wantsInterrupts, wantsSandbox, runs } = plan const messageStore = persistence.stores.messages if (!messageStore) { // validateChatPersistenceStores already throws; this narrows for TypeScript. @@ -434,6 +439,7 @@ export function withPersistence( const provides = [ PersistenceCapability, ...(wantsInterrupts ? [InterruptsCapability] : []), + ...(wantsSandbox ? [SandboxStoreCapability] : []), ] return defineChatMiddleware({ @@ -450,6 +456,10 @@ export function withPersistence( if (wantsInterrupts && persistence.stores.interrupts) { provideInterrupts(ctx, persistence.stores.interrupts) } + if (wantsSandbox) { + const sandbox = (persistence.stores as AIPersistenceStores).sandbox + if (sandbox) provideSandboxStore(ctx, sandbox) + } }, async onConfig(ctx: ChatMiddlewareContext, config: ChatMiddlewareConfig) { @@ -631,10 +641,15 @@ export function withPersistence( * Prisma, D1, memory) do not ship locks, and multi-instance apps compose a * distributed lock separately (e.g. Cloudflare Durable Objects). * + * The `LocksCapability` token is the same one `@tanstack/ai-sandbox` consumes, + * so a lock provided here reaches `withSandbox` automatically when both are in + * the chain. + * * ```ts * middleware: [ * withPersistence(drizzlePersistence(db, opts)), * withLocks(createDurableObjectLockStore(env.AI_LOCKS)), + * withSandbox(sandbox), * ] * ``` */ diff --git a/packages/ai-persistence/src/types.ts b/packages/ai-persistence/src/types.ts index 8923e5d00..9495e3ecf 100644 --- a/packages/ai-persistence/src/types.ts +++ b/packages/ai-persistence/src/types.ts @@ -1,4 +1,9 @@ -import type { ModelMessage, Scope, TokenUsage } from '@tanstack/ai' +import type { + ModelMessage, + Scope, + SandboxStore, + TokenUsage, +} from '@tanstack/ai' // Re-export the shared identity type so app code can import Scope from either // `@tanstack/ai` or `@tanstack/ai-persistence`. See {@link Scope} security notes: @@ -242,8 +247,9 @@ export interface MetadataStore { * * **Not a public product shape.** Prefer the named chat shapes below * ({@link ChatTranscriptStores}, {@link ChatPersistenceStores}, - * {@link ChatWithInterruptsStores}). Locks are not included (see - * {@link withLocks}). + * {@link ChatWithInterruptsStores}, {@link ChatAndSandboxPersistenceStores}). + * Locks are not included (see {@link withLocks}). Optional `sandbox` is for + * durable sandbox resume and is not part of the chat-only shapes. * * @internal Exported from this module for generics; the package root does not * re-export this type — use a named shape or `AIPersistence<{ … }>` instead. @@ -253,6 +259,12 @@ export interface AIPersistenceStores { runs?: RunStore interrupts?: InterruptStore metadata?: MetadataStore + /** + * Optional durable sandbox resume map. When present, {@link withPersistence} + * provides the shared `sandbox-store` capability so `withSandbox` can resume + * across processes. Stays out of the chat-only product shapes. + */ + sandbox?: SandboxStore } /** @@ -320,6 +332,19 @@ export type ChatPersistence = AIPersistence export type ChatWithInterruptsPersistence = AIPersistence +/** + * Full chat durability plus a durable sandbox resume store. Packaged backends + * that include the sandboxes table/model return this shape; chat-only apps use + * {@link ChatPersistence}. + */ +export interface ChatAndSandboxPersistenceStores extends ChatPersistenceStores { + sandbox: SandboxStore +} + +/** {@link AIPersistence} for {@link ChatAndSandboxPersistenceStores}. */ +export type ChatAndSandboxPersistence = + AIPersistence + type StoreKey = keyof AIPersistenceStores type ExactStoreKeys = Exclude extends never @@ -425,6 +450,7 @@ const storeKeys = [ 'runs', 'interrupts', 'metadata', + 'sandbox', ] satisfies Array const storeKeySet = new Set(storeKeys) diff --git a/packages/ai-persistence/tests/capabilities.test.ts b/packages/ai-persistence/tests/capabilities.test.ts index 464f0c0a6..ab2027f52 100644 --- a/packages/ai-persistence/tests/capabilities.test.ts +++ b/packages/ai-persistence/tests/capabilities.test.ts @@ -3,8 +3,10 @@ import { EventType, chat, defineChatMiddleware } from '@tanstack/ai' import type { AnyTextAdapter, ChatMiddlewareContext, + SandboxStore, StreamChunk, } from '@tanstack/ai' +import { InMemorySandboxStore } from '@tanstack/ai' import type { LockStore } from '../src/locks' import { InMemoryLockStore } from '../src/locks' import { memoryPersistence } from '../src/memory' @@ -13,9 +15,11 @@ import { InterruptsCapability, LocksCapability, PersistenceCapability, + SandboxStoreCapability, getInterrupts, getLocks, getPersistence, + getSandboxStore, } from '../src/capabilities' import { createInterruptController } from '../src/interrupts' import type { AIPersistence, InterruptStore } from '../src' @@ -124,6 +128,48 @@ describe('persistence capabilities', () => { expect(seen.locks).toBe(locks) }) + it('provides the sandbox store when present on the persistence bag', async () => { + const persistence = memoryPersistence() + const sandbox = new InMemorySandboxStore() + const withSandboxStore = { + stores: { ...persistence.stores, sandbox }, + } + const seen: { sandbox?: SandboxStore } = {} + + const consumer = defineChatMiddleware({ + name: 'sandbox-consumer', + requires: [SandboxStoreCapability], + setup(ctx: ChatMiddlewareContext) { + seen.sandbox = getSandboxStore(ctx) + }, + }) + + await collect( + chat({ + adapter: mockAdapter([ + { + type: EventType.RUN_STARTED, + runId: 'r1', + threadId: 't1', + timestamp: 1, + }, + { + type: EventType.RUN_FINISHED, + runId: 'r1', + threadId: 't1', + finishReason: 'stop', + timestamp: 1, + }, + ]), + messages: [{ role: 'user', content: 'hi' }], + runId: 'r1', + threadId: 't1', + middleware: [withPersistence(withSandboxStore), consumer], + }) as AsyncIterable, + ) + + expect(seen.sandbox).toBe(sandbox) + }) }) describe('createInterruptController', () => { diff --git a/packages/ai-sandbox/package.json b/packages/ai-sandbox/package.json index 924b25279..5b3c05ccf 100644 --- a/packages/ai-sandbox/package.json +++ b/packages/ai-sandbox/package.json @@ -32,6 +32,10 @@ "./ngrok": { "types": "./dist/esm/ngrok.d.ts", "import": "./dist/esm/ngrok.js" + }, + "./testkit": { + "types": "./dist/esm/testkit/conformance.d.ts", + "import": "./dist/esm/testkit/conformance.js" } }, "files": [ @@ -54,16 +58,21 @@ }, "peerDependencies": { "@ngrok/ngrok": "^1.0.0", - "@tanstack/ai": "workspace:^" + "@tanstack/ai": "workspace:^", + "vitest": "^4.1.10" }, "peerDependenciesMeta": { "@ngrok/ngrok": { "optional": true + }, + "vitest": { + "optional": true } }, "devDependencies": { "@ngrok/ngrok": "^1.7.0", "@tanstack/ai": "workspace:*", - "@vitest/coverage-v8": "4.0.14" + "@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 b1914c4d9..f21158ca7 100644 --- a/packages/ai-sandbox/skills/ai-sandbox/SKILL.md +++ b/packages/ai-sandbox/skills/ai-sandbox/SKILL.md @@ -211,6 +211,42 @@ const policy = defineSandboxPolicy({ provider + workspace hash + tenant so changing the repo/setup/image starts fresh. Ensure order: resume running → restore snapshot → create + bootstrap. +## Persistence (durable resume) + +Resume bookkeeping defaults to in-memory (single-process). For cross-process / +multi-instance resume, `withPersistence` (from `@tanstack/ai-persistence`) +provides a durable `SandboxStore` (which sandbox to resume) and an optional +distributed `LockStore` (serialize ensure); `withSandbox` consumes them. Compose +`withPersistence` **before** `withSandbox`. Same two middlewares whether the +persistence is for chat, sandboxes, or both. `ai-sandbox` never hardcodes storage. + +```typescript +import { chat } from '@tanstack/ai' +import { withSandbox } from '@tanstack/ai-sandbox' +import { withPersistence } from '@tanstack/ai-persistence' +import { sqlitePersistence } from '@tanstack/ai-persistence-drizzle/sqlite' + +// sqlitePersistence carries the sandbox store (+ chat stores) out of the box. +const persistence = sqlitePersistence({ url: './state.db', migrate: true }) + +chat({ + adapter, + messages, + middleware: [withPersistence(persistence), withSandbox(sandbox)], +}) +``` + +- Backends carry the sandbox store in their persistence object: + `sqlitePersistence(...)` / `drizzlePersistence(db)`, `prismaPersistence(prisma)`, + `cloudflarePersistence({ d1, durableObjects })` (edge). The `sandboxes` table / + model ships in the same migration bundle as the chat stores. +- The lock is the shared core `'locks'` token — `withPersistence` provides it + when its store set includes `locks` (e.g. the Cloudflare Durable Object lock), + and `withSandbox` picks it up automatically. +- `SandboxStore`/`LockStore` types live in core `@tanstack/ai`. The store + contract is conformance-tested via `runSandboxStoreConformance` from + `@tanstack/ai-sandbox/testkit`. + ## File-event hooks Watch the workspace for create/change/delete events. Provider-agnostic: native diff --git a/packages/ai-sandbox/src/capabilities.ts b/packages/ai-sandbox/src/capabilities.ts index 8abd7f95d..9a2c470e4 100644 --- a/packages/ai-sandbox/src/capabilities.ts +++ b/packages/ai-sandbox/src/capabilities.ts @@ -5,21 +5,28 @@ * - `SandboxCapability` is PROVIDED by `withSandbox` and REQUIRED by harness * adapters (`requires: [SandboxCapability]`). * - `SandboxStoreCapability` / `LocksCapability` are OPTIONALLY required by - * `withSandbox`. v1 falls back to in-memory defaults; the future persistence - * package PROVIDES durable implementations. + * `withSandbox`. They fall back to in-memory defaults; `withPersistence` + * (from `@tanstack/ai-persistence`) PROVIDES durable implementations. */ -import { createCapability } from '@tanstack/ai' +import { + LocksCapability, + SandboxStoreCapability, + createCapability, +} from '@tanstack/ai' import type { SandboxHandle } from './contracts' -import type { LockStore, SandboxStore } from './store' import type { SandboxPolicy } from './policy' import type { ToolBridgeProvisioner } from './tool-bridge' export const SandboxCapability = createCapability()('sandbox') -export const SandboxStoreCapability = - createCapability()('sandbox-store') - -export const LocksCapability = createCapability()('locks') +/** + * The `'sandbox-store'` and `'locks'` tokens come from core `@tanstack/ai` — the + * SAME references `withPersistence` provides — so a durable store and + * distributed lock reach `ensure` without a sandbox↔persistence dependency. + * Re-exported here so sandbox consumers import them alongside the other + * sandbox capabilities. + */ +export { LocksCapability, SandboxStoreCapability } /** * The active sandbox policy, provided by `withSandbox` from the definition. diff --git a/packages/ai-sandbox/src/store.ts b/packages/ai-sandbox/src/store.ts index 7371cfbd6..547ba98d2 100644 --- a/packages/ai-sandbox/src/store.ts +++ b/packages/ai-sandbox/src/store.ts @@ -1,83 +1,11 @@ /** - * Persistence seams for the sandbox layer. + * Persistence seams for the sandbox layer, re-exported from core `@tanstack/ai`. * - * v1 ships ONLY in-memory implementations (single-process resume). These are - * deliberately pluggable OPTIONAL capabilities so the future persistence - * package can `provide` durable implementations (D1/Postgres/Durable Objects) - * without the sandbox layer changing. Do NOT hardcode storage here. + * `SandboxStore` and `LockStore` are the SAME tokens `@tanstack/ai-persistence` + * provides through `withPersistence`, so a durable store and distributed lock + * supplied by a persistence middleware reach `withSandbox` transparently. They + * live in core (their neutral home) so neither package depends on the other; the + * sandbox layer never hardcodes storage. */ - -/** One persisted sandbox instance, keyed by the compound sandbox instance key. */ -export interface SandboxRecord { - /** Compound key (see computeSandboxKey). */ - key: string - /** Provider name that owns `providerSandboxId`. */ - provider: string - /** Provider-assigned sandbox id used to resume. */ - providerSandboxId: string - /** Most recent snapshot id, when the provider supports snapshots. */ - latestSnapshotId?: string - threadId: string - latestRunId?: string - /** Epoch ms of last write (for keepAlive / GC by the persistence layer). */ - updatedAt: number -} - -/** Maps a compound key to the provider sandbox that should be resumed. */ -export interface SandboxStore { - get: (key: string) => Promise - upsert: (record: SandboxRecord) => Promise - delete: (key: string) => Promise -} - -/** - * Mutual exclusion around sandbox ensure so two concurrent runs for the same - * thread don't both create a sandbox. The in-memory default is single-process; - * the persistence layer provides a distributed lock (e.g. a Durable Object). - */ -export interface LockStore { - withLock: (key: string, fn: () => Promise) => Promise -} - -/** In-memory {@link SandboxStore}. Resume works only within one process. */ -export class InMemorySandboxStore implements SandboxStore { - private readonly map = new Map() - - get(key: string): Promise { - return Promise.resolve(this.map.get(key) ?? null) - } - - upsert(record: SandboxRecord): Promise { - this.map.set(record.key, record) - return Promise.resolve() - } - - delete(key: string): Promise { - this.map.delete(key) - return Promise.resolve() - } -} - -/** - * In-memory {@link LockStore} — a per-key promise chain. Correct within a - * single process; multi-instance correctness needs a distributed lock from the - * persistence layer. - */ -export class InMemoryLockStore implements LockStore { - private readonly chains = new Map>() - - withLock(key: string, fn: () => Promise): Promise { - const prior = this.chains.get(key) ?? Promise.resolve() - // Chain after the prior holder regardless of how it settled. - const run = prior.then(fn, fn) - // Keep the chain alive but swallow rejections so one failure doesn't poison the lock. - this.chains.set( - key, - run.then( - () => undefined, - () => undefined, - ), - ) - return run - } -} +export { InMemoryLockStore, InMemorySandboxStore } from '@tanstack/ai' +export type { LockStore, SandboxStore, SandboxRecord } from '@tanstack/ai' diff --git a/packages/ai-sandbox/src/testkit/conformance.ts b/packages/ai-sandbox/src/testkit/conformance.ts new file mode 100644 index 000000000..6ddcf7cb0 --- /dev/null +++ b/packages/ai-sandbox/src/testkit/conformance.ts @@ -0,0 +1,98 @@ +/** + * Conformance suite for a {@link SandboxStore} implementation. + * + * A durable backend (`@tanstack/ai-persistence-drizzle` / `-prisma` / + * `-cloudflare`) runs this against a fresh store to prove it satisfies the + * get / upsert / delete contract `@tanstack/ai-sandbox`'s ensure algorithm + * relies on — including insert-vs-overwrite and optional-field handling. The + * in-memory reference store runs it too. + * + * Vitest is an OPTIONAL peer dependency: this module is imported only from test + * files, which already run under Vitest. + */ +import { describe, expect, it } from 'vitest' +import type { SandboxRecord, SandboxStore } from '../store' + +function makeRecord(overrides?: Partial): SandboxRecord { + return { + key: 'thread-1', + provider: 'fake', + providerSandboxId: 'sb-1', + threadId: 'thread-1', + updatedAt: 1, + ...overrides, + } +} + +/** + * Assert `makeStore()` produces a spec-compliant {@link SandboxStore}. Each + * `it` gets a fresh store, so implementations may share process state across + * calls without cross-test bleed only if `makeStore` returns an isolated store. + */ +export function runSandboxStoreConformance( + name: string, + makeStore: () => SandboxStore | Promise, +): void { + describe(`SandboxStore conformance: ${name}`, () => { + it('returns null for a missing key', async () => { + const store = await makeStore() + expect(await store.get('absent')).toBeNull() + }) + + it('round-trips an upserted record with all fields', async () => { + const store = await makeStore() + const record = makeRecord({ + latestSnapshotId: 'snap-1', + latestRunId: 'run-1', + }) + await store.upsert(record) + expect(await store.get(record.key)).toEqual(record) + }) + + it('omits absent optional fields on read', async () => { + const store = await makeStore() + const record = makeRecord() + await store.upsert(record) + const loaded = await store.get(record.key) + expect(loaded).toEqual(record) + expect(loaded && 'latestSnapshotId' in loaded).toBe(false) + expect(loaded && 'latestRunId' in loaded).toBe(false) + }) + + it('overwrites an existing record on re-upsert', async () => { + const store = await makeStore() + await store.upsert(makeRecord({ latestSnapshotId: 'snap-1' })) + await store.upsert( + makeRecord({ providerSandboxId: 'sb-2', updatedAt: 2 }), + ) + const loaded = await store.get('thread-1') + expect(loaded?.providerSandboxId).toBe('sb-2') + expect(loaded?.updatedAt).toBe(2) + // The overwrite dropped latestSnapshotId — a durable store must clear it, + // not retain the prior value. + expect(loaded && 'latestSnapshotId' in loaded).toBe(false) + }) + + it('isolates records by key', async () => { + const store = await makeStore() + await store.upsert(makeRecord({ key: 'a', threadId: 'a' })) + await store.upsert( + makeRecord({ key: 'b', threadId: 'b', providerSandboxId: 'sb-b' }), + ) + expect((await store.get('a'))?.providerSandboxId).toBe('sb-1') + expect((await store.get('b'))?.providerSandboxId).toBe('sb-b') + }) + + it('deletes a record', async () => { + const store = await makeStore() + await store.upsert(makeRecord()) + await store.delete('thread-1') + expect(await store.get('thread-1')).toBeNull() + }) + + it('delete of a missing key is a no-op', async () => { + const store = await makeStore() + await expect(store.delete('absent')).resolves.toBeUndefined() + }) + }) +} diff --git a/packages/ai-sandbox/tests/resume-from-store.test.ts b/packages/ai-sandbox/tests/resume-from-store.test.ts new file mode 100644 index 000000000..087b6f467 --- /dev/null +++ b/packages/ai-sandbox/tests/resume-from-store.test.ts @@ -0,0 +1,93 @@ +import { describe, expect, it } from 'vitest' +import { EventType, chat, defineChatMiddleware } from '@tanstack/ai' +import { withSandbox } from '../src/middleware' +import { defineSandbox } from '../src/sandbox' +import { defineWorkspace } from '../src/workspace' +import { + LocksCapability, + SandboxStoreCapability, + provideLocks, + provideSandboxStore, +} from '../src/capabilities' +import { InMemoryLockStore, InMemorySandboxStore } from '../src/store' +import { makeFakeProvider } from './fakes' +import type { + AnyTextAdapter, + LockStore, + SandboxStore, + StreamChunk, +} from '@tanstack/ai' + +function mockAdapter(): AnyTextAdapter { + return { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': {}, + chatStream: ({ runId, threadId }: { runId: string; threadId: string }) => + (async function* (): AsyncGenerator { + yield { type: EventType.RUN_STARTED, runId, threadId, timestamp: 1 } + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + finishReason: 'stop', + timestamp: 1, + } + })(), + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + } as unknown as AnyTextAdapter +} + +/** + * Stand-in for the durable store `withPersistence` provides: a middleware that + * puts a `SandboxStore` (and lock) on the shared capabilities so `withSandbox` + * consumes them in `ensure`. + */ +function provideStores(store: SandboxStore, locks: LockStore) { + return defineChatMiddleware({ + name: 'test-provide-sandbox-stores', + provides: [SandboxStoreCapability, LocksCapability], + setup(ctx) { + provideSandboxStore(ctx, store) + provideLocks(ctx, locks) + }, + }) +} + +async function drain(stream: AsyncIterable): Promise { + for await (const _ of stream) void _ +} + +describe('withSandbox resume from a provided store', () => { + it('resumes the persisted sandbox across runs instead of re-creating it', async () => { + const provider = makeFakeProvider() + const store = new InMemorySandboxStore() + const locks = new InMemoryLockStore() + const sandbox = defineSandbox({ + id: 'repo', + provider, + workspace: defineWorkspace({ source: { type: 'none' } }), + fileEvents: false, + }) + + const run = (runId: string) => + drain( + chat({ + adapter: mockAdapter(), + messages: [{ role: 'user', content: 'hi' }], + runId, + threadId: 'thread-1', + middleware: [provideStores(store, locks), withSandbox(sandbox)], + }), + ) + + await run('run-1') + await run('run-2') + + // The record persisted in the provided store across runs, so the second + // ensure resumed rather than creating a second sandbox. + expect(provider.calls.create).toBe(1) + expect(provider.calls.resume).toBe(1) + }) +}) diff --git a/packages/ai-sandbox/tests/store.conformance.test.ts b/packages/ai-sandbox/tests/store.conformance.test.ts new file mode 100644 index 000000000..e37f6d2fb --- /dev/null +++ b/packages/ai-sandbox/tests/store.conformance.test.ts @@ -0,0 +1,6 @@ +import { runSandboxStoreConformance } from '../src/testkit/conformance' +import { InMemorySandboxStore } from '../src/store' + +// The reference in-memory store must satisfy the same contract the durable +// backends are held to. +runSandboxStoreConformance('in-memory', () => new InMemorySandboxStore()) diff --git a/packages/ai-sandbox/vite.config.ts b/packages/ai-sandbox/vite.config.ts index 27306737e..d6b5b0343 100644 --- a/packages/ai-sandbox/vite.config.ts +++ b/packages/ai-sandbox/vite.config.ts @@ -32,9 +32,13 @@ export default mergeConfig( tanstackViteConfig({ // Core entry + the optional ngrok tool-bridge provisioner subpath // (`@tanstack/ai-sandbox/ngrok`), which lazy-loads the optional `@ngrok/ngrok` - // peer dep so the core never pulls in its native binary. - entry: ['./src/index.ts', './src/ngrok.ts'], + // peer dep so the core never pulls in its native binary + the SandboxStore + // conformance testkit (`@tanstack/ai-sandbox/testkit`). + entry: ['./src/index.ts', './src/ngrok.ts', './src/testkit/conformance.ts'], srcDir: './src', + // The conformance testkit imports Vitest; keep it external so the built + // artifact references the consumer's Vitest at test time. + externalDeps: ['vitest'], cjs: false, }), ) diff --git a/packages/ai/src/activities/chat/middleware/index.ts b/packages/ai/src/activities/chat/middleware/index.ts index ad4fa9dad..668531395 100644 --- a/packages/ai/src/activities/chat/middleware/index.ts +++ b/packages/ai/src/activities/chat/middleware/index.ts @@ -41,3 +41,19 @@ export type { } from './builder' export { validateCapabilities } from './validate' export type { AnyChatMiddleware } from './types' + +export { + LocksCapability, + getLocks, + provideLocks, + InMemoryLockStore, +} from './locks' +export type { LockStore } from './locks' + +export { + SandboxStoreCapability, + getSandboxStore, + provideSandboxStore, + InMemorySandboxStore, +} from './sandbox-store' +export type { SandboxStore, SandboxRecord } from './sandbox-store' diff --git a/packages/ai/src/activities/chat/middleware/locks.ts b/packages/ai/src/activities/chat/middleware/locks.ts new file mode 100644 index 000000000..5cc1c6741 --- /dev/null +++ b/packages/ai/src/activities/chat/middleware/locks.ts @@ -0,0 +1,71 @@ +/** + * Distributed-mutex primitive — the neutral, shared home for the `'locks'` + * capability. + * + * Capability identity is by object reference (see `createCapability`), so this + * ONE `LocksCapability` is the single token both `@tanstack/ai-persistence` + * (which PROVIDES a durable lock via {@link withLocks}) and + * `@tanstack/ai-sandbox` (whose `ensure` CONSUMES it to serialize resume-or-create) + * import and re-export. Defining it here — core, depended on by both and + * depending on neither — is what lets a persistence-provided distributed lock + * reach `withSandbox` without either package depending on the other. + */ +import { createCapability } from './capabilities' + +/** + * Mutual exclusion around a critical section keyed by `key`. A distributed + * backend (e.g. a Cloudflare Durable Object) is the only kind safe across + * instances; the in-memory default is correct within a single process only. + * Lease-backed implementations abort `signal` as soon as ownership can no longer + * be guaranteed; the callback must stop externally visible mutations when it + * aborts. Callbacks that ignore `signal` (e.g. the sandbox `ensure` critical + * section) remain valid — a `() => Promise` is assignable to the + * signal-taking parameter. + */ +export interface LockStore { + withLock: ( + key: string, + fn: (signal: AbortSignal) => Promise, + ) => Promise +} + +/** + * The lock capability. Provided by `withLocks` (from `@tanstack/ai-persistence`) + * or any middleware that calls {@link provideLocks}. + */ +export const LocksCapability = createCapability()('locks') + +/** Destructured accessors: `getLocks(ctx)` / `provideLocks(ctx, store)`. */ +export const [getLocks, provideLocks] = LocksCapability + +/** + * In-memory {@link LockStore} — a per-key promise chain. Correct within a single + * process; multi-instance correctness needs a distributed lock from the backend. + */ +export class InMemoryLockStore implements LockStore { + private readonly chains = new Map>() + + withLock( + key: string, + fn: (signal: AbortSignal) => Promise, + ): Promise { + const prior = this.chains.get(key) ?? Promise.resolve() + const runCriticalSection = () => fn(new AbortController().signal) + // Chain after the prior holder regardless of how it settled. + const run = prior.then(runCriticalSection, runCriticalSection) + // Swallow rejections so one failure doesn't poison the lock, then drop the + // chain entry once this tail is still the latest — otherwise long-lived + // processes accumulate settled promises for every distinct key forever. + const settled = run.then( + () => undefined, + () => undefined, + ) + this.chains.set(key, settled) + void settled.then(() => { + if (this.chains.get(key) === settled) { + this.chains.delete(key) + } + }) + return run + } +} diff --git a/packages/ai/src/activities/chat/middleware/sandbox-store.ts b/packages/ai/src/activities/chat/middleware/sandbox-store.ts new file mode 100644 index 000000000..3a0431929 --- /dev/null +++ b/packages/ai/src/activities/chat/middleware/sandbox-store.ts @@ -0,0 +1,66 @@ +/** + * Durable sandbox-resume seam — the neutral, shared home for the + * `'sandbox-store'` capability. + * + * Capability identity is by object reference (see `createCapability`), so this + * ONE `SandboxStoreCapability` is the single token both `@tanstack/ai-persistence` + * (whose `withPersistence` PROVIDES a durable store when `stores.sandbox` is + * present) and `@tanstack/ai-sandbox` (whose `withSandbox` CONSUMES it in + * `ensure`) import and re-export. Defining it here — core, depended on by both + * and depending on neither — lets a persistence-provided store reach the sandbox + * layer without either package depending on the other (the same arrangement as + * the `'locks'` token). + */ +import { createCapability } from './capabilities' + +/** One persisted sandbox instance, keyed by the compound sandbox instance key. */ +export interface SandboxRecord { + /** Compound key (see `computeSandboxKey` in `@tanstack/ai-sandbox`). */ + key: string + /** Provider name that owns `providerSandboxId`. */ + provider: string + /** Provider-assigned sandbox id used to resume. */ + providerSandboxId: string + /** Most recent snapshot id, when the provider supports snapshots. */ + latestSnapshotId?: string + threadId: string + latestRunId?: string + /** Epoch ms of last write (for keepAlive / GC by the persistence layer). */ + updatedAt: number +} + +/** Maps a compound key to the provider sandbox that should be resumed. */ +export interface SandboxStore { + get: (key: string) => Promise + upsert: (record: SandboxRecord) => Promise + delete: (key: string) => Promise +} + +/** + * The sandbox-store capability. Provided by `withPersistence` when a `sandbox` + * store is present; consumed by `withSandbox`. + */ +export const SandboxStoreCapability = + createCapability()('sandbox-store') + +/** Destructured accessors: `getSandboxStore(ctx)` / `provideSandboxStore(ctx, store)`. */ +export const [getSandboxStore, provideSandboxStore] = SandboxStoreCapability + +/** In-memory {@link SandboxStore}. Resume works only within one process. */ +export class InMemorySandboxStore implements SandboxStore { + private readonly map = new Map() + + get(key: string): Promise { + return Promise.resolve(this.map.get(key) ?? null) + } + + upsert(record: SandboxRecord): Promise { + this.map.set(record.key, record) + return Promise.resolve() + } + + delete(key: string): Promise { + this.map.delete(key) + return Promise.resolve() + } +} diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 4448c4e44..77b103571 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -218,6 +218,27 @@ export type { DefinedChatMiddleware, AnyChatMiddleware, } from './activities/chat/middleware/index' +// Shared distributed-mutex primitive — the neutral home for the `'locks'` +// capability, consumed by both @tanstack/ai-persistence and @tanstack/ai-sandbox. +export { + LocksCapability, + getLocks, + provideLocks, + InMemoryLockStore, +} from './activities/chat/middleware/index' +export type { LockStore } from './activities/chat/middleware/index' +// Durable sandbox-resume store — neutral home for the `'sandbox-store'` +// capability, provided by @tanstack/ai-persistence and consumed by @tanstack/ai-sandbox. +export { + SandboxStoreCapability, + getSandboxStore, + provideSandboxStore, + InMemorySandboxStore, +} from './activities/chat/middleware/index' +export type { + SandboxStore, + SandboxRecord, +} from './activities/chat/middleware/index' // Well-known AG-UI CUSTOM event catalog (agent activity rides on CUSTOM events) export { CUSTOM_EVENT, isCustomEvent } from './custom-events' diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index baa27b4e2..5fe51a643 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -2243,6 +2243,9 @@ importers: '@tanstack/ai-persistence-drizzle': specifier: workspace:* version: link:../ai-persistence-drizzle + '@tanstack/ai-sandbox': + specifier: workspace:* + version: link:../ai-sandbox '@vitest/coverage-v8': specifier: 4.0.14 version: 4.0.14(supports-color@7.2.0)(vitest@4.1.10) @@ -2264,6 +2267,9 @@ importers: '@tanstack/ai-persistence': specifier: workspace:* version: link:../ai-persistence + '@tanstack/ai-sandbox': + specifier: workspace:* + version: link:../ai-sandbox '@vitest/coverage-v8': specifier: 4.0.14 version: 4.0.14(supports-color@7.2.0)(vitest@4.1.10) @@ -2282,6 +2288,9 @@ importers: '@tanstack/ai-persistence': specifier: workspace:* version: link:../ai-persistence + '@tanstack/ai-sandbox': + specifier: workspace:* + version: link:../ai-sandbox '@vitest/coverage-v8': specifier: 4.0.14 version: 4.0.14(supports-color@7.2.0)(vitest@4.1.10) @@ -2406,6 +2415,9 @@ importers: '@vitest/coverage-v8': specifier: 4.0.14 version: 4.0.14(supports-color@7.2.0)(vitest@4.1.10) + vitest: + specifier: ^4.1.10 + version: 4.1.10(@opentelemetry/api@1.9.1)(@types/node@24.10.3)(@vitest/coverage-v8@4.0.14)(happy-dom@20.0.11)(jsdom@27.3.0(postcss@8.5.19))(vite@8.1.4(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.9.0)) packages/ai-sandbox-cloudflare: dependencies: @@ -2845,12 +2857,18 @@ importers: '@tanstack/ai-openrouter': specifier: workspace:* version: link:../../packages/ai-openrouter + '@tanstack/ai-persistence': + specifier: workspace:* + version: link:../../packages/ai-persistence '@tanstack/ai-react': specifier: workspace:* version: link:../../packages/ai-react '@tanstack/ai-react-ui': specifier: workspace:* version: link:../../packages/ai-react-ui + '@tanstack/ai-sandbox': + specifier: workspace:* + version: link:../../packages/ai-sandbox '@tanstack/devtools-event-bus': specifier: ^0.4.1 version: 0.4.1 diff --git a/testing/e2e/package.json b/testing/e2e/package.json index b84b22a92..378b550fd 100644 --- a/testing/e2e/package.json +++ b/testing/e2e/package.json @@ -32,8 +32,10 @@ "@tanstack/ai-ollama": "workspace:*", "@tanstack/ai-openai": "workspace:*", "@tanstack/ai-openrouter": "workspace:*", + "@tanstack/ai-persistence": "workspace:*", "@tanstack/ai-react": "workspace:*", "@tanstack/ai-react-ui": "workspace:*", + "@tanstack/ai-sandbox": "workspace:*", "@tanstack/devtools-event-bus": "^0.4.1", "@tanstack/nitro-v2-vite-plugin": "^1.155.0", "@tanstack/react-ai-devtools": "workspace:*", diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index 4abac986d..54dc78951 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -31,6 +31,7 @@ import { Route as ApiTranscriptionRouteImport } from './routes/api.transcription 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 ApiSandboxDurabilityRouteImport } from './routes/api.sandbox-durability' import { Route as ApiPersistenceDurabilityRouteImport } from './routes/api.persistence-durability' import { Route as ApiOtelUsageRouteImport } from './routes/api.otel-usage' import { Route as ApiOtelMediaRouteImport } from './routes/api.otel-media' @@ -179,6 +180,11 @@ const ApiSummarizeRoute = ApiSummarizeRouteImport.update({ path: '/api/summarize', getParentRoute: () => rootRouteImport, } as any) +const ApiSandboxDurabilityRoute = ApiSandboxDurabilityRouteImport.update({ + id: '/api/sandbox-durability', + path: '/api/sandbox-durability', + getParentRoute: () => rootRouteImport, +} as any) const ApiPersistenceDurabilityRoute = ApiPersistenceDurabilityRouteImport.update({ id: '/api/persistence-durability', @@ -412,6 +418,7 @@ export interface FileRoutesByFullPath { '/api/otel-media': typeof ApiOtelMediaRoute '/api/otel-usage': typeof ApiOtelUsageRoute '/api/persistence-durability': typeof ApiPersistenceDurabilityRoute + '/api/sandbox-durability': typeof ApiSandboxDurabilityRoute '/api/summarize': typeof ApiSummarizeRoute '/api/tool-call-lifecycle-wire': typeof ApiToolCallLifecycleWireRoute '/api/tools-test': typeof ApiToolsTestRoute @@ -472,6 +479,7 @@ export interface FileRoutesByTo { '/api/otel-media': typeof ApiOtelMediaRoute '/api/otel-usage': typeof ApiOtelUsageRoute '/api/persistence-durability': typeof ApiPersistenceDurabilityRoute + '/api/sandbox-durability': typeof ApiSandboxDurabilityRoute '/api/summarize': typeof ApiSummarizeRoute '/api/tool-call-lifecycle-wire': typeof ApiToolCallLifecycleWireRoute '/api/tools-test': typeof ApiToolsTestRoute @@ -533,6 +541,7 @@ export interface FileRoutesById { '/api/otel-media': typeof ApiOtelMediaRoute '/api/otel-usage': typeof ApiOtelUsageRoute '/api/persistence-durability': typeof ApiPersistenceDurabilityRoute + '/api/sandbox-durability': typeof ApiSandboxDurabilityRoute '/api/summarize': typeof ApiSummarizeRoute '/api/tool-call-lifecycle-wire': typeof ApiToolCallLifecycleWireRoute '/api/tools-test': typeof ApiToolsTestRoute @@ -595,6 +604,7 @@ export interface FileRouteTypes { | '/api/otel-media' | '/api/otel-usage' | '/api/persistence-durability' + | '/api/sandbox-durability' | '/api/summarize' | '/api/tool-call-lifecycle-wire' | '/api/tools-test' @@ -655,6 +665,7 @@ export interface FileRouteTypes { | '/api/otel-media' | '/api/otel-usage' | '/api/persistence-durability' + | '/api/sandbox-durability' | '/api/summarize' | '/api/tool-call-lifecycle-wire' | '/api/tools-test' @@ -715,6 +726,7 @@ export interface FileRouteTypes { | '/api/otel-media' | '/api/otel-usage' | '/api/persistence-durability' + | '/api/sandbox-durability' | '/api/summarize' | '/api/tool-call-lifecycle-wire' | '/api/tools-test' @@ -776,6 +788,7 @@ export interface RootRouteChildren { ApiOtelMediaRoute: typeof ApiOtelMediaRoute ApiOtelUsageRoute: typeof ApiOtelUsageRoute ApiPersistenceDurabilityRoute: typeof ApiPersistenceDurabilityRoute + ApiSandboxDurabilityRoute: typeof ApiSandboxDurabilityRoute ApiSummarizeRoute: typeof ApiSummarizeRoute ApiToolCallLifecycleWireRoute: typeof ApiToolCallLifecycleWireRoute ApiToolsTestRoute: typeof ApiToolsTestRoute @@ -941,6 +954,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiSummarizeRouteImport parentRoute: typeof rootRouteImport } + '/api/sandbox-durability': { + id: '/api/sandbox-durability' + path: '/api/sandbox-durability' + fullPath: '/api/sandbox-durability' + preLoaderRoute: typeof ApiSandboxDurabilityRouteImport + parentRoute: typeof rootRouteImport + } '/api/persistence-durability': { id: '/api/persistence-durability' path: '/api/persistence-durability' @@ -1301,6 +1321,7 @@ const rootRouteChildren: RootRouteChildren = { ApiOtelMediaRoute: ApiOtelMediaRoute, ApiOtelUsageRoute: ApiOtelUsageRoute, ApiPersistenceDurabilityRoute: ApiPersistenceDurabilityRoute, + ApiSandboxDurabilityRoute: ApiSandboxDurabilityRoute, ApiSummarizeRoute: ApiSummarizeRoute, ApiToolCallLifecycleWireRoute: ApiToolCallLifecycleWireRoute, ApiToolsTestRoute: ApiToolsTestRoute, diff --git a/testing/e2e/src/routes/api.sandbox-durability.ts b/testing/e2e/src/routes/api.sandbox-durability.ts new file mode 100644 index 000000000..b6c729f2f --- /dev/null +++ b/testing/e2e/src/routes/api.sandbox-durability.ts @@ -0,0 +1,176 @@ +import { createFileRoute } from '@tanstack/react-router' +import { + EventType, + InMemoryLockStore, + InMemorySandboxStore, + chat, +} from '@tanstack/ai' +import { + defineSandbox, + defineWorkspace, + withSandbox, +} from '@tanstack/ai-sandbox' +import { + memoryPersistence, + withLocks, + withPersistence, +} from '@tanstack/ai-persistence' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' +import type { SandboxHandle, SandboxProvider } from '@tanstack/ai-sandbox' + +/** + * Server-side sandbox-persistence durability harness. + * + * Proves that a durable `SandboxStore` (here an in-memory singleton standing in + * for a backend) makes sandbox resume durable ACROSS independent runs: each + * POST is a fresh `chat()` with a brand-new middleware context, and the ONLY + * shared state is the module-singleton store. A second run for the same + * `threadId` must RESUME the sandbox the first run created, not create a + * second one. + * + * Provider-free: a fixed AG-UI stream and a fake sandbox provider, so there is + * no LLM in the loop (exempt from the aimock policy). + */ + +// Module-singleton state — survives between the two HTTP requests. +const sandboxStore = new InMemorySandboxStore() +const locks = new InMemoryLockStore() +const chatStores = memoryPersistence().stores +const persistence = { + stores: { + ...chatStores, + sandbox: sandboxStore, + }, +} +const calls = { create: 0, resume: 0 } + +function fakeHandle(id: string): SandboxHandle { + return { + id, + provider: 'fake', + capabilities: { + fs: true, + exec: true, + env: true, + ports: false, + backgroundProcesses: false, + writableStdin: false, + snapshots: false, + networkPolicy: false, + durableFilesystem: false, + fork: false, + }, + fs: { + read: () => Promise.resolve(''), + readBytes: () => Promise.resolve(new Uint8Array()), + write: () => Promise.resolve(), + list: () => Promise.resolve([]), + mkdir: () => Promise.resolve(), + remove: () => Promise.resolve(), + rename: () => Promise.resolve(), + exists: () => Promise.resolve(false), + }, + git: { + clone: () => Promise.resolve(), + status: () => Promise.resolve(''), + add: () => Promise.resolve(), + commit: () => Promise.resolve(), + push: () => Promise.resolve(), + pull: () => Promise.resolve(), + branch: () => Promise.resolve('main'), + }, + process: { + exec: () => Promise.resolve({ stdout: '', stderr: '', exitCode: 0 }), + spawn: () => Promise.reject(new Error('not supported')), + }, + ports: { connect: () => Promise.reject(new Error('not supported')) }, + env: { set: () => Promise.resolve() }, + destroy: () => Promise.resolve(), + } +} + +const provider: SandboxProvider = { + name: 'fake', + capabilities: () => fakeHandle('probe').capabilities, + create: (input) => { + calls.create++ + return Promise.resolve(fakeHandle(input.id ?? 'fake-sandbox')) + }, + resume: (input) => { + calls.resume++ + return Promise.resolve(fakeHandle(input.id)) + }, + destroy: () => Promise.resolve(), +} + +const sandbox = defineSandbox({ + id: 'durable', + provider, + workspace: defineWorkspace({ source: { type: 'none' } }), + fileEvents: false, +}) + +function fixedRun(threadId: string, runId: string): AsyncIterable { + return (async function* () { + yield { type: EventType.RUN_STARTED, threadId, runId, timestamp: 1 } + yield { + type: EventType.RUN_FINISHED, + threadId, + runId, + finishReason: 'stop', + timestamp: 1, + } + })() +} + +const adapter: AnyTextAdapter = { + kind: 'text', + name: 'fixed', + model: 'test-model', + '~types': {}, + chatStream: ({ runId, threadId }: { runId: string; threadId: string }) => + fixedRun(threadId, runId), + structuredOutput: () => Promise.resolve({ data: {}, rawText: '{}' }), +} as unknown as AnyTextAdapter + +function stringField(body: unknown, key: string): string | undefined { + if (typeof body !== 'object' || body === null || !(key in body)) { + return undefined + } + const value: unknown = (body as Record)[key] + return typeof value === 'string' ? value : undefined +} + +export const Route = createFileRoute('/api/sandbox-durability')({ + server: { + handlers: { + POST: async ({ request }) => { + const body: unknown = await request.json() + const threadId = stringField(body, 'threadId') ?? 'sandbox-thread' + const runId = stringField(body, 'runId') ?? crypto.randomUUID() + + // Drain a full run so setup (ensure) and the terminal hooks all execute. + const stream = chat({ + adapter, + messages: [{ role: 'user', content: 'go' }], + runId, + threadId, + middleware: [ + withPersistence(persistence), + withLocks(locks), + withSandbox(sandbox), + ], + }) + for await (const _ of stream) void _ + + const record = await sandboxStore.get(sandbox.key({ threadId, runId })) + return Response.json({ + create: calls.create, + resume: calls.resume, + providerSandboxId: record?.providerSandboxId ?? null, + latestRunId: record?.latestRunId ?? null, + }) + }, + }, + }, +}) diff --git a/testing/e2e/tests/sandbox-durability.spec.ts b/testing/e2e/tests/sandbox-durability.spec.ts new file mode 100644 index 000000000..d5fd23221 --- /dev/null +++ b/testing/e2e/tests/sandbox-durability.spec.ts @@ -0,0 +1,50 @@ +import { expect, test } from '@playwright/test' + +/** + * Server-side sandbox-persistence durability. + * + * Two independent runs (separate HTTP requests → fresh middleware contexts) for + * the same thread. The only shared state is the harness route's module-singleton + * `SandboxStore` provided via `withSandboxPersistence`. The second run must + * RESUME the sandbox the first created, proving durable resume survives across + * runs rather than re-creating on every request. + * + * Provider-free (fixed AG-UI stream + fake sandbox provider), so it is exempt + * from the aimock policy. + */ +interface DurabilityResult { + create: number + resume: number + providerSandboxId: string | null + latestRunId: string | null +} + +test.describe('sandbox persistence durability (server-side resume)', () => { + test('a second run resumes the persisted sandbox instead of creating a new one', async ({ + request, + }) => { + const threadId = `sandbox-durability-${Date.now()}` + + const first = await request.post('/api/sandbox-durability', { + data: { threadId, runId: 'run-1' }, + }) + expect(first.ok()).toBe(true) + const firstBody = (await first.json()) as DurabilityResult + expect(firstBody.create).toBe(1) + expect(firstBody.resume).toBe(0) + expect(firstBody.providerSandboxId).not.toBeNull() + + const second = await request.post('/api/sandbox-durability', { + data: { threadId, runId: 'run-2' }, + }) + expect(second.ok()).toBe(true) + const secondBody = (await second.json()) as DurabilityResult + + // No second create: the run resumed the sandbox persisted by run-1. + expect(secondBody.create).toBe(1) + expect(secondBody.resume).toBe(1) + expect(secondBody.providerSandboxId).toBe(firstBody.providerSandboxId) + // The store advanced the record's latest run to the resuming run. + expect(secondBody.latestRunId).toBe('run-2') + }) +})