diff --git a/.changeset/generation-persistence.md b/.changeset/generation-persistence.md new file mode 100644 index 000000000..774c0fdbc --- /dev/null +++ b/.changeset/generation-persistence.md @@ -0,0 +1,24 @@ +--- +'@tanstack/ai': minor +'@tanstack/ai-utils': minor +'@tanstack/ai-persistence': minor +'@tanstack/ai-client': minor +'@tanstack/ai-event-client': minor +'@tanstack/ai-react': minor +'@tanstack/ai-solid': minor +'@tanstack/ai-vue': minor +'@tanstack/ai-svelte': minor +'@tanstack/ai-angular': minor +--- + +Add generation persistence, mirroring chat: media generation runs survive a reload or dropped connection, restoring transparently into the normal hook fields, with optional durable storage of the generated bytes. + +**Generation job store (server).** `withGenerationPersistence` records each run in a dedicated `jobs` (`GenerationJobStore`) store, keyed by the job's `jobId` (`runId`), with `threadId` only an optional link — it no longer overloads the chat `RunStore`. The record holds the activity/provider/model, lifecycle status, result metadata, and (when byte storage is on) the durable artifact refs. `memoryPersistence()` ships an in-memory `jobs` store, and `defineGenerationJobStore` / `defineArtifactStore` / `defineBlobStore` type a custom store inline the way `defineMessageStore` / `defineRunStore` already do. + +**Server-side load (`reconstructGeneration`).** A new `reconstructGeneration(persistence, request, options?)` server helper — the generation parallel of `reconstructChat` — reads a `?jobId=` (or `?threadId=`) from the request, authorizes it via an optional `authorize` callback, and returns `{ resumeSnapshot, activeRun }` JSON so a server-authoritative client restores the last run on mount. Requires the `jobs` store. + +**Media byte storage (server).** When the backend also provides both an `artifacts` (`ArtifactStore`) and a `blobs` (`BlobStore`) store, `withGenerationPersistence` writes each generated file's bytes to the blob store (key `artifacts//`), records an `ArtifactRecord`, and attaches `PersistedArtifactRef`s to the result and the job record. A new `artifactUrl` option stamps a durable app-origin serve URL onto each ref (a new `PersistedArtifactRef.url`) and rewrites the live result's media URL to it, so live and restored results both render media from your own origin instead of the provider's expiring link. Extraction is customizable via `extractArtifacts` / `nameArtifact`; `retrieveArtifact` / `retrieveBlob` (and the shared `artifactBlobKey`) serve the bytes back. `memoryPersistence()` ships in-memory `artifacts`/`blobs` stores; the generation activities gained `threadId` / `runId` options. `@tanstack/ai-utils` adds `base64ToUint8Array`. + +**Client (transparent restore).** Generation hooks (`useGenerateImage`, `useGenerateVideo`, `useGenerateAudio`, `useGenerateSpeech`, `useGeneration`, `useSummarize`, `useTranscription`, and their Solid/Vue/Svelte/Angular equivalents) take the same `persistence` option chat does: `true` (server-driven — cache nothing, hydrate the last run for a stable `threadId` on mount) or a storage adapter (client-driven — write a lightweight snapshot under `generation:` as the run streams and read it back on mount). Restore is **invisible**: it repaints the normal `result` / `status` / `error` fields as if the run had just finished, and keeps `resumeState` for the in-flight run identity — there is no `resumeSnapshot` / `pendingArtifacts` / `resultArtifacts` hook field. With byte storage configured, a restored `result` is rebuilt whole, its media resolved to the durable serve URL and its refs on `result.artifacts`; without it, `status` / `error` restore and `result` stays null. The snapshot never holds the generated bytes and never restarts provider work — generation still only begins on `generate(...)`. + +The `persistence` option reuses the same `ChatStorageAdapter` contract as chat, so the shared `localStoragePersistence` / `sessionStoragePersistence` / `indexedDBPersistence` factories work for generations too with no type argument (a bare call is correct on either side). Untrusted snapshots are validated with the new `parseGenerationResumeSnapshot`. diff --git a/docs/config.json b/docs/config.json index 45c038a42..a2ed617a3 100644 --- a/docs/config.json +++ b/docs/config.json @@ -242,7 +242,7 @@ "label": "Overview", "to": "persistence/overview", "addedAt": "2026-07-22", - "updatedAt": "2026-07-27" + "updatedAt": "2026-07-28" }, { "label": "Chat Persistence", @@ -254,7 +254,24 @@ "label": "Client Persistence", "to": "persistence/client-persistence", "addedAt": "2026-07-22", - "updatedAt": "2026-07-27" + "updatedAt": "2026-07-28" + }, + { + "label": "Generation Persistence", + "to": "persistence/generation-persistence", + "addedAt": "2026-07-28", + "updatedAt": "2026-07-28" + }, + { + "label": "Generation Persistence: Advanced", + "to": "persistence/generation-persistence-advanced", + "addedAt": "2026-07-28" + }, + { + "label": "Keep Generated Files", + "to": "persistence/keep-generated-files", + "addedAt": "2026-07-28", + "updatedAt": "2026-07-28" }, { "label": "Controls", @@ -266,7 +283,7 @@ "label": "Build Your Own Adapter", "to": "persistence/build-your-own-adapter", "addedAt": "2026-07-24", - "updatedAt": "2026-07-27" + "updatedAt": "2026-07-28" }, { "label": "Migrations", @@ -377,7 +394,7 @@ "label": "Transcription", "to": "media/transcription", "addedAt": "2026-04-15", - "updatedAt": "2026-07-03" + "updatedAt": "2026-07-28" }, { "label": "Audio Recording", @@ -389,25 +406,25 @@ "label": "Audio Generation", "to": "media/audio-generation", "addedAt": "2026-04-23", - "updatedAt": "2026-06-08" + "updatedAt": "2026-07-28" }, { "label": "Image Generation", "to": "media/image-generation", "addedAt": "2026-04-15", - "updatedAt": "2026-07-07" + "updatedAt": "2026-07-28" }, { "label": "Video Generation", "to": "media/video-generation", "addedAt": "2026-04-15", - "updatedAt": "2026-07-02" + "updatedAt": "2026-07-28" }, { "label": "Generation Hooks", "to": "media/generation-hooks", "addedAt": "2026-04-15", - "updatedAt": "2026-07-10" + "updatedAt": "2026-07-28" } ] }, diff --git a/docs/media/audio-generation.md b/docs/media/audio-generation.md index 2fbb97651..4b323d1a4 100644 --- a/docs/media/audio-generation.md +++ b/docs/media/audio-generation.md @@ -161,6 +161,10 @@ flow. It mirrors the API of `useGenerateSpeech`, `useGenerateImage`, and other media hooks — see [Generation Hooks](./generation-hooks) for the full shape. +> **Note:** For long tracks, keep the run's status and result across a reload or +> a dropped connection — and the audio after the provider's URL expires — with +> [Generation Persistence](../persistence/generation-persistence). + ### Server (streaming SSE route) ```typescript diff --git a/docs/media/generation-hooks.md b/docs/media/generation-hooks.md index 69cb7b7c9..b87e8b133 100644 --- a/docs/media/generation-hooks.md +++ b/docs/media/generation-hooks.md @@ -19,6 +19,11 @@ keywords: TanStack AI provides framework hooks for every generation type: image, audio, speech, transcription, summarization, and video. Each hook connects to a server endpoint and manages loading, error, and result state for you. +> **Surviving reloads and dropped connections:** every generation hook takes the +> same `persistence` option `useChat` does, so a long run's status and result +> come back after a page reload or a dropped connection. See +> [Generation Persistence](../persistence/generation-persistence). + ## Overview Generation hooks share a consistent API across all media types: diff --git a/docs/media/image-generation.md b/docs/media/image-generation.md index 9f20e59ab..6db756904 100644 --- a/docs/media/image-generation.md +++ b/docs/media/image-generation.md @@ -508,6 +508,9 @@ try { TanStack AI provides React hooks and server-side streaming helpers to build full-stack image generation with minimal boilerplate. +> **Note:** To keep a batch across reloads, or to keep the images after the +> provider's URLs expire, add [Generation Persistence](../persistence/generation-persistence). + ### Streaming Mode (Server Route + Client Hook) **Server** — Create an API route that wraps `generateImage` as a streaming response: diff --git a/docs/media/transcription.md b/docs/media/transcription.md index bd92c633a..f752efa14 100644 --- a/docs/media/transcription.md +++ b/docs/media/transcription.md @@ -397,6 +397,10 @@ export async function POST(request: Request) { TanStack AI provides React hooks and server-side streaming helpers to build full-stack audio transcription with minimal boilerplate. +> **Note:** Transcribing a big file can run long — keep its status and result +> across a reload or a dropped connection with +> [Generation Persistence](../persistence/generation-persistence). + ### Streaming Mode (Server Route + Client Hook) **Server** — Create an API route that wraps `generateTranscription` as a streaming response: diff --git a/docs/media/video-generation.md b/docs/media/video-generation.md index 2386de6fa..518ae5757 100644 --- a/docs/media/video-generation.md +++ b/docs/media/video-generation.md @@ -46,6 +46,13 @@ Currently supported: - **Grok (xAI)**: grok-imagine-video (text-to-video + image-to-video) and grok-imagine-video-1.5 (image-to-video only) models - **fal.ai**: MiniMax, Luma, Kling, Hunyuan, and other hosted video models +> **Video runs take minutes — don't lose them to a reload.** This is the +> strongest case for [Generation Persistence](../persistence/generation-persistence): +> it keeps a record of each job so a page reload or dropped connection picks the +> run back up instead of starting over. And because provider video URLs expire, +> [keep the finished clip](../persistence/keep-generated-files) by saving its +> bytes to your own storage. + ## Basic Usage ### Creating a Video Job diff --git a/docs/persistence/build-your-own-adapter.md b/docs/persistence/build-your-own-adapter.md index 66988ce88..794b16880 100644 --- a/docs/persistence/build-your-own-adapter.md +++ b/docs/persistence/build-your-own-adapter.md @@ -36,19 +36,25 @@ const persistence: ChatTranscriptPersistence = { } ``` -Each store is independent. Provide only the ones you need: `messages` for the -transcript, `runs` for run lifecycle, `interrupts` for durable approvals (needs -`runs`), `metadata` for namespaced key/value state. The middleware turns on +Each store is independent. Provide only the ones you need. For chat: `messages` +for the transcript, `runs` for run lifecycle, `interrupts` for durable approvals +(needs `runs`), `metadata` for namespaced key/value state. For generation: +`jobs` for the generation job lifecycle (the counterpart to `runs`, keyed by +`jobId`), plus `artifacts` and `blobs` to keep generated media bytes (see +[Generation & media stores](#generation--media-stores)). The middleware turns on behavior for whatever stores it finds, so a `messages`-only adapter is a valid adapter. -Those four are the *only* keys `stores` accepts — anything else throws -`Unknown AIPersistence store key` at construction. Need a mutex across +Those seven — `messages`, `runs`, `interrupts`, `metadata`, `jobs`, +`artifacts`, `blobs` — are the *only* keys `stores` accepts; anything else +throws `Unknown AIPersistence store key` at construction. Need a mutex across instances? That is `withLocks`; see [Locks](../advanced/locks). Type each store with its `define*Store` helper — `defineMessageStore`, -`defineRunStore`, `defineInterruptStore`, `defineMetadataStore` — as the sections -below do. Each checks the object against the contract inline (autocomplete, no +`defineRunStore`, `defineInterruptStore`, `defineMetadataStore`, +`defineGenerationJobStore`, `defineArtifactStore`, `defineBlobStore` — as the +sections below do. Each checks the object against the contract inline +(autocomplete, no `: MessageStore` annotation) and composes into `defineAIPersistence`, which tracks **exact presence**: the stores you pass are defined, autocompleted keys on `persistence.stores`, and accessing one you did not pass is a compile error. @@ -475,6 +481,485 @@ export async function POST(request: Request) { } ``` +## Generation & media stores + +Everything above builds a **chat** adapter. [Media generation](./generation-persistence) +persists differently: it does not use the chat `runs` store. Instead +`withGenerationPersistence` requires a `jobs` store — a `GenerationJobStore` +keyed by `jobId` (the run/request id a generation mints), the counterpart to +`runs`. A generation has no conversation, so `threadId` is only an optional +*link* on the job record. Keeping the generated bytes is an optional add-on on +top of `jobs`: an `artifacts` store (metadata) and a `blobs` store (the bytes), +which must be provided **together**. + +These are three more tables alongside the four from the schema in step 1: + +```sql +CREATE TABLE IF NOT EXISTS generation_jobs ( + job_id text PRIMARY KEY NOT NULL, + thread_id text, + activity text NOT NULL, + provider text NOT NULL, + model text NOT NULL, + status text NOT NULL, + started_at integer NOT NULL, + finished_at integer, + error_json text, + result_json text, + artifacts_json text, + usage_json text +); +CREATE TABLE IF NOT EXISTS artifacts ( + artifact_id text PRIMARY KEY NOT NULL, + run_id text NOT NULL, + thread_id text NOT NULL, + name text NOT NULL, + mime_type text NOT NULL, + size integer NOT NULL, + external_url text, + created_at integer NOT NULL +); +CREATE TABLE IF NOT EXISTS blobs ( + key text PRIMARY KEY NOT NULL, + bytes blob NOT NULL, + size integer NOT NULL, + etag text NOT NULL, + content_type text, + custom_metadata_json text, + created_at integer NOT NULL, + updated_at integer NOT NULL +); +``` + +### Jobs: idempotent create, patch, latest-for-thread + +`GenerationJobStore` is the generation analogue of `RunStore`. `createOrResume` +is idempotent — a second call for a `jobId` returns the stored record unchanged, +so resuming a job never resets its `startedAt`, `activity`, or status. +`INSERT ... ON CONFLICT DO NOTHING` gives you that. `update` on an unknown +`jobId` is a no-op. `findLatestForThread` returns the job with the greatest +`startedAt` linked to a thread — `reconstructGeneration` calls it to hydrate the +last generation for a thread on a server-driven client's mount. + +```ts +import { DatabaseSync } from 'node:sqlite' +import { defineGenerationJobStore } from '@tanstack/ai-persistence' +import type { + GenerationJobRecord, + GenerationJobStatus, +} from '@tanstack/ai-persistence' + +function toJobStatus(value: unknown): GenerationJobStatus { + switch (value) { + case 'running': + case 'complete': + case 'error': + case 'interrupted': + return value + default: + throw new TypeError(`Unexpected job status: ${String(value)}`) + } +} + +// `node:sqlite` types columns as a SQL-value union, so coerce/narrow each field +// (String / Number / typeof) and JSON-parse the text columns — no cast. +function mapJob(row: Record): GenerationJobRecord { + return { + jobId: String(row.job_id), + ...(typeof row.thread_id === 'string' ? { threadId: row.thread_id } : {}), + activity: String(row.activity), + provider: String(row.provider), + model: String(row.model), + status: toJobStatus(row.status), + startedAt: Number(row.started_at), + ...(row.finished_at != null ? { finishedAt: Number(row.finished_at) } : {}), + ...(typeof row.error_json === 'string' + ? { error: JSON.parse(row.error_json) } + : {}), + ...(typeof row.result_json === 'string' + ? { result: JSON.parse(row.result_json) } + : {}), + ...(typeof row.artifacts_json === 'string' + ? { artifacts: JSON.parse(row.artifacts_json) } + : {}), + ...(typeof row.usage_json === 'string' + ? { usage: JSON.parse(row.usage_json) } + : {}), + } +} + +function createGenerationJobStore(db: DatabaseSync) { + const select = db.prepare('SELECT * FROM generation_jobs WHERE job_id = ?') + const insert = db.prepare( + `INSERT INTO generation_jobs + (job_id, thread_id, activity, provider, model, status, started_at) + VALUES (?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(job_id) DO NOTHING`, + ) + const latest = db.prepare( + `SELECT * FROM generation_jobs WHERE thread_id = ? + ORDER BY started_at DESC LIMIT 1`, + ) + return defineGenerationJobStore({ + async createOrResume(input) { + const existing = select.get(input.jobId) + if (existing) return mapJob(existing) + const status: GenerationJobStatus = input.status ?? 'running' + insert.run( + input.jobId, + input.threadId ?? null, + input.activity, + input.provider, + input.model, + status, + input.startedAt, + ) + return { + jobId: input.jobId, + activity: input.activity, + provider: input.provider, + model: input.model, + status, + startedAt: input.startedAt, + ...(input.threadId !== undefined ? { threadId: input.threadId } : {}), + } + }, + async update(jobId, patch) { + const sets: Array = [] + const params: Array = [] + if (patch.status !== undefined) { + sets.push('status = ?') + params.push(patch.status) + } + if (patch.finishedAt !== undefined) { + sets.push('finished_at = ?') + params.push(patch.finishedAt) + } + if (patch.error !== undefined) { + sets.push('error_json = ?') + params.push(JSON.stringify(patch.error)) + } + if (patch.result !== undefined) { + sets.push('result_json = ?') + params.push(JSON.stringify(patch.result)) + } + if (patch.artifacts !== undefined) { + sets.push('artifacts_json = ?') + params.push(JSON.stringify(patch.artifacts)) + } + if (patch.usage !== undefined) { + sets.push('usage_json = ?') + params.push(JSON.stringify(patch.usage)) + } + // Empty patch, or an unknown job id, touches nothing (UPDATE no-ops). + if (sets.length === 0) return + params.push(jobId) + db.prepare( + `UPDATE generation_jobs SET ${sets.join(', ')} WHERE job_id = ?`, + ).run(...params) + }, + async get(jobId) { + const row = select.get(jobId) + return row ? mapJob(row) : null + }, + // The most recent job linked to a thread. `reconstructGeneration` calls this + // so a server-driven client (`persistence: true`) hydrates the last + // generation for its thread by the stable thread id, without a job id. + async findLatestForThread(threadId) { + const row = latest.get(threadId) + return row ? mapJob(row) : null + }, + }) +} +``` + +### Artifacts: media metadata + +`ArtifactStore` holds one metadata row per generated file — its `runId`, +`mimeType`, `size`, and a `createdAt`. The bytes live in the blob store below. +`save` is an upsert, `list(runId)` returns every artifact for a run (`[]` when +none). `delete` / `deleteForRun` are optional but cheap to add. + +```ts +import { DatabaseSync } from 'node:sqlite' +import { defineArtifactStore } from '@tanstack/ai-persistence' +import type { ArtifactRecord } from '@tanstack/ai-persistence' + +function mapArtifact(row: Record): ArtifactRecord { + return { + artifactId: String(row.artifact_id), + runId: String(row.run_id), + threadId: String(row.thread_id), + name: String(row.name), + mimeType: String(row.mime_type), + size: Number(row.size), + ...(typeof row.external_url === 'string' + ? { externalUrl: row.external_url } + : {}), + createdAt: Number(row.created_at), + } +} + +function createArtifactStore(db: DatabaseSync) { + const upsert = db.prepare( + `INSERT INTO artifacts + (artifact_id, run_id, thread_id, name, mime_type, size, external_url, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(artifact_id) DO UPDATE SET + run_id = excluded.run_id, thread_id = excluded.thread_id, + name = excluded.name, mime_type = excluded.mime_type, + size = excluded.size, external_url = excluded.external_url, + created_at = excluded.created_at`, + ) + const selectOne = db.prepare('SELECT * FROM artifacts WHERE artifact_id = ?') + const byRun = db.prepare( + 'SELECT * FROM artifacts WHERE run_id = ? ORDER BY created_at ASC', + ) + return defineArtifactStore({ + async save(record) { + upsert.run( + record.artifactId, + record.runId, + record.threadId, + record.name, + record.mimeType, + record.size, + record.externalUrl ?? null, + record.createdAt, + ) + }, + async get(artifactId) { + const row = selectOne.get(artifactId) + return row ? mapArtifact(row) : null + }, + async list(runId) { + return byRun.all(runId).map(mapArtifact) + }, + async delete(artifactId) { + db.prepare('DELETE FROM artifacts WHERE artifact_id = ?').run(artifactId) + }, + async deleteForRun(runId) { + db.prepare('DELETE FROM artifacts WHERE run_id = ?').run(runId) + }, + }) +} +``` + +### Blobs: the bytes + +`BlobStore` is a small object store. `withGenerationPersistence` writes each +generated file under the key `artifacts//`, so a +prefix-filtered `list({ prefix: 'artifacts//' })` enumerates a run's +media. `put` accepts any `BlobBody` (a stream, buffer, string, or `Blob`); the +helper below normalizes it to bytes. `list` matches `prefix` literally and pages +with a keyset cursor. + +```ts +import { DatabaseSync } from 'node:sqlite' +import { defineBlobStore } from '@tanstack/ai-persistence' +import type { + BlobBody, + BlobObject, + BlobRecord, +} from '@tanstack/ai-persistence' + +async function toBytes(body: BlobBody): Promise { + if (typeof body === 'string') return new TextEncoder().encode(body) + if (body instanceof ArrayBuffer) return new Uint8Array(body.slice(0)) + if (ArrayBuffer.isView(body)) { + return new Uint8Array(body.buffer, body.byteOffset, body.byteLength).slice() + } + if (body instanceof Blob) { + return new Uint8Array(await body.arrayBuffer()) + } + // ReadableStream: drain it into one buffer. + const reader = body.getReader() + const chunks: Array = [] + let total = 0 + for (;;) { + const { done, value } = await reader.read() + if (done) break + chunks.push(value) + total += value.byteLength + } + const bytes = new Uint8Array(total) + let offset = 0 + for (const chunk of chunks) { + bytes.set(chunk, offset) + offset += chunk.byteLength + } + return bytes +} + +function mapBlobRecord(row: Record): BlobRecord { + return { + key: String(row.key), + ...(row.size != null ? { size: Number(row.size) } : {}), + ...(typeof row.etag === 'string' ? { etag: row.etag } : {}), + ...(typeof row.content_type === 'string' + ? { contentType: row.content_type } + : {}), + ...(typeof row.custom_metadata_json === 'string' + ? { customMetadata: JSON.parse(row.custom_metadata_json) } + : {}), + ...(row.created_at != null ? { createdAt: Number(row.created_at) } : {}), + ...(row.updated_at != null ? { updatedAt: Number(row.updated_at) } : {}), + } +} + +function blobObject(record: BlobRecord, bytes: Uint8Array): BlobObject { + return { + ...record, + body: new ReadableStream({ + start(controller) { + controller.enqueue(bytes.slice()) + controller.close() + }, + }), + arrayBuffer() { + const copy = new ArrayBuffer(bytes.byteLength) + new Uint8Array(copy).set(bytes) + return Promise.resolve(copy) + }, + text: () => Promise.resolve(new TextDecoder().decode(bytes)), + } +} + +function createBlobStore(db: DatabaseSync) { + const upsert = db.prepare( + `INSERT INTO blobs + (key, bytes, size, etag, content_type, custom_metadata_json, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(key) DO UPDATE SET + bytes = excluded.bytes, size = excluded.size, etag = excluded.etag, + content_type = excluded.content_type, + custom_metadata_json = excluded.custom_metadata_json, + updated_at = excluded.updated_at`, + ) + const selectCreated = db.prepare('SELECT created_at FROM blobs WHERE key = ?') + const selectOne = db.prepare('SELECT * FROM blobs WHERE key = ?') + return defineBlobStore({ + async put(key, body, options) { + const bytes = await toBytes(body) + const now = Date.now() + const prior = selectCreated.get(key) + const createdAt = + prior && prior.created_at != null ? Number(prior.created_at) : now + const etag = String(now) + upsert.run( + key, + bytes, + bytes.byteLength, + etag, + options?.contentType ?? null, + options?.customMetadata ? JSON.stringify(options.customMetadata) : null, + createdAt, + now, + ) + return { + key, + size: bytes.byteLength, + etag, + createdAt, + updatedAt: now, + ...(options?.contentType !== undefined + ? { contentType: options.contentType } + : {}), + ...(options?.customMetadata !== undefined + ? { customMetadata: options.customMetadata } + : {}), + } + }, + async get(key) { + const row = selectOne.get(key) + if (!row) return null + const bytes = + row.bytes instanceof Uint8Array ? row.bytes : new Uint8Array() + return blobObject(mapBlobRecord(row), bytes) + }, + async head(key) { + const row = selectOne.get(key) + return row ? mapBlobRecord(row) : null + }, + async delete(key) { + db.prepare('DELETE FROM blobs WHERE key = ?').run(key) + }, + async list(options) { + if (options?.limit === 0) return { objects: [], truncated: false } + // Escape LIKE metacharacters so `prefix` matches literally, then page with + // a keyset cursor (keys strictly greater than the last returned key). + const escaped = (options?.prefix ?? '').replace( + /[\\%_]/g, + (ch) => `\\${ch}`, + ) + const params: Array = [`${escaped}%`] + let where = "key LIKE ? ESCAPE '\\'" + if (options?.cursor !== undefined) { + where += ' AND key > ?' + params.push(options.cursor) + } + let sql = `SELECT * FROM blobs WHERE ${where} ORDER BY key ASC` + const limit = options?.limit + if (limit !== undefined) { + sql += ' LIMIT ?' // fetch one extra row to detect truncation + params.push(limit + 1) + } + const rows = db + .prepare(sql) + .all(...params) + .map(mapBlobRecord) + if (limit !== undefined && rows.length > limit) { + const page = rows.slice(0, limit) + const cursor = page.at(-1)?.key + return { + objects: page, + truncated: true, + ...(cursor !== undefined ? { cursor } : {}), + } + } + return { objects: rows, truncated: false } + }, + }) +} +``` + +### Assemble a generation adapter + +Hand the three stores to `defineAIPersistence` the same way. `jobs` alone is a +valid generation adapter (job records, no byte storage); add `artifacts` + +`blobs` — together — to keep the media: + +```ts +import { DatabaseSync } from 'node:sqlite' +import { defineAIPersistence } from '@tanstack/ai-persistence' +// The three generation store factories and the schema string, from your modules. +import { createArtifactStore } from './artifact-store' +import { createBlobStore } from './blob-store' +import { createGenerationJobStore } from './generation-job-store' +import { GENERATION_SCHEMA_SQL } from './generation-schema' + +export function generationPersistence(options: { + url: string + migrate?: boolean +}) { + const db = new DatabaseSync(options.url) + if (options.migrate) db.exec(GENERATION_SCHEMA_SQL) + return defineAIPersistence({ + stores: { + jobs: createGenerationJobStore(db), + artifacts: createArtifactStore(db), + blobs: createBlobStore(db), + }, + }) +} +``` + +Pass the result to `withGenerationPersistence` on a `generateImage` / +`generateVideo` / … call; see [Generation persistence](./generation-persistence). +You can also fold these stores into an existing chat adapter with +`composePersistence`, so one backend serves both `withPersistence` and +`withGenerationPersistence`. + ## Existing database: map the contracts onto your schema You do not have to create the four tables above. If you already have a database, @@ -701,6 +1186,152 @@ composite identity. A stored `null` is indistinguishable from absence at the typ level, so wrap a value you must persist as `null` (e.g. `{ value: null }`), or reject nullish values outright the way the SQLite store above does. +### GenerationJobStore + +The generation counterpart to `RunStore`, keyed by `jobId` rather than a +conversation `threadId`. `withGenerationPersistence` requires this store, not +`runs`. + +```ts +import type { PersistedArtifactRef, TokenUsage } from '@tanstack/ai' + +type GenerationJobStatus = 'running' | 'complete' | 'error' | 'interrupted' + +interface GenerationJobRecord { + jobId: string + threadId?: string // optional link to a chat conversation + activity: string // 'image' | 'audio' | 'tts' | 'video' | 'transcription' + provider: string + model: string + status: GenerationJobStatus + startedAt: number // epoch ms + finishedAt?: number // epoch ms, set once the job reaches a terminal status + error?: { message: string; code?: string } + result?: unknown // terminal result metadata (ids, urls) — never media bytes + artifacts?: Array // present with an artifacts + blobs backend + usage?: TokenUsage +} + +interface GenerationJobStore { + createOrResume(input: { + jobId: string + activity: string + provider: string + model: string + startedAt: number + threadId?: string + status?: GenerationJobStatus + }): Promise + update( + jobId: string, + patch: Partial< + Pick< + GenerationJobRecord, + 'status' | 'finishedAt' | 'error' | 'result' | 'artifacts' | 'usage' + > + >, + ): Promise + get(jobId: string): Promise + // The most recent job linked to a thread (greatest `startedAt`), or null. + // Optional, but implement it: `reconstructGeneration` feature-detects it to + // hydrate the last generation for a thread by its stable thread id. + findLatestForThread?(threadId: string): Promise +} +``` + +Implement `createOrResume` idempotently: a second call for an existing `jobId` +returns the stored record unchanged (`startedAt` / `activity` / `provider` / +`model` / `threadId` are not mutated), which is what makes resuming a job safe. +`update` against an unknown `jobId` is a no-op. + +### ArtifactStore + +Metadata rows for persisted media. The bytes live in a `BlobStore`; this record +holds the descriptive metadata and an optional `externalUrl` for reference-only +backends. Provide it together with a `BlobStore` to keep generated bytes. + +```ts +interface ArtifactRecord { + artifactId: string + runId: string + threadId: string + name: string + mimeType: string + size: number + externalUrl?: string + createdAt: number // epoch ms +} + +interface ArtifactStore { + save(record: ArtifactRecord): Promise + get(artifactId: string): Promise + list(runId: string): Promise> // [] when the run has none + delete?(artifactId: string): Promise + deleteForRun?(runId: string): Promise +} +``` + +### BlobStore + +A durable object/blob store for the bytes. `withGenerationPersistence` writes +each generated file under the key `artifacts//`. + +```ts +type BlobBody = + | ReadableStream + | ArrayBuffer + | ArrayBufferView + | string + | Blob + +interface BlobRecord { + key: string + size?: number + etag?: string + contentType?: string + customMetadata?: Record + createdAt?: number // epoch ms first written + updatedAt?: number // epoch ms last overwritten +} + +interface BlobObject extends BlobRecord { + arrayBuffer(): Promise + text(): Promise + body?: ReadableStream +} + +interface BlobListPage { + objects: Array + cursor?: string // present only when `truncated` + truncated?: boolean +} + +interface BlobPutOptions { + contentType?: string + customMetadata?: Record +} + +interface BlobListOptions { + prefix?: string + cursor?: string + limit?: number +} + +interface BlobStore { + put(key: string, body: BlobBody, options?: BlobPutOptions): Promise + get(key: string): Promise + head(key: string): Promise + delete(key: string): Promise + list(options?: BlobListOptions): Promise +} +``` + +`list` matches `prefix` literally and case-sensitively (escape SQL `LIKE` +metacharacters). When `limit` is given and more keys match, return +`truncated: true` with a `cursor`; passing that cursor back returns the +strictly-following keys, so paging visits every key exactly once. `limit: 0` +yields an empty, untruncated page. + ## Where to go next - [Controls](./controls): compose stores from different systems. diff --git a/docs/persistence/client-persistence.md b/docs/persistence/client-persistence.md index 29b1103a9..cf65c4a06 100644 --- a/docs/persistence/client-persistence.md +++ b/docs/persistence/client-persistence.md @@ -70,6 +70,12 @@ pointer. On the next load `useChat` reads it and: - **`true`** is server-authoritative. - **`false`** (or omitted) is off: messages live in memory only and a reload starts empty. +The generation hooks (`useGenerateImage` / `useGenerateVideo` / …) make the same +choice: `persistence: true` is server-driven and a storage adapter is +client-driven, so a long media run survives a reload the way a conversation does. +See [Generation persistence](./generation-persistence) for the generation-specific +setup. + ### An adapter: client-authoritative Pass the adapter directly, `persistence: localStoragePersistence()`. The diff --git a/docs/persistence/generation-persistence-advanced.md b/docs/persistence/generation-persistence-advanced.md new file mode 100644 index 000000000..7dc79dd01 --- /dev/null +++ b/docs/persistence/generation-persistence-advanced.md @@ -0,0 +1,118 @@ +--- +title: 'Generation Persistence: Advanced' +id: generation-persistence-advanced +--- + +# Generation Persistence: Advanced + +The deeper parts of [Generation persistence](./generation-persistence): +reconnecting a stream that is still running, the in-flight `resumeState`, seeding +state yourself, securing the hydration endpoint, and exactly what the record +holds. Start with the main guide; come here when you need one of these. + +## Reconnect to a run that is still streaming + +The server-driven endpoint in the main guide already wires this: a `durability` +adapter on `toServerSentEventsResponse`, plus a `GET` that replays the run from +the log when the request carries a resume cursor. On the client there is nothing +to add. A connection dropped mid-generation re-attaches on its own through +`fetchServerSentEvents` or `fetchHttpStream`, the same adapters `useChat` uses. +In production, swap `memoryStream` for `durableStream` from +`@tanstack/ai-durable-stream`, where requests span processes. + +A full page reload is different: the hooks never start or resume a run on mount, +and the record alone cannot re-attach to the stream. It records that a run was +in flight, not a stream position. After a reload, a non-null `resumeState` is +informational: "a run was still going when the page closed." See +[Resumable Streams](../resumable-streams/overview) for the durability contract, +production adapters, and the one-time-side-effects note. + +## The in-flight `resumeState` + +`resumeState` is non-null only while a run is streaming. It is the identity of +the live run (`threadId` / `runId`, plus `pendingArtifacts` while media streams), +and it clears when the run ends. Use `status` / `result` to display a finished +run; use `resumeState` to tell that a run was still going. + +For a client-driven store: + +- **The `id` is the storage key.** The snapshot is written under + `generation:` (plus the adapter's `keyPrefix`), so a stable `id` is what + makes the record findable after a reload. Omitting `id` generates a fresh key + per mount. The `generation:` segment also keeps a chat client with the same id + from colliding. +- `stop()` marks the record no longer resumable; `reset()` deletes it. +- The snapshot never triggers work. A generation starts only on `generate(...)`. + +## Seed the state yourself + +If your app manages storage itself (custom backend, SSR-provided state), read +the value and pass it as `initialResumeSnapshot`, which skips the automatic read. +Validate untrusted values with `parseGenerationResumeSnapshot`: + +```tsx +import { parseGenerationResumeSnapshot } from '@tanstack/ai-client' +import { fetchServerSentEvents, useGenerateImage } from '@tanstack/ai-react' + +function HeroImageGenerator({ stored }: { stored: unknown }) { + const seed = parseGenerationResumeSnapshot(stored) + const image = useGenerateImage({ + id: 'hero-image', + connection: fetchServerSentEvents('/api/generate/image'), + ...(seed ? { initialResumeSnapshot: seed } : {}), + }) + return

{image.status}

+} +``` + +## Secure the hydration endpoint + +`reconstructGeneration(persistence, request)` reads a `?jobId` (preferred) or +`?threadId` from the query and returns the last job's record plus a cursor to any +still-running run. It resolves the latest job for a thread through the store's +`findLatestForThread`. It does **not** enforce tenancy on its own, so on a public +route pass its `authorize` option (session to owned thread/job) before returning +anything: + +```ts group=generation-authorize +import { + memoryPersistence, + reconstructGeneration, +} from '@tanstack/ai-persistence' + +const persistence = memoryPersistence() + +// Replace these with your real session + ownership checks. +async function getSessionUserId(_request: Request): Promise { + return null +} +async function userOwnsThread( + _userId: string, + _threadId: string, +): Promise { + return true +} + +export function GET(request: Request): Response | Promise { + return reconstructGeneration(persistence, request, { + authorize: async (id, req) => { + const userId = await getSessionUserId(req) + return userId !== null && (await userOwnsThread(userId, id)) + }, + }) +} +``` + +## What the record holds + +The record, hydrated from the server or read from the browser snapshot, holds +run identity and result metadata (ids, model, status, a video `jobId`, an expiry +timestamp), never the generated bytes. On its own it does not carry a usable +media URL, so a reload restores `status` and `error` while `result` stays `null`. + +To bring the media back too, persist the bytes on the server with an `artifacts` ++ `blobs` backend and stamp a durable serve URL onto each ref with +`withGenerationPersistence`'s `artifactUrl` (see +[Keep generated files](./keep-generated-files)). Those durable refs travel on the +record, so the restored `result` is rebuilt with its media resolved to your own +origin and its refs on `result.artifacts`. diff --git a/docs/persistence/generation-persistence.md b/docs/persistence/generation-persistence.md new file mode 100644 index 000000000..aa23a7fac --- /dev/null +++ b/docs/persistence/generation-persistence.md @@ -0,0 +1,192 @@ +--- +title: Generation Persistence +id: generation-persistence +--- + +# Generation Persistence + +Media generation takes time, and video can take minutes. If the user reloads the +page or their connection drops mid-run, that run is easy to lose. Generation +persistence keeps a small record of each run and restores it into the hook's +normal `status` / `result` / `error` fields, so a reload reads exactly like a +fresh run. + +Reach for it when a run is long enough that a reload matters: video, batch +images, long audio, a big transcription. For a quick one-shot image you show and +forget, skip it. + +## Choose a mode + +Generation persistence mirrors chat's `persistence` option: + +- **`persistence: true`** — server-driven. The client caches nothing; on mount + it hydrates the last run for its `threadId` from the server. +- **`persistence: `** (e.g. `localStoragePersistence()`) — client-driven. + The client keeps a small snapshot in browser storage. +- **omitted / `false`** — off. The run lives in memory only; a reload starts + empty. + +Either way the server keeps a **generation job** record: +`withGenerationPersistence(persistence)` needs a `jobs` store (a +`GenerationJobStore` keyed by the run's `jobId`; a generation has no conversation +of its own, so `threadId` is only an optional link). `memoryPersistence()` ships +one out of the box; see +[Build your own adapter](./build-your-own-adapter#generation--media-stores) for +your own backend. + +The record never holds the generated bytes, so on its own a reload restores +`status` and `error` while `result` stays `null`. To bring the media back too, +add server byte storage: see [Keep generated files](./keep-generated-files). + +## Server-driven: `persistence: true` + +One `POST` that runs the generation, and a `GET` on the same route that answers +the mount-time hydration with `reconstructGeneration`: + +```ts group=generation-server-driven +import { + generateImage, + generationParamsFromRequest, + memoryStream, + resumeServerSentEventsResponse, + toServerSentEventsResponse, +} from '@tanstack/ai' +import { openaiImage } from '@tanstack/ai-openai' +import { + memoryPersistence, + reconstructGeneration, + withGenerationPersistence, +} from '@tanstack/ai-persistence' + +const persistence = memoryPersistence() + +export async function POST(request: Request) { + const durability = memoryStream(request) + const { input, threadId } = await generationParamsFromRequest('image', request) + + if (typeof input.prompt !== 'string') { + throw new Error('This endpoint accepts text image prompts only.') + } + + const stream = generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: input.prompt, + // Link the job to the thread so the GET can find the last run for it. + ...(threadId !== undefined ? { threadId } : {}), + stream: true, + // `artifactUrl` makes the restored media render from your own origin. It is + // optional — see Keep generated files for the serve route it points at. + middleware: [ + withGenerationPersistence(persistence, { + artifactUrl: (ref) => `/api/generate/image/artifact?id=${ref.artifactId}`, + }), + ], + }) + + return toServerSentEventsResponse(stream, { + durability: { adapter: durability }, + }) +} + +export function GET(request: Request): Response | Promise { + const durability = memoryStream(request) + // A reconnecting client carries a resume cursor; replay the live stream. + if (durability.resumeFrom() !== null) { + return resumeServerSentEventsResponse({ adapter: durability }) + } + // Otherwise this is the mount-time hydration (?threadId). Guard it with the + // `authorize` option in multi-user apps (see the advanced guide). + return reconstructGeneration(persistence, request) +} +``` + +On the client, pass a connection, a stable `threadId`, and `persistence: true`. +The hook is transparent: a reload repaints the last run into the same fields a +fresh run uses, the way `useChat` restores into `messages`: + +```tsx +import { fetchServerSentEvents, useGenerateImage } from '@tanstack/ai-react' + +const connection = fetchServerSentEvents('/api/generate/image') + +export function HeroImageGenerator({ threadId }: { threadId: string }) { + const image = useGenerateImage({ + id: 'hero-image', + threadId, + connection, + persistence: true, + }) + + return ( +
+ + + {image.status === 'success' ? ( +

Last run finished{image.result?.id ? ` (${image.result.id})` : ''}.

+ ) : null} + {image.error ?

Last run failed: {image.error.message}

: null} + {image.result?.images.map((img, index) => + img.url ? : null, + )} +
+ ) +} +``` + +Because the server stamps a durable `artifactUrl`, the restored +`image.result.images[i].url` serves from your own origin, so the image renders +after a reload just as it did live. You never fetch or seed anything: the thread +id is the stable key, and a reload or the same thread on another device follow +the identical path. + +## Client-driven: a storage adapter + +Pass a storage adapter instead, and give the hook a stable `id`. The client +writes the snapshot to browser storage as the run streams and reads it back on +mount, with no server `GET` for hydration. Everything else about the hook is the +same as above: + +```tsx +import { localStoragePersistence } from '@tanstack/ai-client' +import { fetchServerSentEvents, useGenerateImage } from '@tanstack/ai-react' + +// A bare call is all it takes, no type argument. +const snapshots = localStoragePersistence({ keyPrefix: 'my-app:' }) + +function HeroImageGenerator() { + const image = useGenerateImage({ + id: 'hero-image', + connection: fetchServerSentEvents('/api/generate/image'), + persistence: snapshots, + }) + // The render is identical to the server-driven example above. + return

{image.status}

+} +``` + +Two things to keep in mind: + +- **The hook is transparent.** After a reload it repaints `status` + (`'idle'` / `'generating'` / `'success'` / `'error'`), `error`, and `result` + just like a live run. +- **`result` needs byte storage to come back.** A browser snapshot never holds + the bytes, so on its own a reload restores `status` and `error` while `result` + stays `null`. Add server byte storage and `artifactUrl` + ([Keep generated files](./keep-generated-files)) and the restored `result` is + rebuilt, its media served from your own origin. + +## Going further + +- [Keep generated files](./keep-generated-files): store the generated bytes so + the media itself survives, not just the record. +- [Generation persistence: advanced](./generation-persistence-advanced): + reconnecting a dropped stream, the in-flight `resumeState`, seeding state + yourself, and securing the hydration endpoint. diff --git a/docs/persistence/keep-generated-files.md b/docs/persistence/keep-generated-files.md new file mode 100644 index 000000000..3da0ac1c8 --- /dev/null +++ b/docs/persistence/keep-generated-files.md @@ -0,0 +1,125 @@ +--- +title: Keep Generated Files +id: keep-generated-files +--- + +# Keep Generated Files + +Provider URLs for generated media expire. A Sora clip, a batch of images, a long +audio track — the model hands you a URL that stops working after a while, and +once it does the output is gone. To keep the output, save the generated bytes to +your own storage and serve them from your own origin, where they outlive the +provider's link. + +This is a server-side opt-in that layers on **top of** the `jobs` store +[Generation persistence](./generation-persistence) already requires. Byte storage +adds two more stores: an `artifacts` store (metadata) and a `blobs` store (the +bytes). They are a both-or-neither pair — provide both and +`withGenerationPersistence` writes each generated file's bytes to the blob store, +records an `ArtifactRecord`, and attaches durable references to the result and the +job record; provide neither and only the job record is kept. +`memoryPersistence()` ships all three stores (`jobs`, `artifacts`, `blobs`), so it +works out of the box; any backend that implements `ArtifactStore` and `BlobStore` +(see [Build your own adapter](./build-your-own-adapter#generation--media-stores)) +works the same way. + +## Serve the stored bytes + +The bytes land under the blob key `artifacts//`. To fetch a +generated file later (to render it, download it, or hand it to another request), +add a `GET` route that reads the artifact back with the `retrieveArtifact` / +`retrieveBlob` helpers and streams it from your own origin: + +```ts group=generation-bytes +import { + generateImage, + generationParamsFromRequest, + toServerSentEventsResponse, +} from '@tanstack/ai' +import { openaiImage } from '@tanstack/ai-openai' +import { + memoryPersistence, + retrieveArtifact, + retrieveBlob, + withGenerationPersistence, +} from '@tanstack/ai-persistence' + +const persistence = memoryPersistence() + +export async function POST(request: Request) { + const { input } = await generationParamsFromRequest('image', request) + + if (typeof input.prompt !== 'string') { + throw new Error('This endpoint accepts text image prompts only.') + } + + const stream = generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: input.prompt, + stream: true, + middleware: [ + withGenerationPersistence(persistence, { + // Stamp the durable serve URL (the GET route below) onto every + // persisted artifact ref, and rewrite the live result's media field to + // it. Both the live and the restored result then render from your own + // origin instead of the provider's expiring link. + artifactUrl: (ref) => `/api/generate/image?id=${ref.artifactId}`, + }), + ], + }) + + return toServerSentEventsResponse(stream) +} + +// Serve a stored artifact's bytes by id. This is a plain file endpoint: it +// serves one stored file, it does not resume a run or rebuild a conversation. +export async function GET(request: Request) { + const artifactId = new URL(request.url).searchParams.get('id') + if (!artifactId) return new Response('missing id', { status: 400 }) + + const artifact = await retrieveArtifact(persistence, artifactId) + if (!artifact) return new Response('not found', { status: 404 }) + + const blob = await retrieveBlob(persistence, artifact) + if (!blob) return new Response('not found', { status: 404 }) + + return new Response(blob.body ?? (await blob.arrayBuffer()), { + headers: { + 'content-type': artifact.mimeType, + 'content-length': String(artifact.size), + }, + }) +} +``` + +`memoryPersistence` keeps everything in process memory, which is right for +development and tests; point `jobs` / `artifacts` / `blobs` at a durable backend +for production. Control what gets captured with `withGenerationPersistence`'s +`extractArtifacts` (return your own descriptors) and `nameArtifact` (name each +file) options. + +## Wire the durable URL through to the client + +`artifactUrl` is what makes the stored bytes reachable from the client without +any extra plumbing. For each persisted ref it returns the app-origin URL that +serves those bytes (the `GET` route above), and `withGenerationPersistence` +does two things with it: it stamps the URL onto the ref (`ref.url`) and it +rewrites the live result's media field to the same URL — `result.images[i].url` +for images, `result.url` for a video, `result.audio.url` for audio. So the live +result already points at your origin, not the provider's expiring link. + +Those durable refs ride along on `result.artifacts`, and they are what a reload +restores from. In [Generation persistence](./generation-persistence), the +generation hook rebuilds `result` from the persisted refs on mount, resolving +each media field to its durable `ref.url` — so the restored result renders the +same media the live run showed. The refs are the only artifact surface: there +are no separate top-level artifact fields on the hook, final refs live on +`result.artifacts` and in-flight ones on `resumeState.pendingArtifacts`. + +## Where to go next + +- [Generation persistence](./generation-persistence): the two-mode record that + survives a reload or a dropped connection, and the `jobs` store byte storage + builds on. +- [Build your own adapter](./build-your-own-adapter#generation--media-stores): a + custom `ArtifactStore` / `BlobStore` on your own database. diff --git a/docs/persistence/overview.md b/docs/persistence/overview.md index 73d037a20..cf2ef5f50 100644 --- a/docs/persistence/overview.md +++ b/docs/persistence/overview.md @@ -274,6 +274,21 @@ Why this wins over the alternatives: Client-only persistence can't do multi-device and bloats storage. Caching everything client-side duplicates the source of truth. This combination avoids both: one server-resolved `GET` on mount restores history and rejoins any live run, so a reload and a second device follow the identical path. +## Generation persistence + +Everything above is about chat. Media generation — image, audio, TTS, video, and +transcription — persists as a sibling to it: the generation hooks +(`useGenerateImage`, `useGenerateVideo`, …) take the same `persistence: true | +adapter` modes, so a long run's status and result survive a reload or a dropped +connection the same way a conversation does. On the server it is backed by a +`jobs` store (the generation counterpart to chat's `runs`), with optional +`artifacts` + `blobs` stores to keep the generated bytes after the provider's +URLs expire. + +See [Generation persistence](./generation-persistence) for the two modes and the +server wiring, and [Keep generated files](./keep-generated-files) for byte +storage. + ## The store contract Server **state** persistence is a set of stores. Middleware activates behavior @@ -286,6 +301,13 @@ from whichever stores are present (with entrypoint requirements — see | `runs` | Run status, timing, errors, and usage. Required on full `ChatPersistence`. | | `interrupts` | Pending, resolved, or cancelled human/tool waits (needs `runs`). | | `metadata` | App and integration key/value state. | +| `jobs` | Generation run status, result metadata, and artifact refs, keyed by `jobId`. Required by generation persistence. | +| `artifacts` | Generated-file metadata (needs `blobs`). | +| `blobs` | The generated bytes (needs `artifacts`). | + +The last three are the generation counterpart to the chat stores, used by +`withGenerationPersistence` rather than `withPersistence` — see +[Generation persistence](./generation-persistence). Named shapes: `ChatTranscriptStores` (messages floor), `ChatPersistenceStores` (all four), `ChatWithInterruptsStores`. See [Controls](./controls). @@ -307,6 +329,8 @@ matching your database loads itself. The full skill list is in - [Chat persistence](./chat-persistence): the server middleware, the authoritative-history contract, and durable interrupts. - [Client persistence](./client-persistence): client- vs server-authoritative modes (`persistence: true`), reload restore, storage backends, and mid-stream rejoin. +- [Generation persistence](./generation-persistence): the same modes for media runs (image, audio, TTS, video, transcription), backed by a `jobs` store. +- [Keep generated files](./keep-generated-files): save the generated bytes to your own storage so they outlive the provider's expiring URLs. - [Controls](./controls): compose backends per store and choose which stores to run. - [Build your own adapter](./build-your-own-adapter): a complete SQLite example on the core, plus the store interface reference. - [Resumable streams](../resumable-streams/overview): the delivery-durability layer in full. diff --git a/examples/ts-react-chat/src/routes/generations.image.tsx b/examples/ts-react-chat/src/routes/generations.image.tsx index 044570498..8d942129f 100644 --- a/examples/ts-react-chat/src/routes/generations.image.tsx +++ b/examples/ts-react-chat/src/routes/generations.image.tsx @@ -2,10 +2,22 @@ import { useState } from 'react' import { createFileRoute } from '@tanstack/react-router' import { useGenerateImage } from '@tanstack/ai-react' import type { UseGenerateImageReturn } from '@tanstack/ai-react' -import { fetchServerSentEvents } from '@tanstack/ai-client' +import { + fetchServerSentEvents, + localStoragePersistence, +} from '@tanstack/ai-client' import { resolveMediaPrompt } from '@tanstack/ai' import { generateImageFn, generateImageStreamFn } from '../lib/server-fns' +// Reuse the shared web-storage adapter for the lightweight generation resume +// snapshot. Only run identity, status, errors, and result metadata are stored +// — never the generated image bytes. A bare call needs no type argument: the +// adapter defaults to the generation snapshot shape. The client namespaces its +// record under `generation:` and reads it back on mount, repainting the +// hook's normal `status` / `result` / `error` fields so the last run's outcome +// survives a full page reload. +const imageSnapshots = localStoragePersistence({ keyPrefix: 'example:' }) + function StreamingImageGeneration() { const [prompt, setPrompt] = useState('') const [numberOfImages, setNumberOfImages] = useState(1) @@ -69,6 +81,54 @@ function ServerFnImageGeneration() { ) } +function PersistedImageGeneration() { + const [prompt, setPrompt] = useState('') + const [numberOfImages, setNumberOfImages] = useState(1) + + const hookReturn = useGenerateImage({ + id: 'persisted-image', + connection: fetchServerSentEvents('/api/generate/image'), + persistence: imageSnapshots, + }) + + return ( +
+
+ {hookReturn.resumeState ? ( +

+ Run in flight:{' '} + {hookReturn.resumeState.runId} +

+ ) : hookReturn.status === 'success' ? ( +

+ Last run:{' '} + + success + {hookReturn.result?.model ? ` (${hookReturn.result.model})` : ''} + +

+ ) : hookReturn.status === 'error' ? ( +

+ Last run:{' '} + + error{hookReturn.error ? ` — ${hookReturn.error.message}` : ''} + +

+ ) : ( +

No persisted run yet.

+ )} +
+ +
+ ) +} + function ImageGenerationUI({ prompt, setPrompt, @@ -170,9 +230,9 @@ function ImageGenerationUI({ } function ImageGenerationPage() { - const [mode, setMode] = useState<'streaming' | 'direct' | 'server-fn'>( - 'streaming', - ) + const [mode, setMode] = useState< + 'streaming' | 'direct' | 'server-fn' | 'persisted' + >('streaming') return (
@@ -215,6 +275,16 @@ function ImageGenerationPage() { > Server Fn +
@@ -225,8 +295,10 @@ function ImageGenerationPage() { ) : mode === 'direct' ? ( - ) : ( + ) : mode === 'server-fn' ? ( + ) : ( + )} diff --git a/packages/ai-angular/src/index.ts b/packages/ai-angular/src/index.ts index 2c8dbef90..fc92c73d2 100644 --- a/packages/ai-angular/src/index.ts +++ b/packages/ai-angular/src/index.ts @@ -109,4 +109,10 @@ export { type VideoGenerateInput, type VideoGenerateResult, type VideoStatusInfo, + type GenerationPersistence, + type GenerationResumeSnapshot, + type GenerationResumeState, + type GenerationResumeStatus, + type GenerationPendingArtifact, } from '@tanstack/ai-client' +export type { PersistedArtifactRef } from '@tanstack/ai/client' diff --git a/packages/ai-angular/src/inject-generate-audio.ts b/packages/ai-angular/src/inject-generate-audio.ts index 188a032e0..b1d5a4323 100644 --- a/packages/ai-angular/src/inject-generate-audio.ts +++ b/packages/ai-angular/src/inject-generate-audio.ts @@ -1,4 +1,5 @@ import { injectGeneration } from './inject-generation' +import { reconstructAudioResult } from '@tanstack/ai-client' import type { AudioGenerationResult, StreamChunk } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -10,13 +11,22 @@ import type { } from '@tanstack/ai-client' import type { Signal } from '@angular/core' import type { ReactiveOption } from './internal/to-reactive' +import type { + InjectGenerationOptions, + InjectGenerationResult, +} from './inject-generation' /** * Options for the injectGenerateAudio injectable. * * @template TOutput - The output type after optional transform (defaults to AudioGenerationResult) */ -export interface InjectGenerateAudioOptions { +export interface InjectGenerateAudioOptions< + TOutput = AudioGenerationResult, +> extends Pick< + InjectGenerationOptions, + 'persistence' | 'threadId' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for audio generation */ @@ -48,7 +58,9 @@ export interface InjectGenerateAudioOptions { * * @template TOutput - The output type (after optional transform) */ -export interface InjectGenerateAudioResult { +export interface InjectGenerateAudioResult< + TOutput = AudioGenerationResult, +> extends Omit, 'generate'> { /** Trigger audio generation */ generate: (input: AudioGenerateInput) => Promise /** The generation result containing audio, or null */ @@ -59,10 +71,6 @@ export interface InjectGenerateAudioResult { error: Signal /** Current state of the generation */ status: Signal - /** Abort the current generation */ - stop: () => void - /** Clear result, error, and return to idle */ - reset: () => void } /** @@ -109,19 +117,15 @@ export function injectGenerateAudio( hookName: 'injectGenerateAudio', outputKind: 'audio' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - injectGeneration({ - ...options, - devtools, - }) + const generation = injectGeneration< + AudioGenerateInput, + AudioGenerationResult, + TTransformed + >({ + ...options, + devtools, + reconstructResult: reconstructAudioResult, + }) - return { - generate: generate as (input: AudioGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + return generation } diff --git a/packages/ai-angular/src/inject-generate-image.ts b/packages/ai-angular/src/inject-generate-image.ts index 3c3c2e4fe..84c66ac58 100644 --- a/packages/ai-angular/src/inject-generate-image.ts +++ b/packages/ai-angular/src/inject-generate-image.ts @@ -1,4 +1,5 @@ import { injectGeneration } from './inject-generation' +import { reconstructImageResult } from '@tanstack/ai-client' import type { Signal } from '@angular/core' import type { ImageGenerationResult } from '@tanstack/ai' import type { @@ -6,23 +7,26 @@ import type { ImageGenerateInput, InferGenerationOutputFromReturn, } from '@tanstack/ai-client' -import type { InjectGenerationOptions } from './inject-generation' +import type { + InjectGenerationOptions, + InjectGenerationResult, +} from './inject-generation' export type InjectGenerateImageOptions = Omit< InjectGenerationOptions, - 'onResult' + 'onResult' | 'reconstructResult' > & { onResult?: (result: ImageGenerationResult) => TOutput | null | void } -export interface InjectGenerateImageResult { +export interface InjectGenerateImageResult< + TOutput = ImageGenerationResult, +> extends Omit, 'generate'> { generate: (input: ImageGenerateInput) => Promise result: Signal isLoading: Signal error: Signal status: Signal - stop: () => void - reset: () => void } export function injectGenerateImage( @@ -38,18 +42,14 @@ export function injectGenerateImage( hookName: 'injectGenerateImage', outputKind: 'image' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - injectGeneration({ - ...options, - devtools, - }) - return { - generate: generate as (input: ImageGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + const generation = injectGeneration< + ImageGenerateInput, + ImageGenerationResult, + TTransformed + >({ + ...options, + devtools, + reconstructResult: reconstructImageResult, + }) + return generation } diff --git a/packages/ai-angular/src/inject-generate-speech.ts b/packages/ai-angular/src/inject-generate-speech.ts index ce3e08549..cdc1ff8b8 100644 --- a/packages/ai-angular/src/inject-generate-speech.ts +++ b/packages/ai-angular/src/inject-generate-speech.ts @@ -6,23 +6,27 @@ import type { InferGenerationOutputFromReturn, SpeechGenerateInput, } from '@tanstack/ai-client' -import type { InjectGenerationOptions } from './inject-generation' +import type { + InjectGenerationOptions, + InjectGenerationResult, +} from './inject-generation' export type InjectGenerateSpeechOptions = Omit< InjectGenerationOptions, - 'onResult' + 'onResult' | 'reconstructResult' > & { onResult?: (result: TTSResult) => TOutput | null | void } -export interface InjectGenerateSpeechResult { +export interface InjectGenerateSpeechResult extends Omit< + InjectGenerationResult, + 'generate' +> { generate: (input: SpeechGenerateInput) => Promise result: Signal isLoading: Signal error: Signal status: Signal - stop: () => void - reset: () => void } export function injectGenerateSpeech( @@ -38,18 +42,13 @@ export function injectGenerateSpeech( hookName: 'injectGenerateSpeech', outputKind: 'audio' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - injectGeneration({ - ...options, - devtools, - }) - return { - generate: generate as (input: SpeechGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + const generation = injectGeneration< + SpeechGenerateInput, + TTSResult, + TTransformed + >({ + ...options, + devtools, + }) + return generation } diff --git a/packages/ai-angular/src/inject-generate-video.ts b/packages/ai-angular/src/inject-generate-video.ts index c493e39ee..f799d9a1e 100644 --- a/packages/ai-angular/src/inject-generate-video.ts +++ b/packages/ai-angular/src/inject-generate-video.ts @@ -17,6 +17,9 @@ import type { ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, VideoGenerateInput, VideoGenerateResult, @@ -32,6 +35,21 @@ export interface InjectGenerateVideoOptions { id?: string body?: ReactiveOption> devtools?: AIDevtoolsDisplayOptions + /** + * How this generation persists across reloads. + * - Omit / `false`: ephemeral, in-memory only. + * - `true`: server-driven — on mount the client hydrates the last generation + * for its `threadId` from the server (needs a connection with a + * `hydrateGeneration` handler) and repaints it; it never auto-starts a run. + * - a storage adapter: client-driven — the lightweight snapshot is cached under + * `generation:` as a run streams and read back on mount. Media bytes are + * never stored. + */ + persistence?: boolean | GenerationPersistence + /** Thread id for this generation, stable across reloads. Server-driven mode (`persistence: true`) hydrates the last generation under this key. Falls back to `id`. */ + threadId?: string + /** Explicit resume-snapshot seed for apps that manage storage themselves; skips automatic hydration from `persistence`. Later run events merge into it. */ + initialResumeSnapshot?: GenerationResumeSnapshot onResult?: (result: VideoGenerateResult) => TOutput | null | void onError?: (error: Error) => void onProgress?: (progress: number, message?: string) => void @@ -50,6 +68,7 @@ export interface InjectGenerateVideoResult { status: Signal stop: () => void reset: () => void + resumeState: Signal } // `TTransformed` infers from the `onResult` return position so the callback @@ -79,6 +98,10 @@ export function injectGenerateVideo( const isLoading = signal(false) const error = signal(undefined) const status = signal('idle') + const resumeState = signal( + options.initialResumeSnapshot?.resumeState ?? null, + ) + let disposed = false const bodySource = options.body !== undefined ? toReactive(options.body) : undefined @@ -86,6 +109,13 @@ export function injectGenerateVideo( const baseOptions = { id: clientId, ...(bodySource !== undefined && { body: bodySource() }), + ...(options.persistence !== undefined && { + persistence: options.persistence, + }), + ...(options.threadId !== undefined && { threadId: options.threadId }), + ...(options.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: options.initialResumeSnapshot, + }), devtoolsBridgeFactory: createVideoDevtoolsBridge, devtools: { ...options.devtools, @@ -99,17 +129,42 @@ export function injectGenerateVideo( onResult: ((r: VideoGenerateResult) => options.onResult?.(r)) as ( result: VideoGenerateResult, ) => TOutput | null | void, - onError: (e: Error) => options.onError?.(e), - onProgress: (p: number, m?: string) => options.onProgress?.(p, m), - onChunk: (c: StreamChunk) => options.onChunk?.(c), - onJobCreated: (id: string) => options.onJobCreated?.(id), - onStatusUpdate: (s: VideoStatusInfo) => options.onStatusUpdate?.(s), - onResultChange: (r: TOutput | null) => result.set(r), - onLoadingChange: (l: boolean) => isLoading.set(l), - onErrorChange: (e: Error | undefined) => error.set(e), - onStatusChange: (s: GenerationClientState) => status.set(s), - onJobIdChange: (id: string | null) => jobId.set(id), - onVideoStatusChange: (s: VideoStatusInfo | null) => videoStatus.set(s), + onError: (e: Error) => { + if (!disposed) options.onError?.(e) + }, + onProgress: (p: number, m?: string) => { + if (!disposed) options.onProgress?.(p, m) + }, + onChunk: (c: StreamChunk) => { + if (!disposed) options.onChunk?.(c) + }, + onJobCreated: (id: string) => { + if (!disposed) options.onJobCreated?.(id) + }, + onStatusUpdate: (s: VideoStatusInfo) => { + if (!disposed) options.onStatusUpdate?.(s) + }, + onResultChange: (r: TOutput | null) => { + if (!disposed) result.set(r) + }, + onLoadingChange: (l: boolean) => { + if (!disposed) isLoading.set(l) + }, + onErrorChange: (e: Error | undefined) => { + if (!disposed) error.set(e) + }, + onStatusChange: (s: GenerationClientState) => { + if (!disposed) status.set(s) + }, + onJobIdChange: (id: string | null) => { + if (!disposed) jobId.set(id) + }, + onVideoStatusChange: (s: VideoStatusInfo | null) => { + if (!disposed) videoStatus.set(s) + }, + onResumeStateChange: (rs: GenerationResumeState | null) => { + if (!disposed) resumeState.set(rs) + }, } let client: VideoGenerationClient @@ -132,14 +187,26 @@ export function injectGenerateVideo( if (bodySource) { effect( () => { - client.updateOptions({ body: bodySource() }) + client.updateOptions({ + body: bodySource(), + }) }, { injector }, ) } - afterNextRender(() => client.mountDevtools(), { injector }) - destroyRef.onDestroy(() => client.dispose()) + // Mount devtools only. Generation runs are never auto-started after render — + // persisted state is read-only for display. + afterNextRender( + () => { + client.mountDevtools() + }, + { injector }, + ) + destroyRef.onDestroy(() => { + disposed = true + client.dispose() + }) return { generate: (input: VideoGenerateInput) => client.generate(input), @@ -151,5 +218,6 @@ export function injectGenerateVideo( status: status.asReadonly(), stop: () => client.stop(), reset: () => client.reset(), + resumeState: resumeState.asReadonly(), } } diff --git a/packages/ai-angular/src/inject-generation.ts b/packages/ai-angular/src/inject-generation.ts index 177842e01..54a33682f 100644 --- a/packages/ai-angular/src/inject-generation.ts +++ b/packages/ai-angular/src/inject-generation.ts @@ -18,6 +18,10 @@ import type { GenerationClientOptions, GenerationClientState, GenerationFetcher, + GenerationPersistence, + GenerationRestoredResult, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, } from '@tanstack/ai-client' import type { ReactiveOption } from './internal/to-reactive' @@ -35,6 +39,21 @@ export interface InjectGenerationOptions { body?: ReactiveOption> /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** + * How this generation persists across reloads. + * - Omit / `false`: ephemeral, in-memory only. + * - `true`: server-driven — on mount the client hydrates the last generation + * for its `threadId` from the server (needs a connection with a + * `hydrateGeneration` handler) and repaints it; it never auto-starts a run. + * - a storage adapter: client-driven — the lightweight snapshot is cached under + * `generation:` as a run streams and read back on mount. Media bytes are + * never stored. + */ + persistence?: boolean | GenerationPersistence + /** Thread id for this generation, stable across reloads. Server-driven mode (`persistence: true`) hydrates the last generation under this key. Falls back to `id`. */ + threadId?: string + /** Explicit resume-snapshot seed for apps that manage storage themselves; skips automatic hydration from `persistence`. Later run events merge into it. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when a result is received. Can optionally return a transformed value. * @@ -49,11 +68,27 @@ export interface InjectGenerationOptions { onProgress?: (progress: number, message?: string) => void /** Callback for each stream chunk (connect-based adapter mode only) */ onChunk?: (chunk: StreamChunk) => void + /** + * @internal Rebuild a typed result from a restored snapshot, injected by each + * specialized injectable (image / audio / transcription / summarize). + * Forwarded to the client so a client-store / server-hydrate restore repaints + * `result`. + */ + reconstructResult?: (restored: GenerationRestoredResult) => TResult | null } -export interface InjectGenerationResult { +/** + * Return type for the injectGeneration function. + * + * @template TOutput - The output type (after optional transform) + * @template TInput - The input type accepted by `generate` (defaults to any object) + */ +export interface InjectGenerationResult< + TOutput, + TInput extends Record = Record, +> { /** Trigger a generation request */ - generate: (input: Record) => Promise + generate: (input: TInput) => Promise /** The generation result, or null if not yet generated */ result: Signal /** Whether a generation is currently in progress */ @@ -66,6 +101,8 @@ export interface InjectGenerationResult { stop: () => void /** Clear result, error, and return to idle */ reset: () => void + /** Identity of the in-flight run while one is streaming, or null after it ends */ + resumeState: Signal } // `TTransformed` infers from the `onResult` return position (a covariant @@ -83,7 +120,8 @@ export function injectGeneration< onResult?: (result: TResult) => TTransformed }, ): InjectGenerationResult< - InferGenerationOutputFromReturn + InferGenerationOutputFromReturn, + TInput > { assertInInjectionContext(injectGeneration) @@ -97,6 +135,10 @@ export function injectGeneration< const isLoading = signal(false) const error = signal(undefined) const status = signal('idle') + const resumeState = signal( + options.initialResumeSnapshot?.resumeState ?? null, + ) + let disposed = false const bodySource = options.body !== undefined ? toReactive(options.body) : undefined @@ -104,6 +146,16 @@ export function injectGeneration< const clientOptions: GenerationClientOptions = { id: clientId, ...(bodySource !== undefined && { body: bodySource() }), + ...(options.persistence !== undefined && { + persistence: options.persistence, + }), + ...(options.threadId !== undefined && { threadId: options.threadId }), + ...(options.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: options.initialResumeSnapshot, + }), + ...(options.reconstructResult + ? { reconstructResult: options.reconstructResult } + : {}), devtoolsBridgeFactory: createGenerationDevtoolsBridge, devtools: { ...options.devtools, @@ -116,13 +168,30 @@ export function injectGeneration< onResult: ((r: TResult) => options.onResult?.(r)) as ( result: TResult, ) => TOutput | null | void, - onError: (e: Error) => options.onError?.(e), - onProgress: (p: number, m?: string) => options.onProgress?.(p, m), - onChunk: (c: StreamChunk) => options.onChunk?.(c), - onResultChange: (r: TOutput | null) => result.set(r), - onLoadingChange: (l: boolean) => isLoading.set(l), - onErrorChange: (e: Error | undefined) => error.set(e), - onStatusChange: (s: GenerationClientState) => status.set(s), + onError: (e: Error) => { + if (!disposed) options.onError?.(e) + }, + onProgress: (p: number, m?: string) => { + if (!disposed) options.onProgress?.(p, m) + }, + onChunk: (c: StreamChunk) => { + if (!disposed) options.onChunk?.(c) + }, + onResultChange: (r: TOutput | null) => { + if (!disposed) result.set(r) + }, + onLoadingChange: (l: boolean) => { + if (!disposed) isLoading.set(l) + }, + onErrorChange: (e: Error | undefined) => { + if (!disposed) error.set(e) + }, + onStatusChange: (s: GenerationClientState) => { + if (!disposed) status.set(s) + }, + onResumeStateChange: (rs: GenerationResumeState | null) => { + if (!disposed) resumeState.set(rs) + }, } let client: GenerationClient @@ -145,24 +214,35 @@ export function injectGeneration< if (bodySource) { effect( () => { - client.updateOptions({ body: bodySource() }) + client.updateOptions({ + body: bodySource(), + }) }, { injector }, ) } - afterNextRender(() => client.mountDevtools(), { injector }) - destroyRef.onDestroy(() => client.dispose()) + // Mount devtools only. Generation runs are never auto-started after render — + // persisted state is read-only for display. + afterNextRender( + () => { + client.mountDevtools() + }, + { injector }, + ) + destroyRef.onDestroy(() => { + disposed = true + client.dispose() + }) return { - generate: ((input: TInput) => client.generate(input)) as ( - input: Record, - ) => Promise, + generate: (input: TInput) => client.generate(input), result: result.asReadonly(), isLoading: isLoading.asReadonly(), error: error.asReadonly(), status: status.asReadonly(), stop: () => client.stop(), reset: () => client.reset(), + resumeState: resumeState.asReadonly(), } } diff --git a/packages/ai-angular/src/inject-summarize.ts b/packages/ai-angular/src/inject-summarize.ts index fa41a76fa..2edb5e165 100644 --- a/packages/ai-angular/src/inject-summarize.ts +++ b/packages/ai-angular/src/inject-summarize.ts @@ -1,4 +1,5 @@ import { injectGeneration } from './inject-generation' +import { reconstructSummarizeResult } from '@tanstack/ai-client' import type { Signal } from '@angular/core' import type { SummarizationResult } from '@tanstack/ai' import type { @@ -6,23 +7,26 @@ import type { InferGenerationOutputFromReturn, SummarizeGenerateInput, } from '@tanstack/ai-client' -import type { InjectGenerationOptions } from './inject-generation' +import type { + InjectGenerationOptions, + InjectGenerationResult, +} from './inject-generation' export type InjectSummarizeOptions = Omit< InjectGenerationOptions, - 'onResult' + 'onResult' | 'reconstructResult' > & { onResult?: (result: SummarizationResult) => TOutput | null | void } -export interface InjectSummarizeResult { +export interface InjectSummarizeResult< + TOutput = SummarizationResult, +> extends Omit, 'generate'> { generate: (input: SummarizeGenerateInput) => Promise result: Signal isLoading: Signal error: Signal status: Signal - stop: () => void - reset: () => void } export function injectSummarize( @@ -38,20 +42,14 @@ export function injectSummarize( hookName: 'injectSummarize', outputKind: 'text' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - injectGeneration( - { - ...options, - devtools, - }, - ) - return { - generate: generate as (input: SummarizeGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + const generation = injectGeneration< + SummarizeGenerateInput, + SummarizationResult, + TTransformed + >({ + ...options, + devtools, + reconstructResult: reconstructSummarizeResult, + }) + return generation } diff --git a/packages/ai-angular/src/inject-transcription.ts b/packages/ai-angular/src/inject-transcription.ts index ef2066ad2..fe2e240a4 100644 --- a/packages/ai-angular/src/inject-transcription.ts +++ b/packages/ai-angular/src/inject-transcription.ts @@ -1,4 +1,5 @@ import { injectGeneration } from './inject-generation' +import { reconstructTranscriptionResult } from '@tanstack/ai-client' import type { Signal } from '@angular/core' import type { TranscriptionResult } from '@tanstack/ai' import type { @@ -6,7 +7,10 @@ import type { InferGenerationOutputFromReturn, TranscriptionGenerateInput, } from '@tanstack/ai-client' -import type { InjectGenerationOptions } from './inject-generation' +import type { + InjectGenerationOptions, + InjectGenerationResult, +} from './inject-generation' export type InjectTranscriptionOptions = Omit< InjectGenerationOptions< @@ -14,19 +18,19 @@ export type InjectTranscriptionOptions = Omit< TranscriptionResult, TOutput >, - 'onResult' + 'onResult' | 'reconstructResult' > & { onResult?: (result: TranscriptionResult) => TOutput | null | void } -export interface InjectTranscriptionResult { +export interface InjectTranscriptionResult< + TOutput = TranscriptionResult, +> extends Omit, 'generate'> { generate: (input: TranscriptionGenerateInput) => Promise result: Signal isLoading: Signal error: Signal status: Signal - stop: () => void - reset: () => void } export function injectTranscription( @@ -42,22 +46,14 @@ export function injectTranscription( hookName: 'injectTranscription', outputKind: 'text' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - injectGeneration< - TranscriptionGenerateInput, - TranscriptionResult, - TTransformed - >({ - ...options, - devtools, - }) - return { - generate: generate as (input: TranscriptionGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + const generation = injectGeneration< + TranscriptionGenerateInput, + TranscriptionResult, + TTransformed + >({ + ...options, + devtools, + reconstructResult: reconstructTranscriptionResult, + }) + return generation } diff --git a/packages/ai-angular/tests/inject-generation.test.ts b/packages/ai-angular/tests/inject-generation.test.ts index 1ae52833e..2f9068506 100644 --- a/packages/ai-angular/tests/inject-generation.test.ts +++ b/packages/ai-angular/tests/inject-generation.test.ts @@ -6,6 +6,15 @@ import { } from '@angular/platform-browser-dynamic/testing' import { describe, expect, it, vi } from 'vitest' import { injectGeneration } from '../src/inject-generation' +import { injectGenerateVideo } from '../src/inject-generate-video' +import { injectGenerateImage } from '../src/inject-generate-image' +import type { PersistedArtifactRef, StreamChunk } from '@tanstack/ai' +import type { + ConnectConnectionAdapter, + GenerationPersistence, + GenerationResumeSnapshot, + RunAgentInputContext, +} from '@tanstack/ai-client' // Ensure TestBed is initialized in this module's scope, regardless of whether // the setup file's initialization was in a different module context (possible @@ -34,9 +43,109 @@ function renderInjectGeneration(options: any) { return fixture.componentInstance.gen }, flush: () => fixture.detectChanges(), + destroy: () => fixture.destroy(), } } +function renderInjectGenerateVideo(options: any) { + @Component({ standalone: true, template: '' }) + class Host { + gen = injectGenerateVideo(options) + } + const fixture = TestBed.createComponent(Host) + fixture.detectChanges() + return { + get result() { + return fixture.componentInstance.gen + }, + flush: () => fixture.detectChanges(), + destroy: () => fixture.destroy(), + } +} + +function renderInjectGenerateImage(options: any) { + @Component({ standalone: true, template: '' }) + class Host { + gen = injectGenerateImage(options) + } + const fixture = TestBed.createComponent(Host) + fixture.detectChanges() + return { + get result() { + return fixture.componentInstance.gen + }, + flush: () => fixture.detectChanges(), + destroy: () => fixture.destroy(), + } +} + +const videoResumeSnapshot: GenerationResumeSnapshot = { + resumeState: { + threadId: 'thread-resume', + runId: 'run-resume', + }, + status: 'running', +} + +/** + * Storage adapter backed by a Map, so a test can seed a persisted record and + * then assert on what the client read, wrote, and removed. + */ +function createMapPersistence( + seed?: Record, +): { + persistence: GenerationPersistence + store: Map + getItem: ReturnType + setItem: ReturnType + removeItem: ReturnType +} { + const store = new Map( + Object.entries(seed ?? {}), + ) + const getItem = vi.fn((key: string) => store.get(key) ?? null) + const setItem = vi.fn((key: string, value: GenerationResumeSnapshot) => { + store.set(key, value) + }) + const removeItem = vi.fn((key: string) => { + store.delete(key) + }) + return { + persistence: { getItem, setItem, removeItem }, + store, + getItem, + setItem, + removeItem, + } +} + +// Hydration and snapshot removal both run through awaited promise chains, so +// drain the microtask queue rather than awaiting a single tick. +async function flushPromises(): Promise { + for (let i = 0; i < 5; i++) { + await Promise.resolve() + } +} + +function createRunContextCaptureAdapter(chunks: Array): { + adapter: ConnectConnectionAdapter + connect: ReturnType + runContexts: Array +} { + const runContexts: Array = [] + const connect = vi.fn() + const adapter: ConnectConnectionAdapter = { + async *connect(_messages, _data, _signal, runContext) { + connect(runContext) + runContexts.push(runContext) + for (const chunk of chunks) { + yield chunk + } + }, + } + return { adapter, connect, runContexts } +} + describe('injectGeneration', () => { it('initializes idle with a fetcher and generates a result', async () => { const fetcher = vi.fn(async () => ({ value: 42 })) @@ -68,4 +177,198 @@ describe('injectGeneration', () => { expect(result.result()).toEqual({ playable: true }) expect(result.status()).toBe('success') }) + + it('does not auto-fire a generation after render from a persisted running snapshot', async () => { + // Regression guard for the removed generation resume surface. + const snapshot: GenerationResumeSnapshot = { + resumeState: { threadId: 'thread-resume', runId: 'run-resume' }, + status: 'running', + } + const { adapter, connect } = createRunContextCaptureAdapter([]) + const getItem = vi.fn(() => snapshot) + const { result } = renderInjectGeneration({ + id: 'no-auto-fire', + connection: adapter, + persistence: { getItem, setItem: vi.fn(), removeItem: vi.fn() }, + initialResumeSnapshot: snapshot, + }) + + await Promise.resolve() + + expect(connect).not.toHaveBeenCalled() + expect(getItem).not.toHaveBeenCalled() + expect(result.isLoading()).toBe(false) + expect(result.status()).toBe('idle') + // The persisted snapshot remains exposed as read-only state. + expect(result.resumeState()).toEqual(snapshot.resumeState) + }) + + it('hydrates a persisted snapshot from storage on construction', async () => { + const { adapter, connect } = createRunContextCaptureAdapter([]) + const { persistence, getItem } = createMapPersistence({ + 'generation:hydrate-me': { + resumeState: { threadId: 'thread-stored', runId: 'run-stored' }, + status: 'running', + }, + }) + const { result } = renderInjectGeneration({ + id: 'hydrate-me', + connection: adapter, + persistence, + }) + + await flushPromises() + + expect(getItem).toHaveBeenCalledTimes(1) + expect(getItem).toHaveBeenCalledWith('generation:hydrate-me') + // Hydration only surfaces state; it never restarts the run. + expect(connect).not.toHaveBeenCalled() + // A restored running run repaints `status` to `generating` but never + // auto-tails, and its in-flight identity is exposed as `resumeState`. + expect(result.status()).toBe('generating') + expect(result.isLoading()).toBe(false) + expect(result.resumeState()).toEqual({ + threadId: 'thread-stored', + runId: 'run-stored', + }) + }) + + it('clears resume state and removes the persisted record on reset', async () => { + const snapshot: GenerationResumeSnapshot = { + resumeState: { threadId: 'thread-reset', runId: 'run-reset' }, + status: 'running', + } + const { adapter } = createRunContextCaptureAdapter([]) + const { persistence, removeItem, store } = createMapPersistence({ + 'generation:reset-me': snapshot, + }) + const { result } = renderInjectGeneration({ + id: 'reset-me', + connection: adapter, + persistence, + initialResumeSnapshot: snapshot, + }) + + // The seeded in-flight identity is exposed as read-only `resumeState`. + expect(result.resumeState()).toEqual(snapshot.resumeState) + + result.reset() + await flushPromises() + + expect(result.resumeState()).toBeNull() + expect(removeItem).toHaveBeenCalledWith('generation:reset-me') + expect(store.has('generation:reset-me')).toBe(false) + }) +}) + +describe('injectGenerateVideo', () => { + it('does not auto-fire a video generation after render from a persisted running snapshot', async () => { + // Regression guard for the removed generation resume surface (video). + const { adapter, connect } = createRunContextCaptureAdapter([]) + const getItem = vi.fn(() => videoResumeSnapshot) + const { result } = renderInjectGenerateVideo({ + id: 'video-no-auto-fire', + connection: adapter, + persistence: { getItem, setItem: vi.fn(), removeItem: vi.fn() }, + initialResumeSnapshot: videoResumeSnapshot, + }) + + await Promise.resolve() + + expect(connect).not.toHaveBeenCalled() + expect(getItem).not.toHaveBeenCalled() + expect(result.isLoading()).toBe(false) + expect(result.status()).toBe('idle') + // The seeded in-flight identity is exposed as read-only `resumeState`. + expect(result.resumeState()).toEqual(videoResumeSnapshot.resumeState) + }) + + it('hydrates from storage and clears the persisted record on reset', async () => { + const { adapter, connect } = createRunContextCaptureAdapter([]) + const { persistence, getItem, removeItem, store } = createMapPersistence({ + 'generation:video-hydrate': videoResumeSnapshot, + }) + const { result } = renderInjectGenerateVideo({ + id: 'video-hydrate', + connection: adapter, + persistence, + }) + + await flushPromises() + + expect(getItem).toHaveBeenCalledTimes(1) + expect(getItem).toHaveBeenCalledWith('generation:video-hydrate') + expect(connect).not.toHaveBeenCalled() + // A restored running run repaints `status` to `generating` and exposes its + // in-flight identity as `resumeState`, but never auto-tails. + expect(result.status()).toBe('generating') + expect(result.isLoading()).toBe(false) + expect(result.resumeState()).toEqual(videoResumeSnapshot.resumeState) + + result.reset() + await flushPromises() + + expect(result.resumeState()).toBeNull() + expect(removeItem).toHaveBeenCalledWith('generation:video-hydrate') + expect(store.has('generation:video-hydrate')).toBe(false) + }) +}) + +describe('injectGenerateImage', () => { + it('restores a completed image result from a durable artifact url', async () => { + // injectGenerateImage injects `reconstructImageResult`, so a restored + // complete snapshot repaints `result` with the durable serve url — as if the + // run had just finished. + const artifact: PersistedArtifactRef = { + role: 'output', + artifactId: 'artifact-image-1', + threadId: 'thread-img', + runId: 'run-img', + name: 'image.png', + mimeType: 'image/png', + size: 2048, + createdAt: '2026-07-06T00:00:00.000Z', + url: '/api/artifacts/artifact-image-1', + source: { + activity: 'image', + path: 'runs/run-img/image.png', + provider: 'test', + model: 'test-image', + mediaType: 'image', + }, + } + const { adapter, connect } = createRunContextCaptureAdapter([]) + const { persistence, getItem } = createMapPersistence({ + 'generation:img-hydrate': { + resumeState: null, + status: 'complete', + activity: 'image', + result: { + id: 'img-restored', + model: 'test-image', + artifacts: [artifact], + }, + }, + }) + const { result } = renderInjectGenerateImage({ + id: 'img-hydrate', + connection: adapter, + persistence, + }) + + await flushPromises() + + expect(getItem).toHaveBeenCalledWith('generation:img-hydrate') + + // A completed snapshot repaints the normal `status` field to `success`. + expect(result.status()).toBe('success') + expect(connect).not.toHaveBeenCalled() + expect(result.result()).toEqual({ + id: 'img-restored', + model: 'test-image', + images: [{ url: '/api/artifacts/artifact-image-1' }], + artifacts: [artifact], + }) + expect(result.resumeState()).toBeNull() + }) }) diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index 260f17526..20979d3a1 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -453,6 +453,40 @@ async function fetchThreadHydration( } } +/** + * GET the hydration endpoint for a generation thread and parse its JSON + * `{ resumeSnapshot, activeRun }` body. Mirrors {@link fetchThreadHydration} for + * the generation clients: keyed on the stable thread id, it returns the last + * generation's resume snapshot (re-validated client-side before adoption) and — + * if a run is still generating — a cursor. Shared by every fetch/XHR adapter. + */ +async function fetchGenerationHydration( + fetchClient: typeof globalThis.fetch, + url: string, + headers: Record, + credentials: RequestCredentials, + threadId: string, +): Promise { + const response = await fetchClient(withSearchParams(url, { threadId }), { + method: 'GET', + headers: { Accept: 'application/json', ...headers }, + credentials, + }) + assertResponseOk(response) + const data = (await response.json()) as { + resumeSnapshot?: GenerationHydrationResult['resumeSnapshot'] + activeRun?: { runId?: unknown } | null + } + const activeRun = + data.activeRun && typeof data.activeRun.runId === 'string' + ? { runId: data.activeRun.runId } + : null + return { + resumeSnapshot: data.resumeSnapshot ?? null, + activeRun, + } +} + /** Yield SSE stream events (chunk + offset) from a fetch Response body. */ async function* responseToSSEEvents( response: Response, @@ -707,6 +741,35 @@ export interface ConnectConnectionAdapter { abortSignal?: AbortSignal, runContext?: RunAgentInputContext, ) => AsyncIterable + /** + * Fetch server-driven hydration for a generation `threadId`: the last + * generation's resume snapshot, plus a cursor to a run still generating if + * one exists. The generation client calls this itself on mount when + * `persistence: true` (no loader/prop) and repaints the snapshot — it never + * auto-starts a run. Read-only JSON GET (`?threadId`), so it is + * transport-agnostic. Optional and feature-detected exactly like the chat + * `hydrate` handler. + */ + hydrateGeneration?: (threadId: string) => Promise +} + +/** + * Server-resolved hydration for a generation thread. `resumeSnapshot` is the + * last generation's lightweight snapshot (validated client-side before it is + * adopted); `activeRun` is a cursor to a run still generating for the thread + * (or `null`). Structurally matches `@tanstack/ai-persistence`'s + * `reconstructGeneration` response — the client never imports that package. + */ +export interface GenerationHydrationResult { + resumeSnapshot: { + schemaVersion?: 1 + resumeState: { threadId?: string; runId: string } | null + status: 'idle' | 'running' | 'complete' | 'error' + result?: unknown + error?: { message: string; code?: string } + activity?: string + } | null + activeRun: { runId: string } | null } /** @@ -1197,6 +1260,18 @@ export function fetchServerSentEvents( threadId, ) }, + async hydrateGeneration(threadId) { + const resolvedUrl = typeof url === 'function' ? url() : url + const resolvedOptions = + typeof options === 'function' ? await options() : options + return fetchGenerationHydration( + resolvedOptions.fetchClient ?? fetch, + resolvedUrl, + mergeHeaders(resolvedOptions.headers), + resolvedOptions.credentials || 'same-origin', + threadId, + ) + }, } } @@ -1340,6 +1415,18 @@ export function fetchHttpStream( threadId, ) }, + async hydrateGeneration(threadId) { + const resolvedUrl = typeof url === 'function' ? url() : url + const resolvedOptions = + typeof options === 'function' ? await options() : options + return fetchGenerationHydration( + resolvedOptions.fetchClient ?? fetch, + resolvedUrl, + mergeHeaders(resolvedOptions.headers), + resolvedOptions.credentials || 'same-origin', + threadId, + ) + }, } } @@ -1666,6 +1753,19 @@ export function xhrServerSentEvents( threadId, ) }, + async hydrateGeneration(threadId) { + const resolvedUrl = typeof url === 'function' ? url() : url + const resolvedOptions = await resolveXhrConnectionOptions(options) + // Hydration is a non-streaming JSON GET, so fetch is fine even for the + // XHR-backed streaming adapter. + return fetchGenerationHydration( + fetch, + resolvedUrl, + mergeHeaders(resolvedOptions.headers), + resolvedOptions.withCredentials ? 'include' : 'same-origin', + threadId, + ) + }, } } @@ -1736,6 +1836,19 @@ export function xhrHttpStream( threadId, ) }, + async hydrateGeneration(threadId) { + const resolvedUrl = typeof url === 'function' ? url() : url + const resolvedOptions = await resolveXhrConnectionOptions(options) + // Hydration is a non-streaming JSON GET, so fetch is fine even for the + // XHR-backed streaming adapter. + return fetchGenerationHydration( + fetch, + resolvedUrl, + mergeHeaders(resolvedOptions.headers), + resolvedOptions.withCredentials ? 'include' : 'same-origin', + threadId, + ) + }, } } diff --git a/packages/ai-client/src/generation-client.ts b/packages/ai-client/src/generation-client.ts index 38bc56a0b..740b3f426 100644 --- a/packages/ai-client/src/generation-client.ts +++ b/packages/ai-client/src/generation-client.ts @@ -1,9 +1,16 @@ -import { GENERATION_EVENTS } from './generation-types' +import { + GENERATION_EVENTS, + clientStateFromResumeStatus, + createGenerationResultSnapshot, + parseGenerationResumeSnapshot, + updateGenerationResumeSnapshot, +} from './generation-types' import { createNoOpGenerationDevtoolsBridge } from './devtools-noop' import { parseSSEResponse } from './sse-parser' import type { StreamChunk } from '@tanstack/ai/client' import type { ConnectConnectionAdapter, + GenerationHydrationResult, RunAgentInputContext, } from './connection-adapters' import type { @@ -16,6 +23,10 @@ import type { GenerationClientOptions, GenerationClientState, GenerationFetcher, + GenerationResumeSnapshot, + GenerationRestoredResult, + GenerationResumeState, + GenerationPersistence, } from './generation-types' /** @@ -33,6 +44,15 @@ interface GenerationCallbacks { onLoadingChange?: ((isLoading: boolean) => void) | undefined onErrorChange?: ((error: Error | undefined) => void) | undefined onStatusChange?: ((status: GenerationClientState) => void) | undefined + onResumeSnapshotChange?: + | ((snapshot: GenerationResumeSnapshot | undefined) => void) + | undefined + onResumeStateChange?: + | ((resumeState: GenerationResumeState | null) => void) + | undefined + reconstructResult?: + | ((restored: GenerationRestoredResult) => TResult | null) + | undefined } /** @@ -81,6 +101,10 @@ export class GenerationClient< private readonly devtoolsMetadata: AIDevtoolsClientMetadata private readonly devtoolsBridge: GenerationDevtoolsBridge private readonly threadId: string + private readonly resumePersistence: GenerationPersistence | undefined + // Server-driven mode (`persistence: true`): no local snapshot store; on mount + // the client hydrates the last generation for `threadId` from the server. + private readonly serverDriven: boolean = false private body: Record private result: TOutput | null = null private input: TInput | null = null @@ -88,9 +112,15 @@ export class GenerationClient< private isLoading = false private error: Error | undefined = undefined private status: GenerationClientState = 'idle' + private resumeSnapshot: GenerationResumeSnapshot | undefined + private resumeSnapshotPersistenceQueue: Promise = Promise.resolve() + private resumeSnapshotHydration: Promise | undefined + private queuedSnapshotSignature: string | undefined + private resumePersistenceError: Error | undefined = undefined private abortController: AbortController | null = null private readonly callbacksRef: GenerationCallbacks private devtoolsMounted = false + private disposed = false constructor( options: GenerationClientOptions & @@ -103,10 +133,23 @@ export class GenerationClient< ), ) { this.uniqueId = options.id ?? this.generateUniqueId('generation') - this.threadId = this.uniqueId + // The wire/hydration thread key. Server-driven mode needs a stable key, so + // prefer an explicit `threadId`, then `id`, then a generated id. + this.threadId = options.threadId ?? this.uniqueId this.connection = options.connection this.fetcher = options.fetcher this.body = options.body ?? {} + // `persistence` is `false`/omitted (ephemeral), `true` (server-driven: cache + // nothing, hydrate the last generation for `threadId` on mount), or a storage + // adapter (client-driven: cache the lightweight resume snapshot locally). + // Only the server-driven mode leaves `resumePersistence` undefined AND turns + // on mount hydration. + if (options.persistence === true) { + this.serverDriven = true + } else if (options.persistence) { + this.resumePersistence = options.persistence + } + this.resumeSnapshot = options.initialResumeSnapshot this.callbacksRef = { onResult: options.onResult, @@ -117,12 +160,32 @@ export class GenerationClient< onLoadingChange: options.onLoadingChange, onErrorChange: options.onErrorChange, onStatusChange: options.onStatusChange, + onResumeSnapshotChange: options.onResumeSnapshotChange, + onResumeStateChange: options.onResumeStateChange, + reconstructResult: options.reconstructResult, } this.devtoolsMetadata = this.createDevtoolsMetadata(options.devtools) this.devtoolsBridge = ( options.devtoolsBridgeFactory ?? createNoOpGenerationDevtoolsBridge )(this.buildDevtoolsBridgeOptions()) + + // After callbacksRef is assigned: hydration may fire + // onResumeSnapshotChange synchronously if an adapter resolves sync. + this.maybeHydrateResumeSnapshot() + + // Server-driven (`persistence: true`): the client caches nothing locally and + // re-hydrates the last generation for its stable threadId from the server on + // mount. Best-effort and non-blocking; it never auto-starts a run. + if (this.serverDriven && this.connection?.hydrateGeneration) { + this.hydrateFromServer() + } else if (this.serverDriven) { + // `persistence: true` without a hydrate-capable connection can never + // restore anything — warn rather than silently no-op. + console.warn( + '[TanStack AI] `persistence: true` (server-driven) needs a connection that implements `hydrateGeneration` (e.g. `fetchServerSentEvents` / `fetchHttpStream`). With a plain `fetcher` or no connection, nothing is persisted or restored.', + ) + } } private buildDevtoolsBridgeOptions(): GenerationDevtoolsBridgeOptions { @@ -143,6 +206,12 @@ export class GenerationClient< } mountDevtools(): void { + // Mounting revives a disposed client. Framework hooks call this from + // their mount effect, so a dispose → remount cycle (e.g. React + // StrictMode's mount → cleanup → mount replay against the same memoized + // client) leaves the client usable again. + this.disposed = false + this.maybeHydrateResumeSnapshot() if (this.devtoolsMounted) { return } @@ -158,8 +227,9 @@ export class GenerationClient< * while already generating will be a no-op. */ async generate(input: TInput): Promise { - this.mountDevtools() + if (this.disposed) return if (this.isLoading) return + this.mountDevtools() this.input = input this.progress = null @@ -179,11 +249,16 @@ export class GenerationClient< if (signal.aborted) return if (result instanceof Response) { // Server function returned SSE Response — parse stream - await this.processStream(parseSSEResponse(result, signal), runId) + await this.processStream( + parseSSEResponse(result, signal), + runId, + signal, + ) } else { this.devtoolsBridge.ensureRunStarted(runId) this.setResult(result) this.setStatus('success') + this.completePlainFetcherResumeSnapshot(result) } } else if (this.connection) { // Streaming adapter path @@ -194,7 +269,7 @@ export class GenerationClient< signal, this.createRunContext(runId), ) - await this.processStream(stream, runId) + await this.processStream(stream, runId, signal) } else { throw new Error( 'GenerationClient requires either a connection or fetcher option', @@ -217,6 +292,7 @@ export class GenerationClient< const error = err instanceof Error ? err : new Error(String(err)) this.setError(error) this.setStatus('error') + this.recordResumeSnapshotError(error) this.devtoolsBridge.finishRun( this.devtoolsBridge.getActiveRunId() ?? runId, 'run:errored', @@ -225,8 +301,10 @@ export class GenerationClient< ) this.callbacksRef.onError?.(error) } finally { - this.abortController = null - this.setIsLoading(false) + if (this.abortController === abortController) { + this.abortController = null + this.setIsLoading(false) + } } } @@ -236,13 +314,15 @@ export class GenerationClient< private async processStream( source: AsyncIterable, fallbackRunId: string, + signal: AbortSignal, ): Promise { let streamRunId: string | undefined for await (const chunk of source) { - if (this.abortController?.signal.aborted) break + if (signal.aborted) break this.callbacksRef.onChunk?.(chunk) + this.observeResumeSnapshot(chunk) const chunkRunId = 'runId' in chunk && typeof chunk.runId === 'string' ? chunk.runId @@ -307,10 +387,22 @@ export class GenerationClient< this.devtoolsBridge.finishRun(runId, 'run:cancelled', 'cancelled') } } + // A stopped run is no longer resumable. Without this, storage keeps a + // `running` snapshot forever and a reload would chase a dead run. + if (this.resumeSnapshot && this.resumeSnapshot.status === 'running') { + this.resumeSnapshot = { + ...this.resumeSnapshot, + resumeState: null, + status: 'idle', + } + this.notifyResumeSnapshotChanged() + void this.persistResumeSnapshot(this.resumeSnapshot) + } } /** - * Clear the result, error, and return to idle state. + * Clear the result, error, and return to idle state. Also clears the + * resume snapshot, removing any persisted record for this client id. */ reset(): void { this.stop() @@ -320,6 +412,7 @@ export class GenerationClient< this.devtoolsBridge.resetRuns() this.setError(undefined) this.setStatus('idle') + this.clearResumeSnapshot() this.devtoolsBridge.emitState() } @@ -352,6 +445,7 @@ export class GenerationClient< } dispose(): void { + this.disposed = true this.stop() this.devtoolsBridge.dispose() this.devtoolsMounted = false @@ -373,10 +467,41 @@ export class GenerationClient< return this.error } + getResumePersistenceError(): Error | undefined { + return this.resumePersistenceError + } + getStatus(): GenerationClientState { return this.status } + getResumeSnapshot(): GenerationResumeSnapshot | undefined { + return this.resumeSnapshot + ? { + ...this.resumeSnapshot, + ...(this.resumeSnapshot.pendingArtifacts + ? { pendingArtifacts: [...this.resumeSnapshot.pendingArtifacts] } + : {}), + ...(this.resumeSnapshot.result + ? { + result: { + ...this.resumeSnapshot.result, + ...(this.resumeSnapshot.result.artifacts + ? { artifacts: [...this.resumeSnapshot.result.artifacts] } + : {}), + }, + } + : {}), + ...(this.resumeSnapshot.error + ? { error: { ...this.resumeSnapshot.error } } + : {}), + ...(this.resumeSnapshot.lastEvent + ? { lastEvent: { ...this.resumeSnapshot.lastEvent } } + : {}), + } + : undefined + } + // =========================== // Private state setters // =========================== @@ -466,6 +591,303 @@ export class GenerationClient< runId, } } + + private observeResumeSnapshot(chunk: StreamChunk): void { + this.resumeSnapshot = updateGenerationResumeSnapshot( + this.resumeSnapshot, + chunk, + ) + this.notifyResumeSnapshotChanged() + void this.persistResumeSnapshot(this.resumeSnapshot) + } + + /** + * Notify the (internal) snapshot listener AND emit the public resume state. + * The snapshot stays internal (persistence + devtools); the hook consumes + * `resumeState`, mirroring the chat client. + */ + private notifyResumeSnapshotChanged(): void { + this.callbacksRef.onResumeSnapshotChange?.(this.resumeSnapshot) + this.emitResumeState() + } + + /** + * Derive the public `resumeState` from the internal snapshot: the in-flight + * run identity, with any in-flight artifact refs folded under it. `null` once + * no run is in flight. + */ + private emitResumeState(): void { + const snapshot = this.resumeSnapshot + const state = snapshot?.resumeState + const resumeState: GenerationResumeState | null = state + ? { + ...state, + ...(snapshot?.pendingArtifacts && snapshot.pendingArtifacts.length > 0 + ? { pendingArtifacts: [...snapshot.pendingArtifacts] } + : {}), + } + : null + this.callbacksRef.onResumeStateChange?.(resumeState) + } + + /** + * Repaint the normal fields from a restored snapshot (client store or server + * hydrate), so a reload presents the run in `result` / `status` / `error` + * exactly as a just-finished run would, never a bolt-on snapshot object. + * `isLoading` stays false: the client never auto-tails a restored run. The + * snapshot is not re-persisted here (it came from storage / the server). + */ + private repaintFromSnapshot(snapshot: GenerationResumeSnapshot): void { + this.resumeSnapshot = snapshot + this.notifyResumeSnapshotChanged() + this.setStatus(clientStateFromResumeStatus(snapshot.status)) + this.setError( + snapshot.error + ? Object.assign( + new Error(snapshot.error.message), + snapshot.error.code ? { code: snapshot.error.code } : {}, + ) + : undefined, + ) + const restored = this.reconstructRestoredResult(snapshot) + if (restored !== null) this.setResult(restored) + } + + /** + * Build the restorable result shape from the snapshot and hand it to the + * per-activity `reconstructResult` mapper (injected by the specialized + * client/hook, which knows the concrete result type). Returns `null` when no + * mapper is set or it declines — then `result` simply stays null. + */ + private reconstructRestoredResult( + snapshot: GenerationResumeSnapshot, + ): TResult | null { + const build = this.callbacksRef.reconstructResult + if (!build) return null + const result = snapshot.result + const restored: GenerationRestoredResult = { + ...(result?.id !== undefined ? { id: result.id } : {}), + ...(result?.model !== undefined ? { model: result.model } : {}), + ...(result?.status !== undefined ? { status: result.status } : {}), + ...(result?.jobId !== undefined ? { jobId: result.jobId } : {}), + ...(result?.expiresAt !== undefined + ? { expiresAt: result.expiresAt } + : {}), + ...(result?.text !== undefined ? { text: result.text } : {}), + ...(result?.usage !== undefined ? { usage: result.usage } : {}), + ...(snapshot.activity !== undefined + ? { activity: snapshot.activity } + : {}), + artifacts: result?.artifacts ?? [], + } + return build(restored) + } + + /** + * The plain (non-Response) fetcher path never observes stream chunks, so + * the terminal snapshot is built here from the fetcher's own result. A + * stale `error` from a previous run is intentionally dropped — this run + * succeeded. + */ + private completePlainFetcherResumeSnapshot(rawResult: unknown): void { + const previous = this.resumeSnapshot + const result = createGenerationResultSnapshot(rawResult) + this.resumeSnapshot = { + schemaVersion: 1, + resumeState: null, + status: 'complete', + ...(previous?.activity ? { activity: previous.activity } : {}), + ...(previous?.pendingArtifacts && previous.pendingArtifacts.length > 0 + ? { pendingArtifacts: [...previous.pendingArtifacts] } + : {}), + ...(result + ? { result } + : previous?.result + ? { result: { ...previous.result } } + : {}), + } + this.notifyResumeSnapshotChanged() + void this.persistResumeSnapshot(this.resumeSnapshot) + } + + /** + * Records a transport-level failure (network drop, throwing callback) in + * the snapshot. Without this, only a server-emitted RUN_ERROR chunk would + * mark the snapshot `error`, leaving a persisted record that claims the + * run is still in flight. + */ + private recordResumeSnapshotError(error: Error): void { + if (this.resumeSnapshot?.status === 'error') return + if (!this.resumeSnapshot && !this.resumePersistence) return + const previous = this.resumeSnapshot + this.resumeSnapshot = { + schemaVersion: 1, + resumeState: null, + status: 'error', + ...(previous?.activity ? { activity: previous.activity } : {}), + ...(previous?.pendingArtifacts && previous.pendingArtifacts.length > 0 + ? { pendingArtifacts: [...previous.pendingArtifacts] } + : {}), + ...(previous?.result ? { result: { ...previous.result } } : {}), + error: { message: error.message }, + } + this.notifyResumeSnapshotChanged() + void this.persistResumeSnapshot(this.resumeSnapshot) + } + + private clearResumeSnapshot(): void { + this.resumeSnapshot = undefined + this.queuedSnapshotSignature = undefined + this.notifyResumeSnapshotChanged() + if (!this.resumePersistence) { + return + } + this.resumeSnapshotPersistenceQueue = + this.resumeSnapshotPersistenceQueue.then( + () => this.removePersistedResumeSnapshot(), + () => this.removePersistedResumeSnapshot(), + ) + } + + /** + * Storage key for this client's snapshot. The `generation:` segment keeps + * a generation client and a chat client that share an id (and a storage + * adapter with the default key prefix) from overwriting each other. + */ + private get resumeSnapshotKey(): string { + return `generation:${this.uniqueId}` + } + + private maybeHydrateResumeSnapshot(): void { + if (!this.resumePersistence || this.resumeSnapshotHydration) return + // An explicit `initialResumeSnapshot` seed takes precedence over storage. + if (this.resumeSnapshot) return + this.resumeSnapshotHydration = this.hydrateResumeSnapshot() + } + + private async hydrateResumeSnapshot(): Promise { + let stored: unknown + try { + stored = await this.resumePersistence?.getItem(this.resumeSnapshotKey) + } catch (error) { + // A corrupt record (e.g. truncated JSON) or unavailable storage must + // not break construction; the app just starts without a snapshot. + console.warn( + '[TanStack AI] Failed to read persisted generation resume snapshot', + error, + ) + return + } + if (stored === null || stored === undefined) return + const snapshot = parseGenerationResumeSnapshot(stored) + if (!snapshot) return + // Live state wins: adopt the stored snapshot only if nothing has been + // observed since construction. + if (this.resumeSnapshot || this.isLoading || this.status !== 'idle') return + this.repaintFromSnapshot(snapshot) + } + + /** + * Server-driven mount hydration (`persistence: true`). The client holds no + * local snapshot; on mount it asks the server — keyed by the stable threadId — + * for the last generation's resume snapshot, validates it, and repaints it. It + * never auto-starts a run. Best-effort and non-blocking: a failure leaves the + * client empty rather than throwing, and a `generate()` that starts first owns + * the client (hydration then backs off, mirroring the chat client). + */ + private hydrateFromServer(): void { + const hydrate = this.connection?.hydrateGeneration + if (!hydrate) return + // A send that already started owns the client; don't stomp it. + if (this.resumeSnapshot || this.isLoading || this.status !== 'idle') return + void (async () => { + let res: GenerationHydrationResult + try { + res = await hydrate(this.threadId) + } catch { + return + } + if (!res.resumeSnapshot) return + const snapshot = parseGenerationResumeSnapshot(res.resumeSnapshot) + if (!snapshot) return + // Re-check: a send may have started while the fetch was in flight. + if (this.resumeSnapshot || this.isLoading || this.status !== 'idle') + return + this.repaintFromSnapshot(snapshot) + })() + } + + private async persistResumeSnapshot( + snapshot: GenerationResumeSnapshot, + ): Promise { + if (!this.resumePersistence) { + return + } + + // Skip writes that only differ in `lastEvent` — for a long streaming run + // this collapses hundreds of per-chunk writes into the handful where the + // snapshot materially changed. (`lastEvent` in storage is best-effort; + // hydration drops it anyway.) + const signature = resumeSnapshotSignature(snapshot) + if (signature === this.queuedSnapshotSignature) { + return + } + this.queuedSnapshotSignature = signature + + this.resumeSnapshotPersistenceQueue = + this.resumeSnapshotPersistenceQueue.then( + () => this.writeResumeSnapshot(snapshot), + () => this.writeResumeSnapshot(snapshot), + ) + await this.resumeSnapshotPersistenceQueue + } + + private async writeResumeSnapshot( + snapshot: GenerationResumeSnapshot, + ): Promise { + try { + await this.resumePersistence?.setItem(this.resumeSnapshotKey, snapshot) + this.resumePersistenceError = undefined + } catch (error) { + // Warn only on the transition into failure, not once per write. + if (!this.resumePersistenceError) { + console.warn( + '[TanStack AI] Failed to persist generation resume snapshot', + error, + ) + } + this.resumePersistenceError = + error instanceof Error ? error : new Error(String(error)) + // Allow the next snapshot change to retry even if it is materially + // identical to this failed write. + this.queuedSnapshotSignature = undefined + } + } + + private async removePersistedResumeSnapshot(): Promise { + try { + await this.resumePersistence?.removeItem(this.resumeSnapshotKey) + this.resumePersistenceError = undefined + } catch (error) { + if (!this.resumePersistenceError) { + console.warn( + '[TanStack AI] Failed to remove persisted generation resume snapshot', + error, + ) + } + this.resumePersistenceError = + error instanceof Error ? error : new Error(String(error)) + } + } +} + +/** + * Stable serialization of the snapshot minus `lastEvent`, used to detect + * material changes between persistence writes. + */ +function resumeSnapshotSignature(snapshot: GenerationResumeSnapshot): string { + const { lastEvent: _lastEvent, ...significant } = snapshot + return JSON.stringify(significant) } function completeProgressValue( diff --git a/packages/ai-client/src/generation-reconstruct.ts b/packages/ai-client/src/generation-reconstruct.ts new file mode 100644 index 000000000..dcac9ab9e --- /dev/null +++ b/packages/ai-client/src/generation-reconstruct.ts @@ -0,0 +1,89 @@ +import type { + AudioGenerationResult, + ImageGenerationResult, + PersistedArtifactRef, + SummarizationResult, + TranscriptionResult, +} from '@tanstack/ai' +import type { GenerationRestoredResult } from './generation-types' + +/** + * Per-activity `reconstructResult` mappers. On mount restore the generic + * `GenerationClient` hands each specialized hook a {@link GenerationRestoredResult} + * (the metadata that survived persistence plus the durable artifact refs, each + * carrying its serve `url`); the mapper rebuilds the concrete typed result so + * `result` repaints as if the run had just finished, with media resolved to the + * durable serve route rather than the provider's expired link. + * + * A mapper returns `null` when the snapshot cannot rebuild a valid result; then + * `result` stays null while `status` / `error` / `resumeState` still repaint. + */ + +/** Output artifact refs of a given media type that carry a durable serve URL. */ +function mediaUrls( + restored: GenerationRestoredResult, + mediaType: PersistedArtifactRef['source']['mediaType'], +): Array { + return restored.artifacts + .filter( + (a) => + a.role === 'output' && + a.source.mediaType === mediaType && + a.url != null, + ) + .map((a) => a.url as string) +} + +/** image → `{ id, model, images: [{ url }], artifacts }`. */ +export function reconstructImageResult( + restored: GenerationRestoredResult, +): ImageGenerationResult | null { + const urls = mediaUrls(restored, 'image') + if (urls.length === 0) return null + return { + id: restored.id ?? '', + model: restored.model ?? '', + images: urls.map((url) => ({ url })), + ...(restored.artifacts.length > 0 ? { artifacts: restored.artifacts } : {}), + } +} + +/** audio → `{ id, model, audio: { url }, artifacts }`. */ +export function reconstructAudioResult( + restored: GenerationRestoredResult, +): AudioGenerationResult | null { + const [url] = mediaUrls(restored, 'audio') + if (!url) return null + return { + id: restored.id ?? '', + model: restored.model ?? '', + audio: { url }, + ...(restored.artifacts.length > 0 ? { artifacts: restored.artifacts } : {}), + } +} + +/** transcription → `{ id, model, text, artifacts }`. */ +export function reconstructTranscriptionResult( + restored: GenerationRestoredResult, +): TranscriptionResult | null { + if (restored.text === undefined) return null + return { + id: restored.id ?? '', + model: restored.model ?? '', + text: restored.text, + ...(restored.artifacts.length > 0 ? { artifacts: restored.artifacts } : {}), + } +} + +/** summarize → `{ id, model, summary, usage }` (needs persisted `usage`). */ +export function reconstructSummarizeResult( + restored: GenerationRestoredResult, +): SummarizationResult | null { + if (restored.text === undefined || restored.usage === undefined) return null + return { + id: restored.id ?? '', + model: restored.model ?? '', + summary: restored.text, + usage: restored.usage, + } +} diff --git a/packages/ai-client/src/generation-types.ts b/packages/ai-client/src/generation-types.ts index b11e8ca2a..f642e4879 100644 --- a/packages/ai-client/src/generation-types.ts +++ b/packages/ai-client/src/generation-types.ts @@ -1,7 +1,12 @@ -import type { MediaPrompt, StreamChunk } from '@tanstack/ai/client' -import type { TranscriptionResponseFormat } from '@tanstack/ai' +import type { + MediaPrompt, + PersistedArtifactRef, + StreamChunk, +} from '@tanstack/ai/client' +import type { TokenUsage, TranscriptionResponseFormat } from '@tanstack/ai' import type { ConnectConnectionAdapter } from './connection-adapters' import type { AIDevtoolsClientMetadata } from './devtools' +import type { ChatStorageAdapter } from './types' import type { GenerationDevtoolsBridgeFactory, VideoDevtoolsBridgeFactory, @@ -58,6 +63,113 @@ export type InferGenerationOutput = TFn extends ( */ export type GenerationClientState = 'idle' | 'generating' | 'success' | 'error' +export type GenerationResumeStatus = 'idle' | 'running' | 'complete' | 'error' + +/** + * Map a persisted resume status to the client's live state machine on restore: + * complete → success, error → error, running → generating, idle → idle. A + * restored `running` presents as `generating` with `isLoading` false (the client + * never auto-tails a restored run). + */ +export function clientStateFromResumeStatus( + status: GenerationResumeStatus, +): GenerationClientState { + switch (status) { + case 'complete': + return 'success' + case 'error': + return 'error' + case 'running': + return 'generating' + case 'idle': + return 'idle' + } +} + +export interface GenerationResumeState { + threadId: string + runId: string + /** + * Artifact refs observed while the run is still in flight. Non-null only while + * a run is streaming (`resumeState` itself is null once it ends); the final + * refs move onto `result.artifacts` when the run completes. + */ + pendingArtifacts?: Array +} + +export type GenerationPendingArtifact = PersistedArtifactRef + +export interface GenerationResultSnapshot { + id?: string + model?: string + status?: string + jobId?: string + expiresAt?: string + /** + * The text output of a text activity (a transcription's `text` or a summary's + * `summary`). Persisted so a text generation restores its result on reload + * (text is small and not bytes). Absent for media activities, whose output + * restores from `artifacts`. + */ + text?: string + /** Token usage, persisted so a text result that requires it can be rebuilt. */ + usage?: TokenUsage + artifacts?: Array +} + +export interface GenerationErrorSnapshot { + message: string + code?: string +} + +export interface GenerationEventSnapshot { + type: StreamChunk['type'] + name?: string + timestamp?: number +} + +export interface GenerationResumeSnapshot { + /** + * Version of the persisted snapshot shape. Written on every persisted + * snapshot so future shape changes can migrate (or reject) old records. + * Optional so hand-written seeds don't need to set it; absent means `1`. + */ + schemaVersion?: 1 + resumeState: GenerationResumeState | null + status: GenerationResumeStatus + activity?: PersistedArtifactRef['source']['activity'] + pendingArtifacts?: Array + result?: GenerationResultSnapshot + error?: GenerationErrorSnapshot + lastEvent?: GenerationEventSnapshot +} + +/** + * Storage adapter for the lightweight generation resume snapshot. This is the + * same generic {@link ChatStorageAdapter} contract the chat client uses, so the + * `localStoragePersistence` / `sessionStoragePersistence` / `indexedDBPersistence` + * factories work here too. Only the snapshot is ever written — never the + * generated media bytes. + */ +export type GenerationPersistence = ChatStorageAdapter + +/** + * The `persistence` option for a generation client. Mirrors the chat client's + * {@link ChatPersistenceOption}. + * + * - `false` (default) / omitted: ephemeral. Nothing is written; a reload starts + * from empty. + * - `true`: server-driven. Nothing is cached in the browser. On mount the client + * hydrates the last generation for its `threadId` from the server (via the + * connection's `hydrateGeneration` handler) and repaints that snapshot — it + * never auto-starts a run. Requires a connection that implements + * `hydrateGeneration`. + * - a {@link GenerationPersistence} adapter: client-driven. The lightweight + * resume snapshot is cached in the browser under `generation:` as a run + * streams and read back (validated) on mount. + */ +export type GenerationPersistenceOption = boolean | GenerationPersistence + // =========================== // Event Constants // =========================== @@ -70,6 +182,8 @@ export type GenerationClientState = 'idle' | 'generating' | 'success' | 'error' export const GENERATION_EVENTS = { /** The generation result payload */ RESULT: 'generation:result', + /** Persisted artifact refs for generated media */ + ARTIFACTS: 'generation:artifacts', /** Progress update (0-100) with optional message */ PROGRESS: 'generation:progress', /** Video job created with jobId */ @@ -129,12 +243,49 @@ export interface GenerationClientOptions<_TInput, TResult, TOutput = TResult> { /** Unique identifier for this generation client instance */ id?: string + /** + * The thread id for this generation, stable across reloads. Used as the AG-UI + * thread key on the wire AND, in server-driven mode (`persistence: true`), as + * the key the client hydrates the last generation under on mount. Falls back + * to `id`, then to a generated id. + */ + threadId?: string + /** Additional body parameters to send with connect-based adapter requests */ body?: Record /** Metadata used to register this generation hook with TanStack AI Devtools */ devtools?: Partial + /** + * Explicit seed for the lightweight resume snapshot, for apps that manage + * storage themselves. When set, automatic hydration from `persistence` is + * skipped. It does not trigger any generation, but it is **not** inert: it + * seeds the client's live resume snapshot, which subsequent run events merge + * into and which `getResumeSnapshot()` returns and the client re-persists. + * Later reads therefore reflect this seed merged with observed activity, not + * the original value verbatim. + */ + initialResumeSnapshot?: GenerationResumeSnapshot + + /** + * How this generation persists across reloads. See + * {@link GenerationPersistenceOption}. + * + * - Omit or `false`: ephemeral, in-memory only. + * - `true`: server-driven. The client caches nothing and, on mount, hydrates + * the last generation for its `threadId` from the server (needs a connection + * with a `hydrateGeneration` handler). It repaints the snapshot but never + * auto-starts a run. + * - a {@link GenerationPersistence} adapter (any {@link ChatStorageAdapter}, + * including the shared `localStoragePersistence` / `sessionStoragePersistence` + * / `indexedDBPersistence` factories): client-driven. The client writes the + * lightweight snapshot under the key `generation:` as a run streams, and + * reads it back (validated) on construction unless `initialResumeSnapshot` is + * provided. Generated media bytes are never written. + */ + persistence?: GenerationPersistenceOption + /** * Factory that constructs the devtools bridge. Default is a no-op * factory; the real implementation lives in `@tanstack/ai-client/devtools`. @@ -166,6 +317,194 @@ export interface GenerationClientOptions<_TInput, TResult, TOutput = TResult> { onErrorChange?: (error: Error | undefined) => void /** @internal Called when generation status changes */ onStatusChange?: (status: GenerationClientState) => void + /** @internal Called when lightweight resume snapshot changes. Receives `undefined` when the snapshot is cleared by `reset()`. */ + onResumeSnapshotChange?: ( + snapshot: GenerationResumeSnapshot | undefined, + ) => void + /** @internal Called when the in-flight run identity changes. `null` once no run is in flight. Mirrors the chat client's resume-state callback. */ + onResumeStateChange?: (resumeState: GenerationResumeState | null) => void + + /** + * @internal Rebuild a typed result from a restored snapshot, injected by each + * specialized client/hook (which knows the concrete result shape). Called on + * mount restore (client store or server hydrate) so `result` repaints as if the + * run had just finished, with media resolved to the durable serve URL. Returns + * `null` when the snapshot cannot rebuild a result (then `result` stays null; + * `status` / `error` / `resumeState` still repaint). + */ + reconstructResult?: (restored: GenerationRestoredResult) => TResult | null +} + +/** + * The restorable shape handed to a client's `reconstructResult` mapper: the + * result metadata that survived persistence plus the durable artifact refs (each + * carrying its serve {@link PersistedArtifactRef.url}). The specialized client + * turns this into its own typed result (image `images`, video `url`, text + * `text`, ...). + */ +export interface GenerationRestoredResult { + id?: string + model?: string + status?: string + jobId?: string + expiresAt?: string + text?: string + usage?: TokenUsage + activity?: PersistedArtifactRef['source']['activity'] + artifacts: Array +} + +/** + * Reduces one observed stream chunk into the lightweight resume snapshot. + * + * A `RUN_STARTED` chunk begins a fresh run, so stale `result` / `error` / + * `pendingArtifacts` from a previous run are dropped rather than carried into + * the new run's snapshot. + */ +export function updateGenerationResumeSnapshot( + previous: GenerationResumeSnapshot | null | undefined, + chunk: StreamChunk, +): GenerationResumeSnapshot { + const threadId = stringField(chunk, 'threadId') + const runId = stringField(chunk, 'runId') + const carried = chunk.type === 'RUN_STARTED' ? undefined : previous + const previousArtifacts = carried?.pendingArtifacts ?? [] + const next: GenerationResumeSnapshot = { + schemaVersion: 1, + resumeState: carried?.resumeState ?? null, + status: carried?.status ?? 'idle', + ...(carried?.activity ? { activity: carried.activity } : {}), + ...(previousArtifacts.length > 0 + ? { pendingArtifacts: [...previousArtifacts] } + : {}), + ...(carried?.result ? { result: { ...carried.result } } : {}), + ...(carried?.error ? { error: { ...carried.error } } : {}), + lastEvent: createGenerationEventSnapshot(chunk), + } + + if (threadId && runId) { + next.resumeState = { threadId, runId } + next.status = 'running' + } else if (chunk.type === 'RUN_STARTED') { + next.status = 'running' + } + + if (chunk.type === 'CUSTOM') { + if (chunk.name === GENERATION_EVENTS.ARTIFACTS) { + const artifacts = collectArtifactRefs(chunk.value) + if (artifacts.length > 0) { + next.pendingArtifacts = artifacts + next.activity = artifacts[0]?.source.activity + } + } else if (chunk.name === GENERATION_EVENTS.RESULT) { + const result = createGenerationResultSnapshot(chunk.value) + if (result) { + next.result = result + if (result.artifacts && result.artifacts.length > 0) { + next.pendingArtifacts = result.artifacts + next.activity = result.artifacts[0]?.source.activity + } + } + } else if (chunk.name === GENERATION_EVENTS.VIDEO_JOB_CREATED) { + // Capture the job id as soon as the job exists — for a long video run + // this is the one piece of identity worth having after a reload, and + // the terminal `generation:result` may never arrive. + const jobId = isObject(chunk.value) + ? stringField(chunk.value, 'jobId') + : undefined + if (jobId) { + next.result = { ...next.result, jobId } + } + } + } else if (chunk.type === 'RUN_FINISHED') { + next.resumeState = null + next.status = 'complete' + } else if (chunk.type === 'RUN_ERROR') { + next.resumeState = null + next.status = 'error' + next.error = createGenerationErrorSnapshot(chunk) + } + + return next +} + +/** + * Validates an untrusted value (typically read back from a storage adapter) + * into a {@link GenerationResumeSnapshot}, or returns `undefined` when the + * value is not a usable snapshot. + * + * Storage contents are outside the type system — they may be stale, truncated, + * hand-edited, or written by a future version. Every field is re-validated + * with the same narrowing the live chunk reducer uses. `lastEvent` is not + * restored: it describes a transient stream position that has no meaning + * after a reload. + */ +export function parseGenerationResumeSnapshot( + value: unknown, +): GenerationResumeSnapshot | undefined { + if (!isObject(value)) return undefined + + const schemaVersion = Reflect.get(value, 'schemaVersion') + if (schemaVersion !== undefined && schemaVersion !== 1) return undefined + + const status = generationResumeStatusField(value, 'status') + if (!status) return undefined + + const rawResumeState = Reflect.get(value, 'resumeState') + let resumeState: GenerationResumeState | null = null + if (rawResumeState !== null && rawResumeState !== undefined) { + if (!isObject(rawResumeState)) return undefined + const threadId = stringField(rawResumeState, 'threadId') + const runId = stringField(rawResumeState, 'runId') + if (!threadId || !runId) return undefined + resumeState = { threadId, runId } + } + + const snapshot: GenerationResumeSnapshot = { + schemaVersion: 1, + resumeState, + status, + } + + const activity = persistedArtifactActivityField(value, 'activity') + if (activity) snapshot.activity = activity + + const pendingArtifacts = collectArtifactRefs( + Reflect.get(value, 'pendingArtifacts'), + ) + if (pendingArtifacts.length > 0) snapshot.pendingArtifacts = pendingArtifacts + + const result = createGenerationResultSnapshot(Reflect.get(value, 'result')) + if (result) snapshot.result = result + + const rawError = Reflect.get(value, 'error') + if (isObject(rawError)) { + const message = stringField(rawError, 'message') + if (message) { + const code = stringField(rawError, 'code') + snapshot.error = { message, ...(code ? { code } : {}) } + } + } + + return snapshot +} + +function generationResumeStatusField( + value: object, + key: string, +): GenerationResumeStatus | undefined { + const field = stringField(value, key) + if (field === undefined) return undefined + + switch (field) { + case 'idle': + case 'running': + case 'complete': + case 'error': + return field + default: + return undefined + } } // =========================== @@ -200,6 +539,8 @@ export interface VideoGenerateResult { url: string /** When the URL expires, if applicable */ expiresAt?: Date + /** Persisted artifact references for generated assets, when available */ + artifacts?: Array } /** @@ -328,3 +669,249 @@ export interface VideoGenerateInput { /** Model-specific options */ modelOptions?: Record } + +function createGenerationEventSnapshot( + chunk: StreamChunk, +): GenerationEventSnapshot { + const name = stringField(chunk, 'name') + const timestamp = numberField(chunk, 'timestamp') + return { + type: chunk.type, + ...(name ? { name } : {}), + ...(timestamp !== undefined ? { timestamp } : {}), + } +} + +/** @internal Narrows an untrusted result payload into the persisted result snapshot shape. */ +export function createGenerationResultSnapshot( + value: unknown, +): GenerationResultSnapshot | undefined { + if (!isObject(value)) return undefined + + const artifacts = collectArtifactRefs(Reflect.get(value, 'artifacts')) + const snapshot: GenerationResultSnapshot = {} + const id = stringField(value, 'id') + const model = stringField(value, 'model') + const status = stringField(value, 'status') + const jobId = stringField(value, 'jobId') + // A transcription's output is `text`; a summary's is `summary`. Capture either + // under `text` so a text result restores on reload. + const text = stringField(value, 'text') ?? stringField(value, 'summary') + const usage = Reflect.get(value, 'usage') + if (id) snapshot.id = id + if (model) snapshot.model = model + if (status) snapshot.status = status + if (jobId) snapshot.jobId = jobId + if (text) snapshot.text = text + // Passthrough opaque token-usage metadata (untrusted; not deeply validated). + if (isObject(usage)) snapshot.usage = usage as TokenUsage + const expiresAt = Reflect.get(value, 'expiresAt') + if (typeof expiresAt === 'string') { + snapshot.expiresAt = expiresAt + } else if (expiresAt instanceof Date) { + snapshot.expiresAt = expiresAt.toISOString() + } + if (artifacts.length > 0) { + snapshot.artifacts = artifacts + } + + return Object.keys(snapshot).length > 0 ? snapshot : undefined +} + +function createGenerationErrorSnapshot( + chunk: StreamChunk, +): GenerationErrorSnapshot { + const message = + stringField(chunk, 'message') ?? + nestedStringField(chunk, 'error', 'message') ?? + 'An error occurred' + const code = stringField(chunk, 'code') + return { + message, + ...(code ? { code } : {}), + } +} + +function collectArtifactRefs(value: unknown): Array { + if (!Array.isArray(value)) return [] + const refs: Array = [] + for (const item of value) { + const ref = createPersistedArtifactRefSnapshot(item) + if (ref) { + refs.push(ref) + } + } + return refs +} + +function createPersistedArtifactRefSnapshot( + value: unknown, +): PersistedArtifactRef | undefined { + if (!isObject(value)) return undefined + const source = Reflect.get(value, 'source') + if (!isObject(source)) return undefined + + const role = persistedArtifactRoleField(value, 'role') + const artifactId = stringField(value, 'artifactId') + const threadId = stringField(value, 'threadId') + const runId = stringField(value, 'runId') + const name = stringField(value, 'name') + const mimeType = stringField(value, 'mimeType') + const size = numberField(value, 'size') + const createdAt = stringField(value, 'createdAt') + const activity = persistedArtifactActivityField(source, 'activity') + const path = stringField(source, 'path') + const provider = stringField(source, 'provider') + const model = stringField(source, 'model') + if ( + !role || + !artifactId || + !threadId || + !runId || + !name || + !mimeType || + size === undefined || + !createdAt || + !activity || + !path || + !provider || + !model + ) { + return undefined + } + + const externalUrl = durableUrlField(value, 'externalUrl') + const url = serveUrlField(value, 'url') + const mediaType = persistedArtifactMediaTypeField(source, 'mediaType') + const jobId = stringField(source, 'jobId') + const expiresAt = stringField(source, 'expiresAt') + + return { + role, + artifactId, + threadId, + runId, + name, + mimeType, + size, + createdAt, + ...(externalUrl ? { externalUrl } : {}), + ...(url ? { url } : {}), + source: { + activity, + path, + provider, + model, + ...(mediaType ? { mediaType } : {}), + ...(jobId ? { jobId } : {}), + ...(expiresAt ? { expiresAt } : {}), + }, + } +} + +function durableUrlField(value: object, key: string): string | undefined { + const field = stringField(value, key) + if (!field || field.length > 2048) return undefined + try { + const url = new URL(field) + return url.protocol === 'http:' || url.protocol === 'https:' + ? field + : undefined + } catch { + return undefined + } +} + +/** + * Validates an app-origin serve URL, which unlike a provider URL is usually a + * same-origin path (`/api/.../artifact?id=...`). Accepts an absolute http(s) URL + * or a path-absolute same-origin URL (single leading `/`); rejects + * protocol-relative (`//host`), `javascript:` / `data:`, and anything else, since + * this value is rendered as media `src`. + */ +function serveUrlField(value: object, key: string): string | undefined { + const field = stringField(value, key) + if (!field || field.length > 2048) return undefined + // A single leading `/` is a safe same-origin path. Reject protocol-relative + // `//host` AND a backslash bypass (`/\host` — the URL parser treats `\` as `/` + // for http(s), so it would resolve to a foreign origin as an ``). + if (field.startsWith('/') && !field.startsWith('//') && !field.includes('\\')) + return field + try { + const url = new URL(field) + return url.protocol === 'http:' || url.protocol === 'https:' + ? field + : undefined + } catch { + return undefined + } +} + +function persistedArtifactRoleField( + value: object, + key: string, +): PersistedArtifactRef['role'] | undefined { + const field = stringField(value, key) + return field === 'input' || field === 'output' ? field : undefined +} + +function persistedArtifactActivityField( + value: object, + key: string, +): PersistedArtifactRef['source']['activity'] | undefined { + const field = stringField(value, key) + if (field === undefined) return undefined + + switch (field) { + case 'image': + case 'audio': + case 'tts': + case 'video': + case 'transcription': + return field + default: + return undefined + } +} + +function persistedArtifactMediaTypeField( + value: object, + key: string, +): PersistedArtifactRef['source']['mediaType'] | undefined { + const field = stringField(value, key) + if (field === undefined) return undefined + + switch (field) { + case 'image': + case 'audio': + case 'video': + case 'document': + case 'json': + return field + default: + return undefined + } +} + +function nestedStringField( + value: object, + key: string, + nestedKey: string, +): string | undefined { + const nested = Reflect.get(value, key) + return isObject(nested) ? stringField(nested, nestedKey) : undefined +} + +function stringField(value: object, key: string): string | undefined { + const field = Reflect.get(value, key) + return typeof field === 'string' ? field : undefined +} + +function numberField(value: object, key: string): number | undefined { + const field = Reflect.get(value, key) + return typeof field === 'number' ? field : undefined +} + +function isObject(value: unknown): value is object { + return typeof value === 'object' && value !== null +} diff --git a/packages/ai-client/src/index.ts b/packages/ai-client/src/index.ts index e8091f85b..38a5da277 100644 --- a/packages/ai-client/src/index.ts +++ b/packages/ai-client/src/index.ts @@ -70,6 +70,15 @@ export type { InferGenerationOutput, InferGenerationOutputFromReturn, GenerationClientState, + GenerationResumeState, + GenerationResumeStatus, + GenerationResumeSnapshot, + GenerationPendingArtifact, + GenerationResultSnapshot, + GenerationErrorSnapshot, + GenerationEventSnapshot, + GenerationPersistence, + GenerationPersistenceOption, GenerationClientOptions, GenerationFetcher, GenerationFetcherOptions, @@ -83,8 +92,21 @@ export type { TranscriptionGenerateInput, SummarizeGenerateInput, VideoGenerateInput, + GenerationRestoredResult, } from './generation-types' -export { GENERATION_EVENTS } from './generation-types' +export { + GENERATION_EVENTS, + parseGenerationResumeSnapshot, + updateGenerationResumeSnapshot, +} from './generation-types' +// Per-activity result reconstruction mappers (used by the framework hooks to +// repaint a typed `result` on restore) +export { + reconstructImageResult, + reconstructAudioResult, + reconstructTranscriptionResult, + reconstructSummarizeResult, +} from './generation-reconstruct' export { UnsupportedResponseStreamError } from './response-stream' export { clientTools, createChatClientOptions } from './types' // Web storage adapters for durable chat persistence (messages + resume snapshot) @@ -140,6 +162,7 @@ export { StreamReconnectLimitError, type ConnectConnectionAdapter, type ConnectionAdapter, + type GenerationHydrationResult, type FetchConnectionOptions, type ReconnectOptions, type ResumableConnectConnectionAdapter, diff --git a/packages/ai-client/src/storage-adapters.ts b/packages/ai-client/src/storage-adapters.ts index bc67a358c..6cd9447de 100644 --- a/packages/ai-client/src/storage-adapters.ts +++ b/packages/ai-client/src/storage-adapters.ts @@ -1,4 +1,4 @@ -import type { ChatPersistedState, ChatStorageAdapter } from './types' +import type { ChatStorageAdapter } from './types' export interface WebStoragePersistenceOptions { keyPrefix?: string @@ -88,12 +88,12 @@ function createWebStoragePersistence( * adapter can be constructed safely on the server. * * The `serialize` / `deserialize` codec defaults to `JSON.stringify` / - * `JSON.parse`, so the common case needs no codec. `TValue` defaults to - * {@link ChatPersistedState}, so `localStoragePersistence()` drops straight into - * the `persistence` option with no type argument. Pass a codec only for values - * JSON can't round-trip losslessly, and a type argument for non-chat storage. + * `JSON.parse`, so the common case needs no codec. A bare call works for both + * the chat and generation `persistence` options with no type argument; pass an + * explicit `TValue` only when you want a standalone store's own read/write + * typed to a specific shape. */ -export function localStoragePersistence( +export function localStoragePersistence( options: WebStoragePersistenceOptions = {}, ): ChatStorageAdapter { return createWebStoragePersistence('localStorage', options) @@ -102,11 +102,12 @@ export function localStoragePersistence( /** * A `ChatStorageAdapter` backed by `window.sessionStorage` (scoped to the tab * and cleared when it closes). Identical to {@link localStoragePersistence} in - * every other respect: `ChatPersistedState` default `TValue`, `tanstack-ai:` - * default `keyPrefix`, lazy per-operation {@link StorageUnavailableError} on - * SSR, and a JSON codec that defaults to `JSON.stringify` / `JSON.parse`. + * every other respect: a bare call works for both chat and generation + * `persistence` options with no type argument, `tanstack-ai:` default + * `keyPrefix`, lazy per-operation {@link StorageUnavailableError} on SSR, and a + * JSON codec that defaults to `JSON.stringify` / `JSON.parse`. */ -export function sessionStoragePersistence( +export function sessionStoragePersistence( options: WebStoragePersistenceOptions = {}, ): ChatStorageAdapter { return createWebStoragePersistence('sessionStorage', options) @@ -121,9 +122,10 @@ export function sessionStoragePersistence( * * No serialize/deserialize codec is needed or accepted — values are stored via * IndexedDB's native structured clone, so `Date`, `Map`, `ArrayBuffer`, etc. - * round-trip without a JSON step. `TValue` defaults to {@link ChatPersistedState}. + * round-trip without a JSON step. A bare call works for both chat and + * generation `persistence` options with no type argument. */ -export function indexedDBPersistence( +export function indexedDBPersistence( options: IndexedDBPersistenceOptions = {}, ): ChatStorageAdapter { const databaseName = options.databaseName ?? 'tanstack-ai' diff --git a/packages/ai-client/src/video-generation-client.ts b/packages/ai-client/src/video-generation-client.ts index ce8b1fe72..c5e808add 100644 --- a/packages/ai-client/src/video-generation-client.ts +++ b/packages/ai-client/src/video-generation-client.ts @@ -1,9 +1,16 @@ -import { GENERATION_EVENTS } from './generation-types' +import { + GENERATION_EVENTS, + clientStateFromResumeStatus, + createGenerationResultSnapshot, + parseGenerationResumeSnapshot, + updateGenerationResumeSnapshot, +} from './generation-types' import { createNoOpVideoDevtoolsBridge } from './devtools-noop' import { parseSSEResponse } from './sse-parser' import type { StreamChunk } from '@tanstack/ai/client' import type { ConnectConnectionAdapter, + GenerationHydrationResult, RunAgentInputContext, } from './connection-adapters' import type { @@ -15,6 +22,9 @@ import type { import type { GenerationClientState, GenerationFetcher, + GenerationResumeSnapshot, + GenerationPersistence, + GenerationResumeState, VideoGenerateInput, VideoGenerateResult, VideoGenerationClientOptions, @@ -42,6 +52,12 @@ interface VideoCallbacks { onStatusChange?: ((status: GenerationClientState) => void) | undefined onJobIdChange?: ((jobId: string | null) => void) | undefined onVideoStatusChange?: ((status: VideoStatusInfo | null) => void) | undefined + onResumeSnapshotChange?: + | ((snapshot: GenerationResumeSnapshot | undefined) => void) + | undefined + onResumeStateChange?: + | ((resumeState: GenerationResumeState | null) => void) + | undefined } /** @@ -88,6 +104,10 @@ export class VideoGenerationClient { private readonly devtoolsMetadata: AIDevtoolsClientMetadata private readonly devtoolsBridge: VideoDevtoolsBridge private readonly threadId: string + private readonly resumePersistence: GenerationPersistence | undefined + // Server-driven mode (`persistence: true`): no local snapshot store; on mount + // the client hydrates the last generation for `threadId` from the server. + private readonly serverDriven: boolean = false private body: Record private result: TOutput | null = null @@ -98,9 +118,15 @@ export class VideoGenerationClient { private isLoading = false private error: Error | undefined = undefined private status: GenerationClientState = 'idle' + private resumeSnapshot: GenerationResumeSnapshot | undefined + private resumeSnapshotPersistenceQueue: Promise = Promise.resolve() + private resumeSnapshotHydration: Promise | undefined + private queuedSnapshotSignature: string | undefined + private resumePersistenceError: Error | undefined = undefined private abortController: AbortController | null = null private readonly callbacksRef: VideoCallbacks private devtoolsMounted = false + private disposed = false constructor( options: VideoGenerationClientOptions & @@ -113,10 +139,23 @@ export class VideoGenerationClient { ), ) { this.uniqueId = options.id ?? this.generateUniqueId('video') - this.threadId = this.uniqueId + // The wire/hydration thread key. Server-driven mode needs a stable key, so + // prefer an explicit `threadId`, then `id`, then a generated id. + this.threadId = options.threadId ?? this.uniqueId this.connection = options.connection this.fetcher = options.fetcher this.body = options.body ?? {} + // `persistence` is `false`/omitted (ephemeral), `true` (server-driven: cache + // nothing, hydrate the last generation for `threadId` on mount), or a storage + // adapter (client-driven: cache the lightweight resume snapshot locally). + // Only the server-driven mode leaves `resumePersistence` undefined AND turns + // on mount hydration. + if (options.persistence === true) { + this.serverDriven = true + } else if (options.persistence) { + this.resumePersistence = options.persistence + } + this.resumeSnapshot = options.initialResumeSnapshot this.callbacksRef = { onResult: options.onResult, @@ -131,12 +170,31 @@ export class VideoGenerationClient { onStatusChange: options.onStatusChange, onJobIdChange: options.onJobIdChange, onVideoStatusChange: options.onVideoStatusChange, + onResumeSnapshotChange: options.onResumeSnapshotChange, + onResumeStateChange: options.onResumeStateChange, } this.devtoolsMetadata = this.createDevtoolsMetadata(options.devtools) this.devtoolsBridge = ( options.devtoolsBridgeFactory ?? createNoOpVideoDevtoolsBridge )(this.buildDevtoolsBridgeOptions()) + + // After callbacksRef is assigned: hydration may fire + // onResumeSnapshotChange synchronously if an adapter resolves sync. + this.maybeHydrateResumeSnapshot() + + // Server-driven (`persistence: true`): the client caches nothing locally and + // re-hydrates the last generation for its stable threadId from the server on + // mount. Best-effort and non-blocking; it never auto-starts a run. + if (this.serverDriven && this.connection?.hydrateGeneration) { + this.hydrateFromServer() + } else if (this.serverDriven) { + // `persistence: true` without a hydrate-capable connection can never + // restore anything — warn rather than silently no-op. + console.warn( + '[TanStack AI] `persistence: true` (server-driven) needs a connection that implements `hydrateGeneration` (e.g. `fetchServerSentEvents` / `fetchHttpStream`). With a plain `fetcher` or no connection, nothing is persisted or restored.', + ) + } } private buildDevtoolsBridgeOptions(): VideoDevtoolsBridgeOptions { @@ -159,6 +217,12 @@ export class VideoGenerationClient { } mountDevtools(): void { + // Mounting revives a disposed client. Framework hooks call this from + // their mount effect, so a dispose → remount cycle (e.g. React + // StrictMode's mount → cleanup → mount replay against the same memoized + // client) leaves the client usable again. + this.disposed = false + this.maybeHydrateResumeSnapshot() if (this.devtoolsMounted) { return } @@ -173,8 +237,9 @@ export class VideoGenerationClient { * Only one generation can be in-flight at a time. */ async generate(input: VideoGenerateInput): Promise { - this.mountDevtools() + if (this.disposed) return if (this.isLoading) return + this.mountDevtools() this.input = input this.progress = null @@ -200,7 +265,7 @@ export class VideoGenerationClient { signal, this.createRunContext(runId), ) - await this.processStream(stream, runId) + await this.processStream(stream, runId, signal) } else { throw new Error( 'VideoGenerationClient requires either a connection or fetcher option', @@ -218,6 +283,7 @@ export class VideoGenerationClient { const error = err instanceof Error ? err : new Error(String(err)) this.setError(error) this.setStatus('error') + this.recordResumeSnapshotError(error) this.devtoolsBridge.finishRun( this.devtoolsBridge.getActiveRunId() ?? runId, 'run:errored', @@ -226,8 +292,10 @@ export class VideoGenerationClient { ) this.callbacksRef.onError?.(error) } finally { - this.abortController = null - this.setIsLoading(false) + if (this.abortController === abortController) { + this.abortController = null + this.setIsLoading(false) + } } } @@ -247,11 +315,12 @@ export class VideoGenerationClient { if (result instanceof Response) { // Server function returned SSE Response — parse stream - await this.processStream(parseSSEResponse(result, signal), runId) + await this.processStream(parseSSEResponse(result, signal), runId, signal) } else { this.devtoolsBridge.ensureRunStarted(runId) this.setResult(result) this.setStatus('success') + this.completePlainFetcherResumeSnapshot(result) } } @@ -262,13 +331,15 @@ export class VideoGenerationClient { private async processStream( source: AsyncIterable, fallbackRunId: string, + signal: AbortSignal, ): Promise { let streamRunId: string | undefined for await (const chunk of source) { - if (this.abortController?.signal.aborted) break + if (signal.aborted) break this.callbacksRef.onChunk?.(chunk) + this.observeResumeSnapshot(chunk) const chunkRunId = 'runId' in chunk && typeof chunk.runId === 'string' ? chunk.runId @@ -344,10 +415,22 @@ export class VideoGenerationClient { this.devtoolsBridge.finishRun(runId, 'run:cancelled', 'cancelled') } } + // A stopped run is no longer resumable. Without this, storage keeps a + // `running` snapshot forever and a reload would chase a dead run. + if (this.resumeSnapshot && this.resumeSnapshot.status === 'running') { + this.resumeSnapshot = { + ...this.resumeSnapshot, + resumeState: null, + status: 'idle', + } + this.notifyResumeSnapshotChanged() + void this.persistResumeSnapshot(this.resumeSnapshot) + } } /** - * Clear all state and return to idle. + * Clear all state and return to idle. Also clears the resume snapshot, + * removing any persisted record for this client id. */ reset(): void { this.stop() @@ -359,6 +442,7 @@ export class VideoGenerationClient { this.setVideoStatus(null) this.setError(undefined) this.setStatus('idle') + this.clearResumeSnapshot() this.devtoolsBridge.emitState() } @@ -403,6 +487,7 @@ export class VideoGenerationClient { } dispose(): void { + this.disposed = true this.stop() this.devtoolsBridge.dispose() this.devtoolsMounted = false @@ -432,10 +517,41 @@ export class VideoGenerationClient { return this.error } + getResumePersistenceError(): Error | undefined { + return this.resumePersistenceError + } + getStatus(): GenerationClientState { return this.status } + getResumeSnapshot(): GenerationResumeSnapshot | undefined { + return this.resumeSnapshot + ? { + ...this.resumeSnapshot, + ...(this.resumeSnapshot.pendingArtifacts + ? { pendingArtifacts: [...this.resumeSnapshot.pendingArtifacts] } + : {}), + ...(this.resumeSnapshot.result + ? { + result: { + ...this.resumeSnapshot.result, + ...(this.resumeSnapshot.result.artifacts + ? { artifacts: [...this.resumeSnapshot.result.artifacts] } + : {}), + }, + } + : {}), + ...(this.resumeSnapshot.error + ? { error: { ...this.resumeSnapshot.error } } + : {}), + ...(this.resumeSnapshot.lastEvent + ? { lastEvent: { ...this.resumeSnapshot.lastEvent } } + : {}), + } + : undefined + } + // =========================== // Private state setters // =========================== @@ -556,4 +672,289 @@ export class VideoGenerationClient { runId, } } + + private observeResumeSnapshot(chunk: StreamChunk): void { + this.resumeSnapshot = updateGenerationResumeSnapshot( + this.resumeSnapshot, + chunk, + ) + this.notifyResumeSnapshotChanged() + void this.persistResumeSnapshot(this.resumeSnapshot) + } + + /** Notify the internal snapshot listener AND emit the public resume state. */ + private notifyResumeSnapshotChanged(): void { + this.callbacksRef.onResumeSnapshotChange?.(this.resumeSnapshot) + this.emitResumeState() + } + + /** Derive the public `resumeState` from the internal snapshot. */ + private emitResumeState(): void { + const snapshot = this.resumeSnapshot + const state = snapshot?.resumeState + const resumeState: GenerationResumeState | null = state + ? { + ...state, + ...(snapshot?.pendingArtifacts && snapshot.pendingArtifacts.length > 0 + ? { pendingArtifacts: [...snapshot.pendingArtifacts] } + : {}), + } + : null + this.callbacksRef.onResumeStateChange?.(resumeState) + } + + /** + * Repaint the normal fields from a restored snapshot so a reload presents the + * video in `result` / `status` / `error` / `jobId`, never a snapshot object. + * `isLoading` stays false (no auto-tail). Not re-persisted (it came from + * storage / the server). + */ + private repaintFromSnapshot(snapshot: GenerationResumeSnapshot): void { + this.resumeSnapshot = snapshot + this.notifyResumeSnapshotChanged() + this.setStatus(clientStateFromResumeStatus(snapshot.status)) + this.setError( + snapshot.error + ? Object.assign( + new Error(snapshot.error.message), + snapshot.error.code ? { code: snapshot.error.code } : {}, + ) + : undefined, + ) + if (snapshot.result?.jobId) this.setJobId(snapshot.result.jobId) + const restored = this.reconstructVideoResult(snapshot) + if (restored !== null) this.setResult(restored) + } + + /** + * Rebuild a `VideoGenerateResult` from a restored snapshot: the video's bytes + * are served from the durable artifact URL, so the restored result renders + * from your own origin. Returns `null` when there is no durable video artifact. + */ + private reconstructVideoResult( + snapshot: GenerationResumeSnapshot, + ): VideoGenerateResult | null { + const result = snapshot.result + const artifacts = result?.artifacts ?? [] + const output = artifacts.find( + (a) => + a.role === 'output' && a.source.mediaType === 'video' && a.url != null, + ) + if (!output?.url) return null + return { + jobId: result?.jobId ?? '', + status: 'completed', + url: output.url, + ...(result?.expiresAt ? { expiresAt: new Date(result.expiresAt) } : {}), + artifacts, + } + } + + /** + * The plain (non-Response) fetcher path never observes stream chunks, so + * the terminal snapshot is built here from the fetcher's own result. A + * stale `error` from a previous run is intentionally dropped — this run + * succeeded. + */ + private completePlainFetcherResumeSnapshot(rawResult: unknown): void { + const previous = this.resumeSnapshot + const result = createGenerationResultSnapshot(rawResult) + this.resumeSnapshot = { + schemaVersion: 1, + resumeState: null, + status: 'complete', + ...(previous?.activity ? { activity: previous.activity } : {}), + ...(previous?.pendingArtifacts && previous.pendingArtifacts.length > 0 + ? { pendingArtifacts: [...previous.pendingArtifacts] } + : {}), + ...(result + ? { result } + : previous?.result + ? { result: { ...previous.result } } + : {}), + } + this.notifyResumeSnapshotChanged() + void this.persistResumeSnapshot(this.resumeSnapshot) + } + + /** + * Records a transport-level failure (network drop, throwing callback) in + * the snapshot. Without this, only a server-emitted RUN_ERROR chunk would + * mark the snapshot `error`, leaving a persisted record that claims the + * run is still in flight. + */ + private recordResumeSnapshotError(error: Error): void { + if (this.resumeSnapshot?.status === 'error') return + if (!this.resumeSnapshot && !this.resumePersistence) return + const previous = this.resumeSnapshot + this.resumeSnapshot = { + schemaVersion: 1, + resumeState: null, + status: 'error', + ...(previous?.activity ? { activity: previous.activity } : {}), + ...(previous?.pendingArtifacts && previous.pendingArtifacts.length > 0 + ? { pendingArtifacts: [...previous.pendingArtifacts] } + : {}), + ...(previous?.result ? { result: { ...previous.result } } : {}), + error: { message: error.message }, + } + this.notifyResumeSnapshotChanged() + void this.persistResumeSnapshot(this.resumeSnapshot) + } + + private clearResumeSnapshot(): void { + this.resumeSnapshot = undefined + this.queuedSnapshotSignature = undefined + this.notifyResumeSnapshotChanged() + if (!this.resumePersistence) { + return + } + this.resumeSnapshotPersistenceQueue = + this.resumeSnapshotPersistenceQueue.then( + () => this.removePersistedResumeSnapshot(), + () => this.removePersistedResumeSnapshot(), + ) + } + + /** + * Storage key for this client's snapshot. The `generation:` segment keeps + * a generation client and a chat client that share an id (and a storage + * adapter with the default key prefix) from overwriting each other. + */ + private get resumeSnapshotKey(): string { + return `generation:${this.uniqueId}` + } + + private maybeHydrateResumeSnapshot(): void { + if (!this.resumePersistence || this.resumeSnapshotHydration) return + // An explicit `initialResumeSnapshot` seed takes precedence over storage. + if (this.resumeSnapshot) return + this.resumeSnapshotHydration = this.hydrateResumeSnapshot() + } + + private async hydrateResumeSnapshot(): Promise { + let stored: unknown + try { + stored = await this.resumePersistence?.getItem(this.resumeSnapshotKey) + } catch (error) { + // A corrupt record (e.g. truncated JSON) or unavailable storage must + // not break construction; the app just starts without a snapshot. + console.warn( + '[TanStack AI] Failed to read persisted generation resume snapshot', + error, + ) + return + } + if (stored === null || stored === undefined) return + const snapshot = parseGenerationResumeSnapshot(stored) + if (!snapshot) return + // Live state wins: adopt the stored snapshot only if nothing has been + // observed since construction. + if (this.resumeSnapshot || this.isLoading || this.status !== 'idle') return + this.repaintFromSnapshot(snapshot) + } + + /** + * Server-driven mount hydration (`persistence: true`). The client holds no + * local snapshot; on mount it asks the server — keyed by the stable threadId — + * for the last generation's resume snapshot, validates it, and repaints it. It + * never auto-starts a run. Best-effort and non-blocking: a failure leaves the + * client empty rather than throwing, and a `generate()` that starts first owns + * the client (hydration then backs off, mirroring the chat client). + */ + private hydrateFromServer(): void { + const hydrate = this.connection?.hydrateGeneration + if (!hydrate) return + // A send that already started owns the client; don't stomp it. + if (this.resumeSnapshot || this.isLoading || this.status !== 'idle') return + void (async () => { + let res: GenerationHydrationResult + try { + res = await hydrate(this.threadId) + } catch { + return + } + if (!res.resumeSnapshot) return + const snapshot = parseGenerationResumeSnapshot(res.resumeSnapshot) + if (!snapshot) return + // Re-check: a send may have started while the fetch was in flight. + if (this.resumeSnapshot || this.isLoading || this.status !== 'idle') + return + this.repaintFromSnapshot(snapshot) + })() + } + + private async persistResumeSnapshot( + snapshot: GenerationResumeSnapshot, + ): Promise { + if (!this.resumePersistence) { + return + } + + // Skip writes that only differ in `lastEvent` — for a long streaming run + // this collapses hundreds of per-chunk writes into the handful where the + // snapshot materially changed. (`lastEvent` in storage is best-effort; + // hydration drops it anyway.) + const signature = videoResumeSnapshotSignature(snapshot) + if (signature === this.queuedSnapshotSignature) { + return + } + this.queuedSnapshotSignature = signature + + this.resumeSnapshotPersistenceQueue = + this.resumeSnapshotPersistenceQueue.then( + () => this.writeResumeSnapshot(snapshot), + () => this.writeResumeSnapshot(snapshot), + ) + await this.resumeSnapshotPersistenceQueue + } + + private async writeResumeSnapshot( + snapshot: GenerationResumeSnapshot, + ): Promise { + try { + await this.resumePersistence?.setItem(this.resumeSnapshotKey, snapshot) + this.resumePersistenceError = undefined + } catch (error) { + // Warn only on the transition into failure, not once per write. + if (!this.resumePersistenceError) { + console.warn( + '[TanStack AI] Failed to persist generation resume snapshot', + error, + ) + } + this.resumePersistenceError = + error instanceof Error ? error : new Error(String(error)) + // Allow the next snapshot change to retry even if it is materially + // identical to this failed write. + this.queuedSnapshotSignature = undefined + } + } + + private async removePersistedResumeSnapshot(): Promise { + try { + await this.resumePersistence?.removeItem(this.resumeSnapshotKey) + this.resumePersistenceError = undefined + } catch (error) { + if (!this.resumePersistenceError) { + console.warn( + '[TanStack AI] Failed to remove persisted generation resume snapshot', + error, + ) + } + this.resumePersistenceError = + error instanceof Error ? error : new Error(String(error)) + } + } +} + +/** + * Stable serialization of the snapshot minus `lastEvent`, used to detect + * material changes between persistence writes. + */ +function videoResumeSnapshotSignature( + snapshot: GenerationResumeSnapshot, +): string { + const { lastEvent: _lastEvent, ...significant } = snapshot + return JSON.stringify(significant) } diff --git a/packages/ai-client/tests/generation-client.test.ts b/packages/ai-client/tests/generation-client.test.ts index fb921a89a..af5ac443a 100644 --- a/packages/ai-client/tests/generation-client.test.ts +++ b/packages/ai-client/tests/generation-client.test.ts @@ -1,8 +1,39 @@ import { describe, expect, it, vi } from 'vitest' import { EventType } from '@tanstack/ai/client' -import { GenerationClient, UnsupportedResponseStreamError } from '../src' +import { + GenerationClient, + UnsupportedResponseStreamError, + VideoGenerationClient, + reconstructImageResult, +} from '../src' import type { StreamChunk } from '@tanstack/ai/client' -import type { ConnectConnectionAdapter } from '../src/connection-adapters' +import type { PersistedArtifactRef } from '@tanstack/ai/client' +import type { + ConnectConnectionAdapter, + GenerationHydrationResult, +} from '../src/connection-adapters' +import type { GenerationResumeSnapshot, GenerationPersistence } from '../src' + +// A durable output image artifact carrying an app-origin serve URL, used to +// verify the restore-path result repaint via reconstructImageResult. +const restoredImageArtifact: PersistedArtifactRef = { + role: 'output', + artifactId: 'artifact-image-1', + threadId: 'thread-img', + runId: 'run-img', + name: 'image.png', + mimeType: 'image/png', + size: 2048, + createdAt: '2026-07-06T00:00:00.000Z', + url: '/api/artifacts/artifact-image-1', + source: { + activity: 'image', + path: 'runs/run-img/image.png', + provider: 'test', + model: 'test-image', + mediaType: 'image', + }, +} // Helper to create a mock connect-based adapter from StreamChunks function createMockConnection( @@ -17,6 +48,34 @@ function createMockConnection( } } +function createDeferred(): { + promise: Promise + resolve: (value: T | PromiseLike) => void + reject: (reason?: unknown) => void +} { + let resolve!: (value: T | PromiseLike) => void + let reject!: (reason?: unknown) => void + const promise = new Promise((res, rej) => { + resolve = res + reject = rej + }) + return { promise, resolve, reject } +} + +async function waitForCondition(assertion: () => void): Promise { + let lastError: unknown + for (let attempt = 0; attempt < 20; attempt++) { + try { + assertion() + return + } catch (error) { + lastError = error + await new Promise((resolve) => setTimeout(resolve, 0)) + } + } + throw lastError +} + describe('GenerationClient', () => { describe('fetcher mode', () => { it('should generate a result using fetcher', async () => { @@ -433,6 +492,266 @@ describe('GenerationClient', () => { expect(client.getStatus()).toBe('idle') }) + it('should ignore chunks yielded after stop() by an abort-ignoring connection', async () => { + const onResult = vi.fn() + const aborted = createDeferred() + + const connection: ConnectConnectionAdapter = { + async *connect(_msgs, _data, signal) { + yield { + type: EventType.RUN_STARTED as const, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + } + signal?.addEventListener('abort', () => aborted.resolve(undefined), { + once: true, + }) + await aborted.promise + yield { + type: EventType.CUSTOM as const, + name: 'generation:result', + value: { id: 'late-result' }, + timestamp: Date.now(), + } + yield { + type: EventType.RUN_FINISHED as const, + runId: 'run-1', + threadId: 'thread-1', + finishReason: 'stop' as const, + timestamp: Date.now(), + } + }, + } + + const client = new GenerationClient({ + connection, + onResult, + }) + + const generatePromise = client.generate({ prompt: 'test' }) + await new Promise((resolve) => setTimeout(resolve, 0)) + + client.stop() + await generatePromise + + expect(onResult).not.toHaveBeenCalled() + expect(client.getResult()).toBeNull() + expect(client.getStatus()).toBe('idle') + }) + + it('should not let a stopped run clear the controller for a newer generation', async () => { + const firstAborted = createDeferred() + const firstCanFinish = createDeferred() + const secondAborted = createDeferred() + const signals: Array = [] + + const connection: ConnectConnectionAdapter = { + async *connect(_msgs, data, signal) { + signals.push(signal) + if (data?.prompt === 'first') { + yield { + type: EventType.RUN_STARTED as const, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + } + signal?.addEventListener( + 'abort', + () => firstAborted.resolve(undefined), + { once: true }, + ) + await firstAborted.promise + await firstCanFinish.promise + yield { + type: EventType.CUSTOM as const, + name: 'generation:result', + value: { id: 'late-first' }, + timestamp: Date.now(), + } + return + } + + yield { + type: EventType.RUN_STARTED as const, + runId: 'run-2', + threadId: 'thread-1', + timestamp: Date.now(), + } + signal?.addEventListener( + 'abort', + () => secondAborted.resolve(undefined), + { once: true }, + ) + await secondAborted.promise + }, + } + + const client = new GenerationClient({ + connection, + }) + + const firstGenerate = client.generate({ prompt: 'first' }) + await waitForCondition(() => { + expect(signals).toHaveLength(1) + }) + + client.stop() + const secondGenerate = client.generate({ prompt: 'second' }) + await waitForCondition(() => { + expect(signals).toHaveLength(2) + expect(client.getIsLoading()).toBe(true) + }) + + firstCanFinish.resolve(undefined) + await firstGenerate + + expect(client.getIsLoading()).toBe(true) + + client.stop() + expect(signals[1]?.aborted).toBe(true) + + await secondGenerate + expect(client.getIsLoading()).toBe(false) + }) + + it('should ignore video chunks yielded after stop() by an abort-ignoring connection', async () => { + const onResult = vi.fn() + const onStatusUpdate = vi.fn() + const aborted = createDeferred() + + const connection: ConnectConnectionAdapter = { + async *connect(_msgs, _data, signal) { + yield { + type: EventType.RUN_STARTED as const, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + } + signal?.addEventListener('abort', () => aborted.resolve(undefined), { + once: true, + }) + await aborted.promise + yield { + type: EventType.CUSTOM as const, + name: 'generation:result', + value: { id: 'late-video' }, + timestamp: Date.now(), + } + yield { + type: EventType.CUSTOM as const, + name: 'video:status', + value: { status: 'completed', progress: 100 }, + timestamp: Date.now(), + } + yield { + type: EventType.RUN_FINISHED as const, + runId: 'run-1', + threadId: 'thread-1', + finishReason: 'stop' as const, + timestamp: Date.now(), + } + }, + } + + const client = new VideoGenerationClient({ + connection, + onResult, + onStatusUpdate, + }) + + const generatePromise = client.generate({ prompt: 'test' }) + await new Promise((resolve) => setTimeout(resolve, 0)) + + client.stop() + await generatePromise + + expect(onResult).not.toHaveBeenCalled() + expect(onStatusUpdate).not.toHaveBeenCalled() + expect(client.getResult()).toBeNull() + expect(client.getVideoStatus()).toBeNull() + expect(client.getStatus()).toBe('idle') + }) + + it('should not let a stopped video run clear the controller for a newer generation', async () => { + const firstAborted = createDeferred() + const firstCanFinish = createDeferred() + const secondAborted = createDeferred() + const signals: Array = [] + + const connection: ConnectConnectionAdapter = { + async *connect(_msgs, data, signal) { + signals.push(signal) + if (data?.prompt === 'first') { + yield { + type: EventType.RUN_STARTED as const, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + } + signal?.addEventListener( + 'abort', + () => firstAborted.resolve(undefined), + { once: true }, + ) + await firstAborted.promise + await firstCanFinish.promise + yield { + type: EventType.CUSTOM as const, + name: 'generation:result', + value: { + jobId: 'late-first', + status: 'completed', + url: 'https://example.com/late.mp4', + }, + timestamp: Date.now(), + } + return + } + + yield { + type: EventType.RUN_STARTED as const, + runId: 'run-2', + threadId: 'thread-1', + timestamp: Date.now(), + } + signal?.addEventListener( + 'abort', + () => secondAborted.resolve(undefined), + { once: true }, + ) + await secondAborted.promise + }, + } + + const client = new VideoGenerationClient({ + connection, + }) + + const firstGenerate = client.generate({ prompt: 'first' }) + await waitForCondition(() => { + expect(signals).toHaveLength(1) + }) + + client.stop() + const secondGenerate = client.generate({ prompt: 'second' }) + await waitForCondition(() => { + expect(signals).toHaveLength(2) + expect(client.getIsLoading()).toBe(true) + }) + + firstCanFinish.resolve(undefined) + await firstGenerate + + expect(client.getIsLoading()).toBe(true) + + client.stop() + expect(signals[1]?.aborted).toBe(true) + + await secondGenerate + expect(client.getIsLoading()).toBe(false) + }) + it('should not set result if fetcher resolves after stop()', async () => { let resolvePromise: (value: { id: string }) => void const onResult = vi.fn() @@ -1031,4 +1350,729 @@ describe('GenerationClient', () => { expect(states).toEqual(['generating', 'error']) }) }) + + describe('resume snapshot persistence', () => { + it('reports rejected persistence writes without rejecting generation', async () => { + const warningSpy = vi.spyOn(console, 'warn').mockImplementation(() => {}) + const persistenceError = new Error('persistence failed') + const persistence: GenerationPersistence = { + getItem: vi.fn(), + setItem: vi.fn(async () => { + throw persistenceError + }), + removeItem: vi.fn(), + } + const client = new GenerationClient({ + connection: createMockConnection([ + { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + }, + ]), + persistence: persistence, + }) + + await expect(client.generate({ prompt: 'test' })).resolves.toBeUndefined() + + await waitForCondition(() => { + expect(warningSpy).toHaveBeenCalledWith( + '[TanStack AI] Failed to persist generation resume snapshot', + persistenceError, + ) + }) + + warningSpy.mockRestore() + }) + + it('keeps a delayed running write from overwriting a terminal complete snapshot', async () => { + const runningWrite = createDeferred() + let storedSnapshot: GenerationResumeSnapshot | undefined + const persistence: GenerationPersistence = { + getItem: vi.fn(), + setItem: vi.fn(async (_id, snapshot) => { + if (snapshot.status === 'running') { + await runningWrite.promise + } + storedSnapshot = snapshot + }), + removeItem: vi.fn(), + } + const client = new GenerationClient({ + connection: createMockConnection([ + { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + }, + { + type: EventType.RUN_FINISHED, + runId: 'run-1', + threadId: 'thread-1', + finishReason: 'stop', + timestamp: Date.now(), + }, + ]), + persistence: persistence, + }) + + await client.generate({ prompt: 'test' }) + expect(persistence.setItem).toHaveBeenCalledTimes(1) + runningWrite.resolve(undefined) + + await waitForCondition(() => { + expect(persistence.setItem).toHaveBeenCalledTimes(2) + expect(storedSnapshot).toMatchObject({ + status: 'complete', + resumeState: null, + }) + }) + }) + + it('keeps a delayed video running write from overwriting a terminal error snapshot', async () => { + const runningWrite = createDeferred() + let storedSnapshot: GenerationResumeSnapshot | undefined + const persistence: GenerationPersistence = { + getItem: vi.fn(), + setItem: vi.fn(async (_id, snapshot) => { + if (snapshot.status === 'running') { + await runningWrite.promise + } + storedSnapshot = snapshot + }), + removeItem: vi.fn(), + } + const client = new VideoGenerationClient({ + connection: createMockConnection([ + { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + }, + { + type: EventType.RUN_ERROR, + runId: 'run-1', + threadId: 'thread-1', + message: 'Video failed', + timestamp: Date.now(), + }, + ]), + persistence: persistence, + }) + + await client.generate({ prompt: 'test' }) + expect(persistence.setItem).toHaveBeenCalledTimes(1) + runningWrite.resolve(undefined) + + await waitForCondition(() => { + expect(persistence.setItem).toHaveBeenCalledTimes(2) + expect(storedSnapshot).toMatchObject({ + status: 'error', + resumeState: null, + error: { message: 'Video failed' }, + }) + }) + }) + }) + + describe('resume snapshot lifecycle', () => { + function createMapPersistence(seed?: Record): { + store: Map + persistence: GenerationPersistence + } { + const store = new Map(Object.entries(seed ?? {})) + const persistence = { + getItem: vi.fn((key: string) => store.get(key) ?? null), + setItem: vi.fn((key: string, value: unknown) => { + store.set(key, value) + }), + removeItem: vi.fn((key: string) => { + store.delete(key) + }), + // The Map-backed fake is looser than GenerationPersistence's value + // type on purpose: hydration must survive arbitrary stored shapes. + } as unknown as GenerationPersistence + return { store, persistence } + } + + const storedSnapshot: GenerationResumeSnapshot = { + schemaVersion: 1, + resumeState: null, + status: 'complete', + result: { id: 'result-1', model: 'image-model' }, + } + + it('hydrates a stored snapshot under generation: on construction', async () => { + const { persistence } = createMapPersistence({ + 'generation:hero': storedSnapshot, + }) + const onResumeSnapshotChange = vi.fn() + const client = new GenerationClient({ + id: 'hero', + connection: createMockConnection([]), + persistence, + onResumeSnapshotChange, + }) + + await waitForCondition(() => { + expect(client.getResumeSnapshot()).toMatchObject({ + status: 'complete', + result: { id: 'result-1' }, + }) + }) + expect(persistence.getItem).toHaveBeenCalledWith('generation:hero') + expect(onResumeSnapshotChange).toHaveBeenCalledWith( + expect.objectContaining({ status: 'complete' }), + ) + }) + + it('repaints status/result from a stored complete image snapshot via reconstructResult', async () => { + const snapshot: GenerationResumeSnapshot = { + schemaVersion: 1, + resumeState: null, + status: 'complete', + activity: 'image', + result: { + id: 'img-1', + model: 'test-image', + artifacts: [restoredImageArtifact], + }, + } + const { persistence } = createMapPersistence({ + 'generation:img': snapshot, + }) + const onStatusChange = vi.fn() + const onResumeStateChange = vi.fn() + const client = new GenerationClient({ + id: 'img', + connection: createMockConnection([]), + persistence, + // Injected by the image hook/client — rebuilds the typed result from + // the durable artifact refs. + reconstructResult: reconstructImageResult, + onStatusChange, + onResumeStateChange, + }) + + // The completed snapshot repaints `result` (with the durable serve url) + // and `status`, as if the run had just finished. + await waitForCondition(() => { + expect(client.getResult()).toEqual({ + id: 'img-1', + model: 'test-image', + images: [{ url: '/api/artifacts/artifact-image-1' }], + artifacts: [restoredImageArtifact], + }) + }) + expect(client.getStatus()).toBe('success') + // A completed run has no in-flight identity. + expect(onResumeStateChange).toHaveBeenLastCalledWith(null) + }) + + it('emits onResumeStateChange with the in-flight identity from a stored running snapshot', async () => { + const snapshot: GenerationResumeSnapshot = { + schemaVersion: 1, + resumeState: { threadId: 'thread-run', runId: 'run-run' }, + status: 'running', + } + const { persistence } = createMapPersistence({ + 'generation:running': snapshot, + }) + const onResumeStateChange = vi.fn() + const client = new GenerationClient({ + id: 'running', + connection: createMockConnection([]), + persistence, + onResumeStateChange, + }) + + await waitForCondition(() => { + expect(onResumeStateChange).toHaveBeenCalledWith({ + threadId: 'thread-run', + runId: 'run-run', + }) + }) + // A restored running run presents as `generating` but never auto-tails. + expect(client.getStatus()).toBe('generating') + expect(client.getIsLoading()).toBe(false) + }) + + it('skips hydration when an explicit initialResumeSnapshot seed is provided', async () => { + const { persistence } = createMapPersistence({ + 'generation:hero': storedSnapshot, + }) + const seed: GenerationResumeSnapshot = { + resumeState: null, + status: 'error', + error: { message: 'seeded' }, + } + const client = new GenerationClient({ + id: 'hero', + connection: createMockConnection([]), + persistence, + initialResumeSnapshot: seed, + }) + + await new Promise((resolve) => setTimeout(resolve, 0)) + expect(persistence.getItem).not.toHaveBeenCalled() + expect(client.getResumeSnapshot()).toMatchObject({ + status: 'error', + error: { message: 'seeded' }, + }) + }) + + it('ignores invalid stored values and read failures without throwing', async () => { + const warningSpy = vi.spyOn(console, 'warn').mockImplementation(() => {}) + const { persistence } = createMapPersistence({ + 'generation:hero': { status: 'not-a-status' }, + }) + const client = new GenerationClient({ + id: 'hero', + connection: createMockConnection([]), + persistence, + }) + await new Promise((resolve) => setTimeout(resolve, 0)) + expect(client.getResumeSnapshot()).toBeUndefined() + + const throwingPersistence = { + getItem: vi.fn(() => { + throw new Error('corrupt JSON') + }), + setItem: vi.fn(), + removeItem: vi.fn(), + } as unknown as GenerationPersistence + const client2 = new GenerationClient({ + id: 'hero', + connection: createMockConnection([]), + persistence: throwingPersistence, + }) + await new Promise((resolve) => setTimeout(resolve, 0)) + expect(client2.getResumeSnapshot()).toBeUndefined() + warningSpy.mockRestore() + }) + + it('marks the snapshot no longer resumable when stop() aborts a run', async () => { + const gate = createDeferred() + const connection: ConnectConnectionAdapter = { + async *connect() { + yield { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + } satisfies StreamChunk + await gate.promise + }, + } + const { store, persistence } = createMapPersistence() + const client = new GenerationClient({ + id: 'hero', + connection, + persistence, + }) + + const generatePromise = client.generate({ prompt: 'test' }) + await waitForCondition(() => { + expect(client.getResumeSnapshot()?.status).toBe('running') + }) + client.stop() + gate.resolve(undefined) + await generatePromise + + expect(client.getResumeSnapshot()).toMatchObject({ + status: 'idle', + resumeState: null, + }) + await waitForCondition(() => { + expect(store.get('generation:hero')).toMatchObject({ + status: 'idle', + resumeState: null, + }) + }) + }) + + it('records a transport-level failure as an error snapshot', async () => { + const connection: ConnectConnectionAdapter = { + // eslint-disable-next-line require-yield -- the stream fails before producing any chunk + async *connect() { + throw new Error('network dropped') + }, + } + const { store, persistence } = createMapPersistence() + const client = new GenerationClient({ + id: 'hero', + connection, + persistence, + }) + + await client.generate({ prompt: 'test' }) + + expect(client.getResumeSnapshot()).toMatchObject({ + status: 'error', + resumeState: null, + error: { message: 'network dropped' }, + }) + await waitForCondition(() => { + expect(store.get('generation:hero')).toMatchObject({ + status: 'error', + error: { message: 'network dropped' }, + }) + }) + }) + + it('records a complete snapshot for plain fetcher runs with no seed', async () => { + const { store, persistence } = createMapPersistence() + const client = new GenerationClient({ + id: 'hero', + fetcher: async () => ({ id: 'result-9', model: 'image-model' }), + persistence, + }) + + await client.generate({ prompt: 'test' }) + + expect(client.getResumeSnapshot()).toMatchObject({ + status: 'complete', + resumeState: null, + result: { id: 'result-9', model: 'image-model' }, + }) + await waitForCondition(() => { + expect(store.get('generation:hero')).toMatchObject({ + status: 'complete', + }) + }) + }) + + it('reset() clears the snapshot and removes the persisted record', async () => { + const { store, persistence } = createMapPersistence() + const onResumeSnapshotChange = vi.fn() + const client = new GenerationClient({ + id: 'hero', + fetcher: async () => ({ id: 'result-9' }), + persistence, + onResumeSnapshotChange, + }) + + await client.generate({ prompt: 'test' }) + await waitForCondition(() => { + expect(store.has('generation:hero')).toBe(true) + }) + + client.reset() + + expect(client.getResumeSnapshot()).toBeUndefined() + expect(onResumeSnapshotChange).toHaveBeenLastCalledWith(undefined) + await waitForCondition(() => { + expect(persistence.removeItem).toHaveBeenCalledWith('generation:hero') + expect(store.has('generation:hero')).toBe(false) + }) + }) + + it('skips writes whose snapshot only differs in lastEvent', async () => { + const { persistence } = createMapPersistence() + const client = new GenerationClient({ + id: 'hero', + connection: createMockConnection([ + { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: 'generation:progress', + value: { progress: 10 }, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: 'generation:progress', + value: { progress: 20 }, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + }, + { + type: EventType.RUN_FINISHED, + runId: 'run-1', + threadId: 'thread-1', + finishReason: 'stop', + timestamp: Date.now(), + }, + ]), + persistence, + }) + + await client.generate({ prompt: 'test' }) + await waitForCondition(() => { + // One write for the running state, one for the terminal state — the + // two progress chunks change nothing material. + expect(persistence.setItem).toHaveBeenCalledTimes(2) + }) + }) + + it('revives after dispose when devtools remount (StrictMode replay)', async () => { + const connectSpy = vi.fn(async function* () { + yield { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + } satisfies StreamChunk + yield { + type: EventType.RUN_FINISHED, + runId: 'run-1', + threadId: 'thread-1', + finishReason: 'stop', + timestamp: Date.now(), + } satisfies StreamChunk + }) + const client = new GenerationClient({ + connection: { connect: connectSpy }, + }) + + // StrictMode: mount → cleanup → mount against the same memoized client. + client.mountDevtools() + client.dispose() + client.mountDevtools() + + await client.generate({ prompt: 'test' }) + expect(connectSpy).toHaveBeenCalledTimes(1) + expect(client.getStatus()).toBe('success') + }) + + it('ignores generate() on a client that is disposed and not remounted', async () => { + const connectSpy = vi.fn(async function* () { + yield { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + } satisfies StreamChunk + }) + const client = new GenerationClient({ + connection: { connect: connectSpy }, + }) + + client.mountDevtools() + client.dispose() + + await client.generate({ prompt: 'test' }) + expect(connectSpy).not.toHaveBeenCalled() + }) + }) + + describe('server-driven persistence (persistence: true)', () => { + const completedHydration: GenerationHydrationResult = { + resumeSnapshot: { + schemaVersion: 1, + resumeState: null, + status: 'complete', + result: { id: 'result-1', model: 'image-model' }, + }, + activeRun: null, + } + + function createHydratingConnection(result: GenerationHydrationResult): { + connection: ConnectConnectionAdapter + hydrateGeneration: ReturnType + } { + const hydrateGeneration = vi.fn(async () => result) + return { + connection: { async *connect() {}, hydrateGeneration }, + hydrateGeneration, + } + } + + it('adopts the server snapshot on mount, keyed by threadId, without a local store', async () => { + const { connection, hydrateGeneration } = + createHydratingConnection(completedHydration) + const onResumeSnapshotChange = vi.fn() + const client = new GenerationClient({ + threadId: 'thread-server', + connection, + persistence: true, + onResumeSnapshotChange, + }) + + await waitForCondition(() => { + expect(client.getResumeSnapshot()).toMatchObject({ + status: 'complete', + result: { id: 'result-1' }, + }) + }) + expect(hydrateGeneration).toHaveBeenCalledWith('thread-server') + expect(onResumeSnapshotChange).toHaveBeenCalledWith( + expect.objectContaining({ status: 'complete' }), + ) + }) + + it('repaints result from the server snapshot via reconstructResult (image)', async () => { + const { connection } = createHydratingConnection({ + resumeSnapshot: { + schemaVersion: 1, + resumeState: null, + status: 'complete', + activity: 'image', + result: { + id: 'srv-img', + model: 'test-image', + artifacts: [restoredImageArtifact], + }, + }, + activeRun: null, + }) + const onResumeStateChange = vi.fn() + const client = new GenerationClient({ + threadId: 'thread-img-server', + connection, + persistence: true, + reconstructResult: reconstructImageResult, + onResumeStateChange, + }) + + await waitForCondition(() => { + expect(client.getResult()).toEqual({ + id: 'srv-img', + model: 'test-image', + images: [{ url: '/api/artifacts/artifact-image-1' }], + artifacts: [restoredImageArtifact], + }) + }) + expect(client.getStatus()).toBe('success') + expect(onResumeStateChange).toHaveBeenLastCalledWith(null) + }) + + it('does nothing when the connection exposes no hydrateGeneration', async () => { + const client = new GenerationClient({ + threadId: 'thread-server', + connection: createMockConnection([]), + persistence: true, + }) + + await new Promise((resolve) => setTimeout(resolve, 0)) + expect(client.getResumeSnapshot()).toBeUndefined() + }) + + it('ignores an invalid server snapshot without throwing', async () => { + const { connection } = createHydratingConnection({ + // Structurally invalid status — must be rejected by the client's parser. + resumeSnapshot: { + resumeState: null, + status: 'bogus', + } as unknown as GenerationHydrationResult['resumeSnapshot'], + activeRun: null, + }) + const client = new GenerationClient({ + threadId: 'thread-server', + connection, + persistence: true, + }) + + await new Promise((resolve) => setTimeout(resolve, 0)) + expect(client.getResumeSnapshot()).toBeUndefined() + }) + + it('does not stomp a run that a generate() started before hydration resolves', async () => { + const hydrateGeneration = vi.fn( + async (): Promise => completedHydration, + ) + const connection: ConnectConnectionAdapter = { + async *connect() { + yield { + type: EventType.RUN_STARTED, + runId: 'run-live', + threadId: 'thread-server', + timestamp: Date.now(), + } satisfies StreamChunk + yield { + type: EventType.RUN_FINISHED, + runId: 'run-live', + threadId: 'thread-server', + finishReason: 'stop', + timestamp: Date.now(), + } satisfies StreamChunk + }, + hydrateGeneration, + } + const client = new GenerationClient({ + threadId: 'thread-server', + connection, + persistence: true, + }) + + // Start a live run immediately; it owns the client and hydration backs off. + await client.generate({ prompt: 'test' }) + await new Promise((resolve) => setTimeout(resolve, 0)) + + expect(client.getStatus()).toBe('success') + // The live run's terminal snapshot (from run-live), not the server's. + expect(client.getResumeSnapshot()).toMatchObject({ status: 'complete' }) + expect(client.getResumeSnapshot()?.result?.id).toBeUndefined() + }) + + it('leaves the client-driven adapter path unchanged (writes locally, no hydrateGeneration)', async () => { + const hydrateGeneration = vi.fn() + const store = new Map() + const persistence = { + getItem: vi.fn((key: string) => store.get(key) ?? null), + setItem: vi.fn((key: string, value: unknown) => { + store.set(key, value) + }), + removeItem: vi.fn((key: string) => { + store.delete(key) + }), + } as unknown as GenerationPersistence + const client = new GenerationClient({ + id: 'local-hero', + connection: { + async *connect() { + yield { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + } satisfies StreamChunk + yield { + type: EventType.RUN_FINISHED, + runId: 'run-1', + threadId: 'thread-1', + finishReason: 'stop', + timestamp: Date.now(), + } satisfies StreamChunk + }, + hydrateGeneration, + }, + persistence, + }) + + await client.generate({ prompt: 'test' }) + await waitForCondition(() => { + expect(persistence.setItem).toHaveBeenCalled() + }) + // An adapter is client-driven: the server hydrate path is never used. + expect(hydrateGeneration).not.toHaveBeenCalled() + expect(store.has('generation:local-hero')).toBe(true) + }) + + it('adopts the server snapshot for the video client too', async () => { + const { connection, hydrateGeneration } = + createHydratingConnection(completedHydration) + const client = new VideoGenerationClient({ + threadId: 'thread-video', + connection, + persistence: true, + }) + + await waitForCondition(() => { + expect(client.getResumeSnapshot()).toMatchObject({ status: 'complete' }) + }) + expect(hydrateGeneration).toHaveBeenCalledWith('thread-video') + }) + }) }) diff --git a/packages/ai-client/tests/generation-resume-state.test.ts b/packages/ai-client/tests/generation-resume-state.test.ts new file mode 100644 index 000000000..ab8408d74 --- /dev/null +++ b/packages/ai-client/tests/generation-resume-state.test.ts @@ -0,0 +1,404 @@ +import { describe, expect, it } from 'vitest' +import { EventType } from '@tanstack/ai/client' +import { + GENERATION_EVENTS, + parseGenerationResumeSnapshot, + updateGenerationResumeSnapshot, +} from '../src/generation-types' +import type { PersistedArtifactRef, StreamChunk } from '@tanstack/ai/client' +import type { GenerationResumeSnapshot } from '../src/generation-types' + +const artifactRef: PersistedArtifactRef = { + role: 'output', + artifactId: 'artifact-1', + threadId: 'thread-1', + runId: 'run-1', + name: 'image.png', + mimeType: 'image/png', + size: 1234, + createdAt: '2026-07-06T00:00:00.000Z', + source: { + activity: 'image', + path: 'thread-1/run-1/image.png', + provider: 'test', + model: 'image-model', + mediaType: 'image', + }, +} + +function reduceChunks( + chunks: ReadonlyArray, + initial?: GenerationResumeSnapshot, +): GenerationResumeSnapshot { + let snapshot = initial + for (const chunk of chunks) { + snapshot = updateGenerationResumeSnapshot(snapshot, chunk) + } + if (!snapshot) { + throw new Error('Expected at least one generation event') + } + return snapshot +} + +describe('generation resume state reducer', () => { + it('tracks thread and run from persisted generation events', () => { + const snapshot = reduceChunks([ + { + type: EventType.RUN_STARTED, + threadId: 'thread-1', + runId: 'run-1', + timestamp: 1, + }, + { + type: EventType.CUSTOM, + name: GENERATION_EVENTS.PROGRESS, + value: { progress: 50, message: 'Halfway' }, + threadId: 'thread-1', + runId: 'run-1', + timestamp: 2, + }, + ]) + + expect(snapshot).toMatchObject({ + resumeState: { + threadId: 'thread-1', + runId: 'run-1', + }, + status: 'running', + lastEvent: { + type: EventType.CUSTOM, + name: GENERATION_EVENTS.PROGRESS, + }, + }) + }) + + it('stores generation artifacts as persisted refs only', () => { + const snapshot = reduceChunks([ + { + type: EventType.CUSTOM, + name: GENERATION_EVENTS.ARTIFACTS, + value: [ + artifactRef, + { + b64Json: 'raw-media-bytes', + url: 'data:image/png;base64,raw-media-bytes', + }, + ], + threadId: 'thread-1', + runId: 'run-1', + timestamp: 1, + }, + ]) + + expect(snapshot.pendingArtifacts).toEqual([artifactRef]) + expect(JSON.stringify(snapshot)).not.toContain('raw-media-bytes') + expect(snapshot.activity).toBe('image') + }) + + it('sanitizes artifact refs before storing them in resume snapshots', () => { + const unsafeArtifactRef = { + ...artifactRef, + externalUrl: 'data:image/png;base64,raw-artifact-bytes', + b64Json: 'raw-artifact-bytes', + blob: new Blob(['raw-artifact-bytes']), + url: `https://example.com/${'x'.repeat(4096)}`, + source: { + ...artifactRef.source, + extraRawField: 'raw-artifact-bytes', + }, + } + + const snapshot = reduceChunks([ + { + type: EventType.CUSTOM, + name: GENERATION_EVENTS.ARTIFACTS, + value: [unsafeArtifactRef], + threadId: 'thread-1', + runId: 'run-1', + timestamp: 1, + }, + { + type: EventType.CUSTOM, + name: GENERATION_EVENTS.RESULT, + value: { + id: 'result-1', + artifacts: [unsafeArtifactRef], + }, + threadId: 'thread-1', + runId: 'run-1', + timestamp: 2, + }, + ]) + + expect(snapshot.pendingArtifacts).toEqual([artifactRef]) + expect(snapshot.result?.artifacts).toEqual([artifactRef]) + expect(JSON.stringify(snapshot)).not.toContain('raw-artifact-bytes') + expect(JSON.stringify(snapshot)).not.toContain('b64Json') + expect(JSON.stringify(snapshot)).not.toContain('extraRawField') + }) + + it('stores terminal result metadata without raw generated bytes', () => { + const snapshot = reduceChunks([ + { + type: EventType.CUSTOM, + name: GENERATION_EVENTS.RESULT, + value: { + id: 'result-1', + model: 'image-model', + images: [ + { + b64Json: 'raw-image-bytes', + url: 'data:image/png;base64,raw-image-bytes', + revisedPrompt: 'clean prompt', + }, + ], + artifacts: [artifactRef], + }, + threadId: 'thread-1', + runId: 'run-1', + timestamp: 2, + }, + { + type: EventType.RUN_FINISHED, + threadId: 'thread-1', + runId: 'run-1', + finishReason: 'stop', + timestamp: 3, + }, + ]) + + expect(snapshot.resumeState).toBeNull() + expect(snapshot.status).toBe('complete') + expect(snapshot.result).toMatchObject({ + id: 'result-1', + model: 'image-model', + artifacts: [artifactRef], + }) + expect(JSON.stringify(snapshot)).not.toContain('raw-image-bytes') + expect(JSON.stringify(snapshot)).not.toContain('b64Json') + }) + + it('keeps a durable artifact externalUrl and strips non-durable or oversized ones', () => { + const durableUrl = 'https://cdn.example.com/artifacts/image.png' + const durable = reduceChunks([ + { + type: EventType.CUSTOM, + name: GENERATION_EVENTS.ARTIFACTS, + value: [{ ...artifactRef, externalUrl: durableUrl }], + threadId: 'thread-1', + runId: 'run-1', + timestamp: 1, + }, + ]) + expect(durable.pendingArtifacts?.[0]?.externalUrl).toBe(durableUrl) + + const unsafeUrls = [ + 'data:image/png;base64,raw-image-bytes', + 'blob:https://example.com/raw-image-bytes', + `https://example.com/${'x'.repeat(4096)}`, + 'not a url', + ] + for (const externalUrl of unsafeUrls) { + const snapshot = reduceChunks([ + { + type: EventType.CUSTOM, + name: GENERATION_EVENTS.ARTIFACTS, + value: [{ ...artifactRef, externalUrl }], + threadId: 'thread-1', + runId: 'run-1', + timestamp: 1, + }, + ]) + expect(snapshot.pendingArtifacts).toEqual([artifactRef]) + expect(snapshot.pendingArtifacts?.[0]).not.toHaveProperty('externalUrl') + } + }) + + it('clears resume state and stores lightweight error metadata on terminal errors', () => { + const snapshot = reduceChunks([ + { + type: EventType.RUN_STARTED, + threadId: 'thread-1', + runId: 'run-1', + timestamp: 1, + }, + { + type: EventType.RUN_ERROR, + threadId: 'thread-1', + runId: 'run-1', + message: 'Generation failed', + code: 'provider_error', + error: { message: 'legacy message' }, + timestamp: 2, + }, + ]) + + expect(snapshot).toMatchObject({ + resumeState: null, + status: 'error', + error: { + message: 'Generation failed', + code: 'provider_error', + }, + }) + }) + + it('captures the video job id from video:job:created before any result arrives', () => { + const snapshot = reduceChunks([ + { + type: EventType.RUN_STARTED, + threadId: 'thread-1', + runId: 'run-1', + timestamp: 1, + }, + { + type: EventType.CUSTOM, + name: GENERATION_EVENTS.VIDEO_JOB_CREATED, + value: { jobId: 'job-42' }, + threadId: 'thread-1', + runId: 'run-1', + timestamp: 2, + }, + ]) + + expect(snapshot.status).toBe('running') + expect(snapshot.result?.jobId).toBe('job-42') + }) + + it('merges an initial seed and lets later run events update it', () => { + const seed: GenerationResumeSnapshot = { + resumeState: null, + status: 'error', + error: { message: 'previous failure' }, + pendingArtifacts: [artifactRef], + } + + const midRun = reduceChunks( + [ + { + type: EventType.CUSTOM, + name: GENERATION_EVENTS.PROGRESS, + value: { progress: 10 }, + threadId: 'thread-1', + runId: 'run-2', + timestamp: 1, + }, + ], + seed, + ) + // Without a RUN_STARTED boundary the seed's fields are carried forward. + expect(midRun.resumeState).toEqual({ threadId: 'thread-1', runId: 'run-2' }) + expect(midRun.error).toEqual({ message: 'previous failure' }) + expect(midRun.pendingArtifacts).toEqual([artifactRef]) + }) + + it('drops stale error, result, and artifacts when a new run starts', () => { + const seed: GenerationResumeSnapshot = { + resumeState: null, + status: 'error', + error: { message: 'previous failure' }, + result: { id: 'old-result' }, + pendingArtifacts: [artifactRef], + } + + const snapshot = reduceChunks( + [ + { + type: EventType.RUN_STARTED, + threadId: 'thread-1', + runId: 'run-2', + timestamp: 1, + }, + ], + seed, + ) + + expect(snapshot.status).toBe('running') + expect(snapshot.resumeState).toEqual({ + threadId: 'thread-1', + runId: 'run-2', + }) + expect(snapshot.error).toBeUndefined() + expect(snapshot.result).toBeUndefined() + expect(snapshot.pendingArtifacts).toBeUndefined() + }) +}) + +describe('parseGenerationResumeSnapshot', () => { + it('round-trips a reducer-produced snapshot through JSON, dropping lastEvent', () => { + const produced = reduceChunks([ + { + type: EventType.RUN_STARTED, + threadId: 'thread-1', + runId: 'run-1', + timestamp: 1, + }, + { + type: EventType.CUSTOM, + name: GENERATION_EVENTS.RESULT, + value: { + id: 'result-1', + model: 'image-model', + artifacts: [artifactRef], + }, + threadId: 'thread-1', + runId: 'run-1', + timestamp: 2, + }, + { + type: EventType.RUN_FINISHED, + threadId: 'thread-1', + runId: 'run-1', + finishReason: 'stop', + timestamp: 3, + }, + ]) + + const parsed = parseGenerationResumeSnapshot( + JSON.parse(JSON.stringify(produced)), + ) + expect(parsed).toBeDefined() + expect(parsed?.status).toBe('complete') + expect(parsed?.resumeState).toBeNull() + expect(parsed?.result).toEqual(produced.result) + expect(parsed?.lastEvent).toBeUndefined() + }) + + it('rejects garbage, invalid statuses, malformed resume state, and future schema versions', () => { + expect(parseGenerationResumeSnapshot(undefined)).toBeUndefined() + expect(parseGenerationResumeSnapshot('running')).toBeUndefined() + expect(parseGenerationResumeSnapshot({})).toBeUndefined() + expect( + parseGenerationResumeSnapshot({ resumeState: null, status: 'paused' }), + ).toBeUndefined() + expect( + parseGenerationResumeSnapshot({ + resumeState: { threadId: 'thread-1' }, + status: 'running', + }), + ).toBeUndefined() + expect( + parseGenerationResumeSnapshot({ + schemaVersion: 2, + resumeState: null, + status: 'idle', + }), + ).toBeUndefined() + }) + + it('strips unknown fields and re-validates artifact refs from storage', () => { + const parsed = parseGenerationResumeSnapshot({ + resumeState: { threadId: 'thread-1', runId: 'run-1' }, + status: 'running', + pendingArtifacts: [artifactRef, { b64Json: 'raw-bytes' }], + error: { message: 'boom', code: 'E1', extra: 'dropped' }, + injected: 'dropped', + }) + + expect(parsed).toBeDefined() + expect(parsed?.pendingArtifacts).toEqual([artifactRef]) + expect(parsed?.error).toEqual({ message: 'boom', code: 'E1' }) + expect(JSON.stringify(parsed)).not.toContain('raw-bytes') + expect(JSON.stringify(parsed)).not.toContain('dropped') + }) +}) diff --git a/packages/ai-event-client/src/index.ts b/packages/ai-event-client/src/index.ts index 062faabdd..7465168be 100644 --- a/packages/ai-event-client/src/index.ts +++ b/packages/ai-event-client/src/index.ts @@ -614,6 +614,8 @@ export interface SummarizeUsageEvent extends BaseEventContext { /** Emitted when an image request starts. */ export interface ImageRequestStartedEvent extends BaseEventContext { requestId: string + threadId?: string + runId?: string provider: string model: string prompt: string @@ -630,6 +632,8 @@ export interface ImageRequestStartedEvent extends BaseEventContext { /** Emitted when an image request completes. */ export interface ImageRequestCompletedEvent extends BaseEventContext { requestId: string + threadId?: string + runId?: string provider: string model: string images: Array<{ url?: string; b64Json?: string }> @@ -639,6 +643,8 @@ export interface ImageRequestCompletedEvent extends BaseEventContext { /** Emitted when image usage metrics are available. */ export interface ImageUsageEvent extends BaseEventContext { requestId: string + threadId?: string + runId?: string model: string usage: TokenUsage } @@ -650,6 +656,8 @@ export interface ImageUsageEvent extends BaseEventContext { /** Emitted when a speech request starts. */ export interface SpeechRequestStartedEvent extends BaseEventContext { requestId: string + threadId?: string + runId?: string provider: string model: string text: string @@ -661,6 +669,8 @@ export interface SpeechRequestStartedEvent extends BaseEventContext { /** Emitted when a speech request completes. */ export interface SpeechRequestCompletedEvent extends BaseEventContext { requestId: string + threadId?: string + runId?: string provider: string model: string audio: string @@ -673,6 +683,8 @@ export interface SpeechRequestCompletedEvent extends BaseEventContext { /** Emitted when speech usage metrics are available. */ export interface SpeechUsageEvent extends BaseEventContext { requestId: string + threadId?: string + runId?: string model: string usage: TokenUsage } @@ -684,6 +696,8 @@ export interface SpeechUsageEvent extends BaseEventContext { /** Emitted when a transcription request starts. */ export interface TranscriptionRequestStartedEvent extends BaseEventContext { requestId: string + threadId?: string + runId?: string provider: string model: string language?: string @@ -694,6 +708,8 @@ export interface TranscriptionRequestStartedEvent extends BaseEventContext { /** Emitted when a transcription request completes. */ export interface TranscriptionRequestCompletedEvent extends BaseEventContext { requestId: string + threadId?: string + runId?: string provider: string model: string text: string @@ -704,6 +720,8 @@ export interface TranscriptionRequestCompletedEvent extends BaseEventContext { /** Emitted when transcription usage metrics are available. */ export interface TranscriptionUsageEvent extends BaseEventContext { requestId: string + threadId?: string + runId?: string model: string usage: TokenUsage } @@ -715,6 +733,8 @@ export interface TranscriptionUsageEvent extends BaseEventContext { /** Emitted when an audio generation request starts. */ export interface AudioRequestStartedEvent extends BaseEventContext { requestId: string + threadId?: string + runId?: string provider: string model: string prompt: string @@ -743,6 +763,8 @@ export type AudioRequestCompletedAudio = /** Emitted when an audio generation request completes. */ export interface AudioRequestCompletedEvent extends BaseEventContext { requestId: string + threadId?: string + runId?: string provider: string model: string audio: AudioRequestCompletedAudio @@ -752,6 +774,8 @@ export interface AudioRequestCompletedEvent extends BaseEventContext { /** Emitted when an audio generation request fails. */ export interface AudioRequestErrorEvent extends BaseEventContext { requestId: string + threadId?: string + runId?: string provider: string model: string error: { message: string; name?: string } @@ -761,6 +785,8 @@ export interface AudioRequestErrorEvent extends BaseEventContext { /** Emitted when a speech generation request fails. */ export interface SpeechRequestErrorEvent extends BaseEventContext { requestId: string + threadId?: string + runId?: string provider: string model: string error: { message: string; name?: string } @@ -770,6 +796,8 @@ export interface SpeechRequestErrorEvent extends BaseEventContext { /** Emitted when a transcription request fails. */ export interface TranscriptionRequestErrorEvent extends BaseEventContext { requestId: string + threadId?: string + runId?: string provider: string model: string error: { message: string; name?: string } @@ -779,6 +807,8 @@ export interface TranscriptionRequestErrorEvent extends BaseEventContext { /** Emitted when audio usage metrics are available. */ export interface AudioUsageEvent extends BaseEventContext { requestId: string + threadId?: string + runId?: string model: string usage: TokenUsage } @@ -790,6 +820,8 @@ export interface AudioUsageEvent extends BaseEventContext { /** Emitted when a video request starts. */ export interface VideoRequestStartedEvent extends BaseEventContext { requestId: string + threadId?: string + runId?: string provider: string model: string requestType: 'create' | 'status' | 'url' @@ -802,6 +834,8 @@ export interface VideoRequestStartedEvent extends BaseEventContext { /** Emitted when a video request completes. */ export interface VideoRequestCompletedEvent extends BaseEventContext { requestId: string + threadId?: string + runId?: string provider: string model: string requestType: 'create' | 'status' | 'url' @@ -816,6 +850,8 @@ export interface VideoRequestCompletedEvent extends BaseEventContext { /** Emitted when video usage metrics are available. */ export interface VideoUsageEvent extends BaseEventContext { requestId: string + threadId?: string + runId?: string model: string usage: TokenUsage } diff --git a/packages/ai-persistence/package.json b/packages/ai-persistence/package.json index f86e0563d..670a52c35 100644 --- a/packages/ai-persistence/package.json +++ b/packages/ai-persistence/package.json @@ -48,6 +48,9 @@ "test:types": "tsc", "test:oxlint": "oxlint src --type-aware" }, + "dependencies": { + "@tanstack/ai-utils": "workspace:*" + }, "peerDependencies": { "@tanstack/ai": "workspace:*", "vitest": "^4.1.10" diff --git a/packages/ai-persistence/skills/ai-persistence/SKILL.md b/packages/ai-persistence/skills/ai-persistence/SKILL.md index 7463890fe..28c94f48e 100644 --- a/packages/ai-persistence/skills/ai-persistence/SKILL.md +++ b/packages/ai-persistence/skills/ai-persistence/SKILL.md @@ -42,14 +42,25 @@ drives them, an in-memory reference backend, and a conformance testkit. It does stores against whatever you already run — Postgres, SQLite, D1, Mongo — and hand the result to `withPersistence`. The core never inspects your tables. -| Ships in the package | What it is | -| --------------------------------------------------------------------------- | ------------------------------------------------------ | -| `MessageStore` / `RunStore` / `InterruptStore` / `MetadataStore` | The four state contracts | -| `withPersistence` / `withGenerationPersistence` | Chat + generation middleware | -| `memoryPersistence()` | In-process reference backend (dev, tests) | -| `reconstructChat` | Server hydrate route helper | -| `LockStore` / `withLocks` / `InMemoryLockStore` (from `@tanstack/ai/locks`) | Coordination, **not** this package — see ai-core/locks | -| `@tanstack/ai-persistence/testkit` | `runPersistenceConformance` compatibility gate | +| Ships in the package | What it is | +| --------------------------------------------------------------------------- | ----------------------------------------------------------- | +| `MessageStore` / `RunStore` / `InterruptStore` / `MetadataStore` | The four **chat** state contracts | +| `GenerationJobStore` / `ArtifactStore` / `BlobStore` | The **generation** contracts (job lifecycle + bytes) | +| `withPersistence` / `withGenerationPersistence` | Chat + generation middleware | +| `memoryPersistence()` | In-process reference backend, all seven stores (dev, tests) | +| `reconstructChat` / `reconstructGeneration` | Server hydrate route helpers (chat / generation) | +| `retrieveArtifact` / `retrieveBlob` / `artifactBlobKey` | Serve persisted generation-media bytes back | +| `LockStore` / `withLocks` / `InMemoryLockStore` (from `@tanstack/ai/locks`) | Coordination, **not** this package — see ai-core/locks | +| `@tanstack/ai-persistence/testkit` | `runPersistenceConformance` gate (chat state stores) | + +**Chat vs generation stores.** Chat persistence keys on `threadId` and uses +`messages` + optional `runs` / `interrupts` / `metadata`. Generation persistence +keys on `jobId` and uses `jobs` (required by `withGenerationPersistence`) plus an +optional `artifacts` + `blobs` **pair** — provide both or neither — to store the +generated media bytes at blob key `artifacts//`. `threadId` on +a generation is only an optional _link_ to a chat, never the job's identity. To +build the R2/D1-backed byte stores for a Worker, see +**ai-persistence/build-cloudflare-artifact-store**. ## Sub-skills @@ -64,12 +75,13 @@ Adding persistence to an app? Pick the recipe that matches what it already runs — each one writes a single `chat-persistence.ts` against the app's existing database client and schema: -| The app runs... | Read | -| ------------------------------------------------ | ------------------------------------------------ | -| Drizzle ORM (SQLite / Postgres / MySQL) | ai-persistence/build-drizzle-adapter/SKILL.md | -| Prisma | ai-persistence/build-prisma-adapter/SKILL.md | -| Cloudflare Workers + D1 (± Durable Object locks) | ai-persistence/build-cloudflare-adapter/SKILL.md | -| Anything else — raw `pg`, Kysely, SQLite, Mongo | ai-persistence/build-custom-adapter/SKILL.md | +| The app runs... | Read | +| ---------------------------------------------------- | ------------------------------------------------------- | +| Drizzle ORM (SQLite / Postgres / MySQL) | ai-persistence/build-drizzle-adapter/SKILL.md | +| Prisma | ai-persistence/build-prisma-adapter/SKILL.md | +| Cloudflare Workers + D1 (± Durable Object locks) | ai-persistence/build-cloudflare-adapter/SKILL.md | +| Cloudflare Workers + R2/D1 for generated media bytes | ai-persistence/build-cloudflare-artifact-store/SKILL.md | +| Anything else — raw `pg`, Kysely, SQLite, Mongo | ai-persistence/build-custom-adapter/SKILL.md | ## State persistence has two halves diff --git a/packages/ai-persistence/skills/ai-persistence/build-cloudflare-artifact-store/SKILL.md b/packages/ai-persistence/skills/ai-persistence/build-cloudflare-artifact-store/SKILL.md new file mode 100644 index 000000000..bc17818a2 --- /dev/null +++ b/packages/ai-persistence/skills/ai-persistence/build-cloudflare-artifact-store/SKILL.md @@ -0,0 +1,406 @@ +--- +name: ai-persistence/build-cloudflare-artifact-store +description: Use when a Cloudflare Worker needs durable byte storage for TanStack AI generated media (images, audio, video, transcripts) — writes a BlobStore backed by R2 and an ArtifactStore backed by D1 (or KV), composes them onto the generation persistence so withGenerationPersistence persists artifact bytes, and serves them back from a Worker GET route. Includes one-line sketches for S3, GCS, Vercel Blob, Supabase, and a dev filesystem BlobStore. +--- + +# Cloudflare Artifact + Blob Store + +`withGenerationPersistence(persistence)` needs only `stores.jobs` to track a +generation's lifecycle. Add `stores.artifacts` (metadata) **and** `stores.blobs` +(the bytes) — both, or neither — and the middleware also persists the generated +media: image/audio/TTS/video/transcription bytes land at blob key +`artifacts//`, with an `ArtifactRecord` row describing each. + +The deliverable is **one file in the Worker** — e.g. +`src/lib/generation-persistence.ts` — exporting a factory that builds an +`AIPersistence` from the request's R2 + D1 bindings, plus a GET route that serves +artifact bytes with `retrieveArtifact` / `retrieveBlob`. + +Read the sibling **ai-persistence/build-cloudflare-adapter** skill for the +per-request-binding rule, `wrangler` config shape, D1 migration workflow, and the +chat (jobs/messages) side. This skill covers only the two byte-storage stores and +how to compose them. + +## The two contracts, verbatim + +Both come from `@tanstack/ai-persistence`. `defineBlobStore` / `defineArtifactStore` +type an object literal inline (autocomplete + contract checking, no separate +annotation). + +```ts +// BlobStore — the byte layer. R2 backs it. +interface BlobStore { + put( + key: string, + body: BlobBody, + options?: BlobPutOptions, + ): Promise + get(key: string): Promise // metadata + byte accessors + head(key: string): Promise // metadata only + delete(key: string): Promise // no-op if absent + list(options?: BlobListOptions): Promise +} + +// ArtifactStore — the metadata layer. D1 (or KV) backs it. +interface ArtifactStore { + save(record: ArtifactRecord): Promise // insert or overwrite + get(artifactId: string): Promise + list(runId: string): Promise> // [] when none + delete?(artifactId: string): Promise // OPTIONAL + deleteForRun?(runId: string): Promise // OPTIONAL +} +``` + +`BlobBody` is `ReadableStream | ArrayBuffer | ArrayBufferView | +string | Blob` — which is exactly what `R2Bucket.put` accepts, so the body flows +straight through with no conversion. `BlobPutOptions` is +`{ contentType?, customMetadata? }`; `BlobListOptions` is +`{ prefix?, cursor?, limit? }`; `BlobListPage` is +`{ objects: BlobRecord[], cursor?, truncated? }`. + +## 1. BlobStore backed by R2 + +`R2Object` carries `size`, `etag`, `httpMetadata.contentType`, `customMetadata`, +and `uploaded` (a `Date`). `BlobRecord` wants `createdAt` / `updatedAt` as epoch +ms — R2 tracks only the single `uploaded` instant, so map it to both. `get` / +`head` are the byte-body vs metadata-only split; `R2ObjectBody` already exposes +`body`, `arrayBuffer()`, and `text()`, so a `BlobObject` is essentially the R2 +object plus the mapped metadata. + +```ts ignore +import { defineBlobStore } from '@tanstack/ai-persistence' +import type { BlobObject, BlobRecord } from '@tanstack/ai-persistence' + +function toRecord(obj: R2Object): BlobRecord { + const uploaded = obj.uploaded.getTime() + return { + key: obj.key, + size: obj.size, + etag: obj.etag, + ...(obj.httpMetadata?.contentType + ? { contentType: obj.httpMetadata.contentType } + : {}), + ...(obj.customMetadata ? { customMetadata: obj.customMetadata } : {}), + createdAt: uploaded, + updatedAt: uploaded, + } +} + +export function r2BlobStore(bucket: R2Bucket) { + return defineBlobStore({ + async put(key, body, options) { + const obj = await bucket.put(key, body, { + ...(options?.contentType + ? { httpMetadata: { contentType: options.contentType } } + : {}), + ...(options?.customMetadata + ? { customMetadata: options.customMetadata } + : {}), + }) + // R2.put returns null only when an onlyIf precondition fails — not used here. + if (!obj) throw new Error(`R2 put failed for ${key}`) + return toRecord(obj) + }, + + async get(key): Promise { + const obj = await bucket.get(key) + if (!obj) return null + return { + ...toRecord(obj), + body: obj.body, + arrayBuffer: () => obj.arrayBuffer(), + text: () => obj.text(), + } + }, + + async head(key) { + const obj = await bucket.head(key) + return obj ? toRecord(obj) : null + }, + + async delete(key) { + await bucket.delete(key) + }, + + async list(options) { + const page = await bucket.list({ + ...(options?.prefix !== undefined ? { prefix: options.prefix } : {}), + ...(options?.cursor !== undefined ? { cursor: options.cursor } : {}), + ...(options?.limit !== undefined ? { limit: options.limit } : {}), + // R2 omits httpMetadata/customMetadata from list rows unless asked. + include: ['httpMetadata', 'customMetadata'], + }) + return { + objects: page.objects.map(toRecord), + ...(page.truncated ? { cursor: page.cursor, truncated: true } : {}), + } + }, + }) +} +``` + +Invariants that matter (asserted by the conformance testkit): + +- `get` / `head` return `null` for a missing key; `delete` is a silent no-op. +- `put` **overwrites** an existing key. +- `list` filters by `prefix` literally (R2 prefix is a literal byte prefix — no + glob), returns keys in ascending order, and pages via the opaque `cursor` when + `truncated`. R2's own cursor is opaque and satisfies this directly. `limit: 0` + must yield an empty, untruncated page — R2 treats `limit: 0` as "use the + default", so special-case it: `if (options?.limit === 0) return { objects: [] }`. + +## 2. ArtifactStore backed by D1 + +`ArtifactRecord` is `{ artifactId, runId, threadId, name, mimeType, size, +externalUrl?, createdAt }` (`createdAt` epoch ms). One flat table, keyed by +`artifact_id`, indexed by `run_id` for `list`. + +```sql +CREATE TABLE IF NOT EXISTS generation_artifacts ( + artifact_id text PRIMARY KEY NOT NULL, + run_id text NOT NULL, + thread_id text NOT NULL, + name text NOT NULL, + mime_type text NOT NULL, + size integer NOT NULL, + external_url text, + created_at integer NOT NULL +); +CREATE INDEX IF NOT EXISTS generation_artifacts_run ON generation_artifacts (run_id); +``` + +```ts ignore +import { defineArtifactStore } from '@tanstack/ai-persistence' +import type { ArtifactRecord } from '@tanstack/ai-persistence' + +interface ArtifactRow { + artifact_id: string + run_id: string + thread_id: string + name: string + mime_type: string + size: number + external_url: string | null + created_at: number +} + +function fromRow(row: ArtifactRow): ArtifactRecord { + return { + artifactId: row.artifact_id, + runId: row.run_id, + threadId: row.thread_id, + name: row.name, + mimeType: row.mime_type, + size: row.size, + ...(row.external_url != null ? { externalUrl: row.external_url } : {}), + createdAt: row.created_at, + } +} + +export function d1ArtifactStore(db: D1Database) { + return defineArtifactStore({ + async save(record) { + // Insert or overwrite (artifact ids are unique). + await db + .prepare( + `INSERT INTO generation_artifacts + (artifact_id, run_id, thread_id, name, mime_type, size, external_url, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(artifact_id) DO UPDATE SET + run_id = excluded.run_id, thread_id = excluded.thread_id, + name = excluded.name, mime_type = excluded.mime_type, + size = excluded.size, external_url = excluded.external_url, + created_at = excluded.created_at`, + ) + .bind( + record.artifactId, + record.runId, + record.threadId, + record.name, + record.mimeType, + record.size, + record.externalUrl ?? null, + record.createdAt, + ) + .run() + }, + + async get(artifactId) { + const row = await db + .prepare(`SELECT * FROM generation_artifacts WHERE artifact_id = ?`) + .bind(artifactId) + .first() + return row ? fromRow(row) : null + }, + + async list(runId) { + const { results } = await db + .prepare(`SELECT * FROM generation_artifacts WHERE run_id = ?`) + .bind(runId) + .all() + return results.map(fromRow) + }, + + async delete(artifactId) { + await db + .prepare(`DELETE FROM generation_artifacts WHERE artifact_id = ?`) + .bind(artifactId) + .run() + }, + + async deleteForRun(runId) { + await db + .prepare(`DELETE FROM generation_artifacts WHERE run_id = ?`) + .bind(runId) + .run() + }, + }) +} +``` + +Omitting `external_url` from the record when the column is `NULL` keeps records +comparing cleanly against the reference in-memory store. **KV alternative:** if +you have no D1, back `save`/`get` with `KV.put(artifactId, JSON.stringify(record))` +/ `KV.get(artifactId, 'json')`, and maintain a `run:` index key (a JSON +array of artifact ids) for `list` — KV has no query, so `list` needs that +secondary index. + +## 3. Compose and wire + +Bindings are per-request on Workers, so export a **factory**. Combine the byte +stores with a jobs store (and, if this Worker also does chat, the chat stores). +Either build the whole `AIPersistence` with `defineAIPersistence`, or layer the +artifact stores onto an existing chat persistence with `composePersistence`: + +```ts ignore +import { + defineAIPersistence, + composePersistence, + withGenerationPersistence, +} from '@tanstack/ai-persistence' +import { r2BlobStore } from './r2-blob-store' +import { d1ArtifactStore } from './d1-artifact-store' +import { d1GenerationJobStore } from './d1-job-store' // your GenerationJobStore + +/** Call inside a request handler — bindings are not available at module scope. */ +export function generationPersistence(env: Env) { + return defineAIPersistence({ + stores: { + jobs: d1GenerationJobStore(env.DB), + artifacts: d1ArtifactStore(env.DB), + blobs: r2BlobStore(env.ARTIFACTS_BUCKET), + }, + }) +} + +// …or add bytes to a persistence that already has chat + jobs: +// composePersistence(chatAndJobsPersistence(env), { +// overrides: { +// artifacts: d1ArtifactStore(env.DB), +// blobs: r2BlobStore(env.ARTIFACTS_BUCKET), +// }, +// }) +``` + +`withGenerationPersistence` throws if exactly one of `artifacts` / `blobs` is +present — provide both or neither. Wire it as generation middleware: + +```ts ignore +import { generateImage, toServerSentEventsResponse } from '@tanstack/ai' +import { openaiImage } from '@tanstack/ai-openai' +import { withGenerationPersistence } from '@tanstack/ai-persistence' +import { generationPersistence } from './lib/generation-persistence' + +export default { + async fetch(request: Request, env: Env) { + const { prompt, threadId } = await request.json() + const stream = generateImage({ + adapter: openaiImage('gpt-image-1'), + prompt, + threadId, // optional link recorded on the job + artifacts + stream: true, + middleware: [withGenerationPersistence(generationPersistence(env))], + }) + return toServerSentEventsResponse(stream) + }, +} +``` + +## 4. Serve the bytes back + +A GET route resolves an `artifactId` to its record and its stored bytes. +`retrieveArtifact` returns the `ArtifactRecord` (or `null` → 404); +`retrieveBlob` returns the `BlobObject` (metadata + a streamable `body`). Both +key off `artifactBlobKey({ runId, artifactId })` internally, so you never build +the key yourself. + +```ts ignore +import { retrieveArtifact, retrieveBlob } from '@tanstack/ai-persistence' +import { generationPersistence } from './lib/generation-persistence' + +export async function GET(request: Request, env: Env) { + const artifactId = new URL(request.url).searchParams.get('id') ?? '' + const persistence = generationPersistence(env) + + // Authorize before serving — derive the owner from the session, never trust + // a client-supplied id. (The record carries runId/threadId to check against.) + const record = await retrieveArtifact(persistence, artifactId) + if (!record) return new Response('Not found', { status: 404 }) + + const blob = await retrieveBlob(persistence, record) // pass the record: no 2nd lookup + if (!blob?.body) return new Response('Not found', { status: 404 }) + + return new Response(blob.body, { + headers: { + 'content-type': record.mimeType, + 'content-length': String(record.size), + }, + }) +} +``` + +To hydrate a **server-driven generation client** (`persistence: true` + a stable +`threadId`) on mount, also expose `reconstructGeneration(persistence, request)` +on a GET that reads `?threadId=` / `?jobId=` — see `ai-core/client-persistence`. + +## Other backends — the contract is tiny, here's how each maps + +`BlobStore` is five methods over an object store. Any of these backs it; swap the +factory, keep everything else. `put` maps to the SDK's upload, `get` to a +download that exposes `body`/`arrayBuffer`/`text`, `head` to a metadata fetch, +`delete` to a delete, `list` to a prefixed, cursor-paged list. + +| Backend | npm | One-line sketch | +| ------------------------- | ----------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| **AWS S3** | `@aws-sdk/client-s3` | `put`→`PutObjectCommand`; `get`→`GetObjectCommand` (`Body` is a stream → `body`, `.transformToByteArray()`/`.transformToString()`); `head`→`HeadObjectCommand`; `delete`→`DeleteObjectCommand`; `list`→`ListObjectsV2Command` (`Prefix`, `ContinuationToken`↔`cursor`, `MaxKeys`↔`limit`, `IsTruncated`↔`truncated`). | +| **Google Cloud Storage** | `@google-cloud/storage` | `bucket.file(key)`: `put`→`.save(body, { contentType, metadata })`; `get`→`.createReadStream()` for `body` + `.download()` for bytes; `head`→`.getMetadata()`; `delete`→`.delete({ ignoreNotFound: true })`; `list`→`bucket.getFiles({ prefix, maxResults, pageToken })`. | +| **Vercel Blob** | `@vercel/blob` | `put`→`put(key, body, { access: 'public', contentType })`; `get`→`fetch(head(key).url)` (stream `res.body`); `head`→`head(key)` (returns `null`→catch as absent); `delete`→`del(key)`; `list`→`list({ prefix, cursor, limit })` (`hasMore`↔`truncated`). | +| **Supabase Storage** | `@supabase/supabase-js` | `storage.from(bucket)`: `put`→`.upload(key, body, { contentType, upsert: true })`; `get`→`.download(key)` (returns a `Blob` → `body`/`arrayBuffer`/`text`); `head`→`.info(key)` or list-one; `delete`→`.remove([key])`; `list`→`.list(prefix, { limit })` (offset/limit paging → synthesize a `cursor`). | +| **Filesystem (dev only)** | `node:fs/promises` | Root each key under a dir: `put`→`mkdir(dirname, { recursive: true })` + `writeFile`; `get`→`createReadStream` for `body` + `readFile`; `head`→`stat` (`size`, `mtimeMs`→`updatedAt`); `delete`→`rm(path, { force: true })`; `list`→recursive `readdir` filtered by prefix, sorted, sliced by `limit`, cursor = last key. Not for production — no concurrency guarantees. | + +For each: `contentType` and `customMetadata` ride the SDK's own metadata fields; +`BlobRecord.createdAt`/`updatedAt` come from the object's stored timestamps +(epoch ms); return `null` from `get`/`head` on a not-found rather than throwing. + +## Verify + +The shared `runPersistenceConformance` testkit currently covers the four **state** +stores only (`messages`, `runs`, `interrupts`, `metadata`) — it does **not** +exercise `blobs`, `artifacts`, or `jobs`. So the byte stores need their **own** +tests. Run them against a Miniflare R2 + D1 binding with the migration applied, +reset between runs (see **ai-persistence/build-cloudflare-adapter** for the +`cloudflare:test` harness pattern). Assert against the reference in-memory stores +from `memoryPersistence()` as the oracle, covering at minimum: + +- `put` then `get` round-trips bytes and metadata; `get`/`head` return `null` for + a missing key; `delete` is a silent no-op on an absent key. +- `put` overwrites an existing key (and its `contentType`/`customMetadata`). +- `list` filters by `prefix`, returns ascending keys, pages correctly through the + `cursor` when `truncated`, and returns an empty untruncated page for `limit: 0`. +- The `ArtifactStore`: `save` is insert-or-overwrite, `get` returns `null` when + absent, `list(runId)` returns `[]` for an unknown run, and (if implemented) + `deleteForRun` removes exactly that run's rows. + +An end-to-end check is the strongest signal: run `generateImage` through +`withGenerationPersistence(generationPersistence(env))`, then confirm the blob +exists at `artifacts//` and `retrieveBlob` streams it back. diff --git a/packages/ai-persistence/src/index.ts b/packages/ai-persistence/src/index.ts index 9007aa83b..41964440d 100644 --- a/packages/ai-persistence/src/index.ts +++ b/packages/ai-persistence/src/index.ts @@ -6,6 +6,9 @@ export { defineRunStore, defineInterruptStore, defineMetadataStore, + defineGenerationJobStore, + defineArtifactStore, + defineBlobStore, } from './types' export type { MessageStore, @@ -23,6 +26,20 @@ export type { ChatTranscriptPersistence, ChatPersistence, ChatWithInterruptsPersistence, + // Generation job store contract + GenerationJobStatus, + GenerationJobRecord, + GenerationJobStore, + // Generation artifact + blob store contracts + ArtifactRecord, + ArtifactStore, + BlobBody, + BlobRecord, + BlobObject, + BlobListPage, + BlobPutOptions, + BlobListOptions, + BlobStore, AIPersistence, AIPersistenceOverrides, ComposedAIPersistenceStores, @@ -33,14 +50,37 @@ export type { // AIPersistenceStores is intentionally NOT re-exported — use a named chat // shape or AIPersistence<{ messages: MessageStore, … }>. -// Middleware (state only — locks live in @tanstack/ai as withLocks) +// Core artifact wire types (re-exported for convenience) +export type { + PersistedArtifactActivity, + PersistedArtifactRef, + PersistedArtifactRole, +} from '@tanstack/ai' + +// Middleware (chat state — locks live in @tanstack/ai as withLocks) export { withPersistence, withGenerationPersistence } from './middleware' +export type { + WithPersistenceOptions, + GenerationArtifactDescriptor, + GenerationArtifactExtractionInput, + GenerationArtifactNameInput, +} from './middleware' // Server helper: rehydrate a thread's messages for a client load export { reconstructChat } from './reconstruct' export type { ReconstructChatOptions } from './reconstruct' -// Reference in-memory implementation (state stores only) +// Server helper: rehydrate the last generation job for a client load +export { reconstructGeneration } from './reconstruct-generation' +export type { + ReconstructedGeneration, + ReconstructGenerationOptions, +} from './reconstruct-generation' + +// Server helpers: retrieve a persisted generation artifact + its bytes +export { retrieveArtifact, retrieveBlob, artifactBlobKey } from './retrieve' + +// Reference in-memory implementation export { memoryPersistence } from './memory' // Interrupt controller diff --git a/packages/ai-persistence/src/memory.ts b/packages/ai-persistence/src/memory.ts index 70f8368f1..905ba29d1 100644 --- a/packages/ai-persistence/src/memory.ts +++ b/packages/ai-persistence/src/memory.ts @@ -1,7 +1,15 @@ import { defineAIPersistence } from './types' import type { ModelMessage } from '@tanstack/ai' import type { - ChatPersistence, + ArtifactRecord, + ArtifactStore, + BlobBody, + BlobListOptions, + BlobObject, + BlobRecord, + BlobStore, + GenerationJobRecord, + GenerationJobStore, InterruptRecord, InterruptStore, MessageStore, @@ -61,6 +69,52 @@ class MemoryRunStore implements RunStore { } } +class MemoryGenerationJobStore implements GenerationJobStore { + private readonly jobs = new Map() + createOrResume( + input: Pick< + GenerationJobRecord, + 'jobId' | 'activity' | 'provider' | 'model' | 'startedAt' + > & { threadId?: string; status?: GenerationJobRecord['status'] }, + ): Promise { + const existing = this.jobs.get(input.jobId) + if (existing) return Promise.resolve(existing) + const record: GenerationJobRecord = { + jobId: input.jobId, + activity: input.activity, + provider: input.provider, + model: input.model, + status: input.status ?? 'running', + startedAt: input.startedAt, + ...(input.threadId !== undefined ? { threadId: input.threadId } : {}), + } + this.jobs.set(record.jobId, record) + return Promise.resolve(record) + } + update( + jobId: string, + patch: Partial< + Pick< + GenerationJobRecord, + 'status' | 'finishedAt' | 'error' | 'result' | 'artifacts' | 'usage' + > + >, + ): Promise { + const existing = this.jobs.get(jobId) + if (existing) this.jobs.set(jobId, { ...existing, ...patch }) + return Promise.resolve() + } + get(jobId: string): Promise { + return Promise.resolve(this.jobs.get(jobId) ?? null) + } + findLatestForThread(threadId: string): Promise { + const linked = [...this.jobs.values()] + .filter((job) => job.threadId === threadId) + .sort((a, b) => b.startedAt - a.startedAt) + return Promise.resolve(linked[0] ?? null) + } +} + function byRequestedAt(a: InterruptRecord, b: InterruptRecord): number { return a.requestedAt - b.requestedAt } @@ -169,20 +223,226 @@ class MemoryMetadataStore implements MetadataStore { } } +class MemoryArtifactStore implements ArtifactStore { + private readonly artifacts = new Map() + save(record: ArtifactRecord): Promise { + this.artifacts.set(record.artifactId, { ...record }) + return Promise.resolve() + } + get(artifactId: string): Promise { + return Promise.resolve(this.artifacts.get(artifactId) ?? null) + } + list(runId: string): Promise> { + return Promise.resolve( + [...this.artifacts.values()].filter((a) => a.runId === runId), + ) + } + delete(artifactId: string): Promise { + this.artifacts.delete(artifactId) + return Promise.resolve() + } + deleteForRun(runId: string): Promise { + for (const artifact of this.artifacts.values()) { + if (artifact.runId === runId) this.artifacts.delete(artifact.artifactId) + } + return Promise.resolve() + } +} + +interface MemoryBlobEntry { + record: BlobRecord + bytes: Uint8Array +} + +const textEncoder = new TextEncoder() +const textDecoder = new TextDecoder() + +function copyBytes(bytes: Uint8Array): Uint8Array { + return new Uint8Array(bytes) +} + +function bytesToArrayBuffer(bytes: Uint8Array): ArrayBuffer { + const buffer = new ArrayBuffer(bytes.byteLength) + new Uint8Array(buffer).set(bytes) + return buffer +} + +async function bytesFromStream( + stream: ReadableStream, +): Promise { + const reader = stream.getReader() + const chunks: Array = [] + let total = 0 + try { + // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition + while (true) { + const { done, value } = await reader.read() + if (done) break + chunks.push(copyBytes(value)) + total += value.byteLength + } + } finally { + reader.releaseLock() + } + + const bytes = new Uint8Array(total) + let offset = 0 + for (const chunk of chunks) { + bytes.set(chunk, offset) + offset += chunk.byteLength + } + return bytes +} + +async function bytesFromBlobBody(body: BlobBody): Promise { + if (typeof body === 'string') { + return textEncoder.encode(body) + } + if (body instanceof ArrayBuffer) { + return new Uint8Array(body.slice(0)) + } + if (ArrayBuffer.isView(body)) { + return copyBytes( + new Uint8Array(body.buffer, body.byteOffset, body.byteLength), + ) + } + if (typeof Blob !== 'undefined' && body instanceof Blob) { + return new Uint8Array(await body.arrayBuffer()) + } + if (typeof ReadableStream !== 'undefined' && body instanceof ReadableStream) { + return bytesFromStream(body) + } + throw new TypeError('Unsupported blob body.') +} + +function blobRecordSnapshot(record: BlobRecord): BlobRecord { + return { + ...record, + ...(record.customMetadata + ? { customMetadata: { ...record.customMetadata } } + : {}), + } +} + +function blobObject(record: BlobRecord, bytes: Uint8Array): BlobObject { + const copied = copyBytes(bytes) + return { + ...blobRecordSnapshot(record), + body: new ReadableStream({ + start(controller) { + controller.enqueue(copyBytes(copied)) + controller.close() + }, + }), + arrayBuffer: () => Promise.resolve(bytesToArrayBuffer(copied)), + text: () => Promise.resolve(textDecoder.decode(copied)), + } +} + +class MemoryBlobStore implements BlobStore { + private readonly blobs = new Map() + private nextEtag = 1 + + async put( + key: string, + body: BlobBody, + options?: { + contentType?: string + customMetadata?: Record + }, + ): Promise { + const bytes = await bytesFromBlobBody(body) + const existing = this.blobs.get(key) + const now = Date.now() + const record: BlobRecord = { + key, + size: bytes.byteLength, + etag: String(this.nextEtag++), + contentType: + options?.contentType ?? + (typeof Blob !== 'undefined' && body instanceof Blob + ? body.type || undefined + : undefined), + customMetadata: options?.customMetadata + ? { ...options.customMetadata } + : undefined, + createdAt: existing?.record.createdAt ?? now, + updatedAt: now, + } + this.blobs.set(key, { record, bytes: copyBytes(bytes) }) + return blobRecordSnapshot(record) + } + + get(key: string): Promise { + const entry = this.blobs.get(key) + return Promise.resolve(entry ? blobObject(entry.record, entry.bytes) : null) + } + + head(key: string): Promise { + const entry = this.blobs.get(key) + return Promise.resolve(entry ? blobRecordSnapshot(entry.record) : null) + } + + delete(key: string): Promise { + this.blobs.delete(key) + return Promise.resolve() + } + + list(options?: BlobListOptions): Promise<{ + objects: Array + cursor?: string + truncated?: boolean + }> { + const limit = options?.limit + if (limit === 0) { + return Promise.resolve({ objects: [], truncated: false }) + } + const keys = [...this.blobs.keys()] + .filter((key) => key.startsWith(options?.prefix ?? '')) + .filter((key) => options?.cursor === undefined || key > options.cursor) + .sort() + const pageKeys = limit === undefined ? keys : keys.slice(0, limit) + const objects = pageKeys.map((key) => { + const blob = this.blobs.get(key) + if (blob === undefined) { + throw new Error(`Missing blob for listed key: ${key}`) + } + return blobRecordSnapshot(blob.record) + }) + const truncated = limit !== undefined && keys.length > limit + return Promise.resolve({ + objects, + ...(truncated ? { cursor: pageKeys.at(-1), truncated } : {}), + }) + } +} + +interface MemoryPersistenceStores { + messages: MessageStore + runs: RunStore + jobs: GenerationJobStore + interrupts: InterruptStore + metadata: MetadataStore + artifacts: ArtifactStore + blobs: BlobStore +} + /** - * In-process reference backend for **chat** state stores. + * In-process reference backend for the full state + generation store set. * - * Returns {@link ChatPersistence} (messages + runs + interrupts + metadata). + * Returns messages + runs + jobs + interrupts + metadata + artifacts + blobs. * Locks are not included — use `InMemoryLockStore` + `withLocks` from * `@tanstack/ai` when a test or single-process app needs coordination. */ -export function memoryPersistence(): ChatPersistence { - return defineAIPersistence({ - stores: { - messages: new MemoryMessageStore(), - runs: new MemoryRunStore(), - interrupts: new MemoryInterruptStore(), - metadata: new MemoryMetadataStore(), - }, - }) +export function memoryPersistence() { + const stores: MemoryPersistenceStores = { + messages: new MemoryMessageStore(), + runs: new MemoryRunStore(), + jobs: new MemoryGenerationJobStore(), + interrupts: new MemoryInterruptStore(), + metadata: new MemoryMetadataStore(), + artifacts: new MemoryArtifactStore(), + blobs: new MemoryBlobStore(), + } + return defineAIPersistence({ stores }) } diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index cd86cfab9..c1843f512 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -1,4 +1,5 @@ import { defineChatMiddleware } from '@tanstack/ai' +import { base64ToUint8Array } from '@tanstack/ai-utils' import { InterruptsCapability, PersistenceCapability, @@ -23,6 +24,9 @@ import type { GenerationMiddleware, GenerationMiddlewareContext, ModelMessage, + PersistedArtifactActivity, + PersistedArtifactRef, + PersistedArtifactRole, RunAgentResumeItem, StreamChunk, ToolApprovalResolution, @@ -31,10 +35,64 @@ import type { import type { AIPersistence, AIPersistenceStores, + ArtifactRecord, + BlobBody, ChatTranscriptStores, InterruptRecord, RunStore, } from './types' +import { artifactBlobKey } from './retrieve' + +export interface WithPersistenceOptions { + extractArtifacts?: ( + input: GenerationArtifactExtractionInput, + ) => + | Array + | Promise> + nameArtifact?: (input: GenerationArtifactNameInput) => string + /** + * Map a freshly-persisted artifact ref to the durable app-origin URL that + * serves its bytes (your `GET` route around `retrieveArtifact` / + * `retrieveBlob`). The returned URL is stamped onto `ref.url` and written into + * the result's media field, so both the live and the restored result render + * durable media from your own origin instead of the provider's expiring link. + * Return `undefined` to leave a ref without a durable URL. + */ + artifactUrl?: (ref: PersistedArtifactRef) => string | undefined +} + +export interface GenerationArtifactDescriptor { + role: PersistedArtifactRole + path: string + mediaType?: PersistedArtifactRef['source']['mediaType'] + mimeType?: string + bytes?: BlobBody + url?: string + json?: unknown + name?: string + jobId?: string + expiresAt?: string | Date +} + +export interface GenerationArtifactExtractionInput { + activity: PersistedArtifactActivity + provider: string + model: string + threadId: string + runId: string + inputs: unknown + result: unknown +} + +export interface GenerationArtifactNameInput { + descriptor: GenerationArtifactDescriptor + activity: PersistedArtifactActivity + provider: string + model: string + threadId: string + runId: string + index: number +} interface RunStateEntry { merged: boolean @@ -259,18 +317,529 @@ function interruptPayload(interrupt: unknown): Record { : { value: interrupt } } +// --------------------------------------------------------------------------- +// Generation artifact extraction / persistence +// --------------------------------------------------------------------------- + +function isArtifactRef(value: unknown): value is PersistedArtifactRef { + const record = objectValue(value) + return !!record && typeof record.artifactId === 'string' +} + +function mediaActivity( + activity: GenerationMiddlewareContext['activity'], +): PersistedArtifactActivity | undefined { + return activity === 'image' || + activity === 'audio' || + activity === 'tts' || + activity === 'video' || + activity === 'transcription' + ? activity + : undefined +} + +function parseDataUrl( + value: string, +): { mimeType: string; bytes: Uint8Array } | undefined { + const match = /^data:([^;,]+)?(;base64)?,(.*)$/s.exec(value) + if (!match) return undefined + const mimeType = match[1] || 'application/octet-stream' + const payload = decodeURIComponent(match[3] ?? '') + return { + mimeType, + bytes: match[2] + ? base64ToUint8Array(payload) + : new TextEncoder().encode(payload), + } +} + +function extensionForMime(mimeType: string | undefined): string { + if (mimeType === undefined) return 'bin' + + switch (mimeType) { + case 'image/png': + return 'png' + case 'image/jpeg': + return 'jpg' + case 'audio/wav': + return 'wav' + case 'audio/mpeg': + return 'mp3' + case 'audio/mp3': + return 'mp3' + case 'video/mp4': + return 'mp4' + case 'application/json': + return 'json' + default: + return 'bin' + } +} + +function defaultArtifactName( + descriptor: GenerationArtifactDescriptor, + activity: PersistedArtifactActivity, + index: number, +): string { + const ext = extensionForMime(descriptor.mimeType) + return `${activity}-${descriptor.role}-${descriptor.mediaType ?? 'artifact'}-${index}.${ext}` +} + +function sourcePartDescriptors( + part: unknown, + role: PersistedArtifactRole, + path: string, +): Array { + const record = objectValue(part) + const type = stringField(record ?? {}, 'type') + const source = objectValue(record?.source) + if ( + !record || + !source || + (type !== 'image' && type !== 'audio' && type !== 'video') + ) { + return [] + } + const sourceType = stringField(source, 'type') + const mimeType = stringField(source, 'mimeType') ?? `${type}/mpeg` + if (sourceType === 'data') { + const value = stringField(source, 'value') + if (!value) return [] + return [ + { + role, + path, + mediaType: type, + mimeType, + bytes: base64ToUint8Array(value), + }, + ] + } + if (sourceType === 'url') { + const value = stringField(source, 'value') + if (!value) return [] + return [{ role, path, mediaType: type, mimeType, url: value }] + } + return [] +} + +function promptInputDescriptors( + inputs: unknown, +): Array { + const prompt = objectValue(inputs)?.prompt + if (!Array.isArray(prompt)) return [] + + const counts: Record = { image: 0, audio: 0, video: 0 } + const descriptors: Array = [] + for (const part of prompt) { + const type = stringField(objectValue(part) ?? {}, 'type') + if (type !== 'image' && type !== 'audio' && type !== 'video') continue + const index = counts[type] ?? 0 + counts[type] = index + 1 + descriptors.push( + ...sourcePartDescriptors(part, 'input', `prompt.${type}s.${index}`), + ) + } + return descriptors +} + +function generatedMediaDescriptor(args: { + role: PersistedArtifactRole + path: string + mediaType: 'image' | 'audio' | 'video' + mimeType: string + media: unknown + jobId?: string + expiresAt?: string | Date +}): GenerationArtifactDescriptor | undefined { + const media = objectValue(args.media) + if (!media) return undefined + const b64Json = stringField(media, 'b64Json') + if (b64Json) { + return { + role: args.role, + path: args.path, + mediaType: args.mediaType, + mimeType: stringField(media, 'contentType') ?? args.mimeType, + bytes: base64ToUint8Array(b64Json), + jobId: args.jobId, + expiresAt: args.expiresAt, + } + } + const url = stringField(media, 'url') + if (url) { + return { + role: args.role, + path: args.path, + mediaType: args.mediaType, + mimeType: stringField(media, 'contentType') ?? args.mimeType, + url, + jobId: args.jobId, + expiresAt: args.expiresAt, + } + } + return undefined +} + +function builtInArtifactDescriptors( + activity: PersistedArtifactActivity, + inputs: unknown, + result: unknown, +): Array { + const descriptors = promptInputDescriptors(inputs) + const output = objectValue(result) + if (!output) return descriptors + + if (activity === 'image' && Array.isArray(output.images)) { + output.images.forEach((image, index) => { + const descriptor = generatedMediaDescriptor({ + role: 'output', + path: `images.${index}`, + mediaType: 'image', + mimeType: 'image/png', + media: image, + }) + if (descriptor) descriptors.push(descriptor) + }) + } + + if (activity === 'audio') { + const descriptor = generatedMediaDescriptor({ + role: 'output', + path: 'audio', + mediaType: 'audio', + mimeType: 'audio/mpeg', + media: output.audio, + }) + if (descriptor) descriptors.push(descriptor) + } + + if (activity === 'tts') { + const audio = stringField(output, 'audio') + if (audio) { + const format = stringField(output, 'format') + descriptors.push({ + role: 'output', + path: 'audio', + mediaType: 'audio', + mimeType: + stringField(output, 'contentType') ?? + (format ? `audio/${format}` : 'audio/mpeg'), + bytes: base64ToUint8Array(audio), + }) + } + } + + if (activity === 'video' && typeof output.url === 'string') { + descriptors.push({ + role: 'output', + path: 'video', + mediaType: 'video', + mimeType: 'video/mp4', + url: output.url, + jobId: stringField(output, 'jobId'), + expiresAt: + output.expiresAt instanceof Date ? output.expiresAt : undefined, + }) + } + + if (activity === 'transcription') { + const audio = objectValue(inputs)?.audio + if (typeof audio === 'string') { + const data = parseDataUrl(audio) + descriptors.push({ + role: 'input', + path: 'audio', + mediaType: 'audio', + mimeType: data?.mimeType ?? 'audio/mpeg', + bytes: data?.bytes ?? base64ToUint8Array(audio), + }) + } else if (audio instanceof ArrayBuffer) { + descriptors.push({ + role: 'input', + path: 'audio', + mediaType: 'audio', + mimeType: 'audio/mpeg', + bytes: audio.slice(0), + }) + } else if (typeof Blob !== 'undefined' && audio instanceof Blob) { + descriptors.push({ + role: 'input', + path: 'audio', + mediaType: 'audio', + mimeType: audio.type || 'audio/mpeg', + bytes: audio, + }) + } + if (Array.isArray(output.segments) || Array.isArray(output.words)) { + descriptors.push({ + role: 'output', + path: 'transcription', + mediaType: 'json', + mimeType: 'application/json', + json: output, + }) + } + } + + return descriptors +} + +async function descriptorBody( + descriptor: GenerationArtifactDescriptor, +): Promise<{ + body: BlobBody + size: number + mimeType: string + externalUrl?: string +}> { + if (descriptor.json !== undefined) { + const body = JSON.stringify(descriptor.json) + return { + body, + size: new TextEncoder().encode(body).byteLength, + mimeType: descriptor.mimeType ?? 'application/json', + } + } + + if (descriptor.bytes !== undefined) { + const body = descriptor.bytes + let size: number + if (typeof body === 'string') { + size = new TextEncoder().encode(body).byteLength + } else if (body instanceof ArrayBuffer) { + size = body.byteLength + } else if (ArrayBuffer.isView(body)) { + size = body.byteLength + } else if (typeof Blob !== 'undefined' && body instanceof Blob) { + size = body.size + } else { + size = 0 + } + return { + body, + size, + mimeType: descriptor.mimeType ?? 'application/octet-stream', + } + } + + if (descriptor.url) { + const data = parseDataUrl(descriptor.url) + if (data) { + return { + body: data.bytes, + size: data.bytes.byteLength, + mimeType: descriptor.mimeType ?? data.mimeType, + } + } + const response = await fetch(descriptor.url) + if (!response.ok) { + throw new Error( + `Failed to persist artifact from ${descriptor.url}: HTTP ${response.status}`, + ) + } + const mimeType = + descriptor.mimeType ?? + response.headers.get('content-type') ?? + 'application/octet-stream' + // Stream the body straight into the blob store instead of buffering the + // whole artifact in memory. `size` is left 0 (unknown up front); the store + // records the actual byte length as it drains the stream. Fall back to + // buffering only when the response has no body to stream. + if (response.body) { + return { + body: response.body, + size: 0, + mimeType, + externalUrl: descriptor.url, + } + } + const body = await response.arrayBuffer() + return { + body, + size: body.byteLength, + mimeType, + externalUrl: descriptor.url, + } + } + + throw new Error( + `Artifact descriptor ${descriptor.path} has no bytes, url, or json.`, + ) +} + +async function persistGenerationArtifacts( + persistence: AIPersistence, + opts: WithPersistenceOptions | undefined, + ctx: GenerationMiddlewareContext, + result: unknown, +): Promise> { + const activity = mediaActivity(ctx.activity) + if (!activity) return [] + + const threadId = ctx.threadId ?? ctx.requestId + const runId = ctx.runId ?? ctx.requestId + const extractionInput: GenerationArtifactExtractionInput = { + activity, + provider: ctx.provider, + model: ctx.model, + threadId, + runId, + inputs: ctx.artifactInputs, + result, + } + const extracted = + opts?.extractArtifacts !== undefined + ? await opts.extractArtifacts(extractionInput) + : builtInArtifactDescriptors(activity, ctx.artifactInputs, result) + + if (extracted.length === 0) return [] + + const existingRefs = extracted.filter(isArtifactRef) + const descriptors = extracted.filter( + (item): item is GenerationArtifactDescriptor => !isArtifactRef(item), + ) + if (descriptors.length === 0) return existingRefs + + if (!persistence.stores.artifacts || !persistence.stores.blobs) { + throw new Error( + 'Generation artifact persistence requires stores.artifacts and stores.blobs.', + ) + } + + const refs: Array = [...existingRefs] + for (const [index, descriptor] of descriptors.entries()) { + const artifactId = ctx.createId('artifact') + const { body, size, mimeType, externalUrl } = + await descriptorBody(descriptor) + const key = artifactBlobKey({ runId, artifactId }) + const stored = await persistence.stores.blobs.put(key, body, { + contentType: mimeType, + customMetadata: { + runId, + threadId, + role: descriptor.role, + activity, + path: descriptor.path, + }, + }) + // For streamed downloads the descriptor size is unknown (0); the store + // reports the real byte length once it has drained the stream. + const resolvedSize = size || stored.size || 0 + const createdAtMs = Date.now() + const name = + opts?.nameArtifact?.({ + descriptor: { ...descriptor, mimeType }, + activity, + provider: ctx.provider, + model: ctx.model, + threadId, + runId, + index, + }) ?? + descriptor.name ?? + defaultArtifactName({ ...descriptor, mimeType }, activity, index) + const record: ArtifactRecord = { + artifactId, + runId, + threadId, + name, + mimeType, + size: resolvedSize, + externalUrl, + createdAt: createdAtMs, + } + await persistence.stores.artifacts.save(record) + refs.push({ + role: descriptor.role, + artifactId, + threadId, + runId, + name, + mimeType, + size: resolvedSize, + createdAt: new Date(createdAtMs).toISOString(), + ...(externalUrl ? { externalUrl } : {}), + source: { + activity, + path: descriptor.path, + provider: ctx.provider, + model: ctx.model, + mediaType: descriptor.mediaType, + jobId: descriptor.jobId, + expiresAt: + descriptor.expiresAt instanceof Date + ? descriptor.expiresAt.toISOString() + : descriptor.expiresAt, + }, + }) + } + + // Stamp the durable app-origin serve URL onto every ref that lacks one, so + // clients render + restore media from your own origin, not the provider link. + if (opts?.artifactUrl) { + for (let i = 0; i < refs.length; i++) { + const ref = refs[i] + if (ref && !ref.url) { + const url = opts.artifactUrl(ref) + if (url) refs[i] = { ...ref, url } + } + } + } + + return refs +} + +/** + * Rewrite the live result's media fields to each output ref's durable serve URL + * (`ref.url`), so the live result matches what a reload restores. Keyed off the + * ref's `source.path`: `images.` → `result.images[i].url`, `video` → + * `result.url`, `audio` (object) → `result.audio.url`. tts (a base64 string) and + * transcription (json) have no media-URL field, so they are left as-is; their + * durable bytes are reachable via `result.artifacts`. A no-op when no ref has a + * `url`. + */ +function applyDurableMediaUrls( + result: Record, + refs: Array, +): Record { + let next = result + for (const ref of refs) { + if (ref.role !== 'output' || !ref.url) continue + const path = ref.source.path + if (path.startsWith('images.')) { + const index = Number(path.slice('images.'.length)) + const images = next.images + if (Array.isArray(images) && objectValue(images[index])) { + const cloned = [...images] + cloned[index] = { ...objectValue(images[index]), url: ref.url } + next = { ...next, images: cloned } + } + } else if (path === 'video') { + next = { ...next, url: ref.url } + } else if (path === 'audio' && objectValue(next.audio)) { + next = { ...next, audio: { ...objectValue(next.audio), url: ref.url } } + } + } + return next +} + // --------------------------------------------------------------------------- // Shared store / feature plan // --------------------------------------------------------------------------- interface PersistencePlan { wantsInterrupts: boolean + wantsArtifactPersistence: boolean runs: AIPersistence['stores']['runs'] } function resolvePersistencePlan(persistence: AIPersistence): PersistencePlan { return { wantsInterrupts: persistence.stores.interrupts !== undefined, + wantsArtifactPersistence: + persistence.stores.artifacts !== undefined && + persistence.stores.blobs !== undefined, runs: persistence.stores.runs, } } @@ -310,13 +879,22 @@ type InvalidChatPersistence = ? StoreIsDefinitelyAbsent : false +/** + * Generation entrypoint invalid when `jobs` is known-absent, or when exactly one + * of `artifacts` / `blobs` is present (artifact persistence needs both). + */ +type InvalidGenerationPersistence = + StoreIsDefinitelyAbsent extends true + ? true + : StoreIsDefinitelyPresent extends true + ? StoreIsDefinitelyAbsent + : StoreIsDefinitelyPresent extends true + ? StoreIsDefinitelyAbsent + : false + type ValidChatPersistence = InvalidChatPersistence extends true ? never : unknown -/** Generation entrypoint invalid when `runs` is known-absent. */ -type InvalidGenerationPersistence = - StoreIsDefinitelyAbsent extends true ? true : false - type ValidGenerationPersistence = InvalidGenerationPersistence extends true ? never : unknown @@ -622,61 +1200,128 @@ export function withPersistence( // --------------------------------------------------------------------------- /** - * Generation-only persistence middleware. Tracks run status (run records) for - * image, audio, TTS, video, and transcription activities. - * - * Requires `stores.runs`. + * Generation-only persistence middleware. Tracks generation job status (job + * records keyed by `jobId`) and, when `stores.artifacts` + `stores.blobs` are + * both provided, persists the generated media for image, audio, TTS, video, and + * transcription activities. * - * ⚠️ TEMPORARY / WRONG SHAPE — do not extend this design. + * Requires `stores.jobs`. A generation activity has no conversation, so the job + * is keyed on its own `jobId` (`ctx.runId ?? ctx.requestId`). `ctx.threadId` is + * carried through only as an OPTIONAL *link* to a chat when the caller supplies + * one — it is never the job's primary identity and is never faked from the + * request id. * - * Generation jobs must **not** fake `threadId = requestId`. `threadId` is the - * shared conversation key ({@link Scope.threadId} / chat middleware); a - * generation activity has no conversation. Dual-keying chat {@link RunStore} - * with `(runId: requestId, threadId: requestId)` pollutes chat run queries and - * confuses `findActiveRun(threadId)`. - * - * The follow-up generation-persistence work should introduce a dedicated job - * store (e.g. `GenerationJobStore` keyed by `jobId` / `requestId`) and optional - * later artifact storage — not reuse chat `RunStore` / `MessageStore`. An - * optional `threadId` on a job is only a *link* to a chat when the product - * needs it, never the job's primary identity. + * On success the terminal result metadata (ids, urls — never media bytes) and, + * when artifact persistence is on, the persisted artifact refs are captured onto + * the job record so a server-authoritative client can hydrate the last + * generation for a thread via {@link reconstructGeneration}. */ -export function withGenerationPersistence< - TStores extends AIPersistenceStores & { runs: RunStore }, ->( +export function withGenerationPersistence( persistence: AIPersistence & ValidGenerationPersistence, + opts?: WithPersistenceOptions, +): GenerationMiddleware +export function withGenerationPersistence( + persistence: AIPersistence, + opts?: WithPersistenceOptions, ): GenerationMiddleware { validateGenerationPersistenceStores(persistence) - const runStore = persistence.stores.runs - if (!runStore) { + const { wantsArtifactPersistence } = resolvePersistencePlan(persistence) + const jobs = persistence.stores.jobs + if (!jobs) { // validateGenerationPersistenceStores already throws; this narrows for TypeScript. - throw new Error('Generation persistence requires stores.runs.') + throw new Error('Generation persistence requires stores.jobs.') } + const jobIdOf = (ctx: GenerationMiddlewareContext): string => + ctx.runId ?? ctx.requestId + return { name: 'generation-persistence', async onStart(ctx: GenerationMiddlewareContext) { - // STOPGAP ONLY — see function JSDoc. Do not copy this pattern. - await createOrResumeRun(runStore, ctx.requestId, ctx.requestId) + const jobId = jobIdOf(ctx) + const threadId = ctx.threadId + await jobs.createOrResume({ + jobId, + activity: ctx.activity, + provider: ctx.provider, + model: ctx.model, + startedAt: Date.now(), + ...(threadId !== undefined ? { threadId } : {}), + }) + + // Extract + persist artifact bytes (media → blobs, metadata → artifacts) + // and merge the resulting refs onto the result. Gated on artifact stores. + if (wantsArtifactPersistence) { + ctx.resultTransforms?.push(async (result) => { + const refs = await persistGenerationArtifacts( + persistence, + opts, + ctx, + result, + ) + if (refs.length === 0) return undefined + const base = objectValue(result) ?? {} + const existing = base.artifacts + const withArtifacts = { + ...base, + artifacts: [...(Array.isArray(existing) ? existing : []), ...refs], + } + // Point the live result's media at the durable serve URL (when + // `artifactUrl` stamped one), so live and restored results match. + return applyDurableMediaUrls(withArtifacts, refs) + }) + } + + // Always capture the terminal result metadata + any artifact refs onto the + // job record. Registered AFTER the artifact transform so it observes the + // fully-merged result (with the artifact refs attached). `result` is + // metadata/urls only — the media bytes already live in the blob store. + ctx.resultTransforms?.push(async (result) => { + const rawArtifacts = objectValue(result)?.artifacts + const artifacts = Array.isArray(rawArtifacts) + ? rawArtifacts.filter(isArtifactRef) + : [] + await jobs.update(jobId, { + result, + ...(artifacts.length > 0 ? { artifacts } : {}), + }) + return undefined + }) }, async onFinish( ctx: GenerationMiddlewareContext, info: GenerationFinishInfo, ) { - await completeRun(runStore, ctx.requestId, info.usage) + await jobs.update(jobIdOf(ctx), { + status: 'complete', + finishedAt: Date.now(), + ...(info.usage ? { usage: info.usage } : {}), + }) }, async onError(ctx: GenerationMiddlewareContext, info: GenerationErrorInfo) { - await failRun(runStore, ctx.requestId, info.error) + await jobs.update(jobIdOf(ctx), { + status: 'error', + finishedAt: Date.now(), + error: { + message: + info.error instanceof Error + ? info.error.message + : String(info.error), + }, + }) }, async onAbort( ctx: GenerationMiddlewareContext, _info: GenerationAbortInfo, ) { - await interruptRun(runStore, ctx.requestId) + await jobs.update(jobIdOf(ctx), { + status: 'interrupted', + finishedAt: Date.now(), + }) }, } } diff --git a/packages/ai-persistence/src/reconstruct-generation.ts b/packages/ai-persistence/src/reconstruct-generation.ts new file mode 100644 index 000000000..2c49d4215 --- /dev/null +++ b/packages/ai-persistence/src/reconstruct-generation.ts @@ -0,0 +1,184 @@ +import { validateReconstructGenerationStores } from './types' +import type { AIPersistence, GenerationJobRecord } from './types' + +/** + * The JSON body `reconstructGeneration` returns and a server-authoritative + * client hydrates from on mount. + * + * `resumeSnapshot` mirrors the last generation job for the requested thread (or + * a specific job id): its terminal/running `status`, the `result` metadata and + * `error` it recorded, the `activity` it ran, and a `resumeState` cursor + * (present only while the job is still `running`) the client can use to tail the + * live generation. `null` when there is no matching job. + * + * `activeRun` is `{ runId }` when the resolved job is still `running`, else + * `null` — the parallel of {@link ReconstructedChat.activeRun}. + */ +export interface ReconstructedGeneration { + resumeSnapshot: { + schemaVersion: 1 + resumeState: { threadId?: string; runId: string } | null + status: 'idle' | 'running' | 'complete' | 'error' + result?: unknown + error?: { message: string; code?: string } + activity?: string + } | null + activeRun: { runId: string } | null +} + +export interface ReconstructGenerationOptions { + /** Query parameter carrying the thread id. Defaults to `threadId`. */ + param?: string + /** Query parameter carrying the job id. Defaults to `jobId`. */ + jobParam?: string + /** + * Authorize access to the requested generation before loading it. + * + * ⚠️ Without this, any caller who knows or guesses `?threadId=` / `?jobId=` + * receives the generation's status and result metadata. Multi-user / + * multi-tenant deployments **must** supply an authorization check (session → + * owned thread/job) or resolve a validated id in the route. + * + * Called with whichever id was supplied — the `jobId` when present, else the + * `threadId`. Return: + * - `true` to allow the load + * - `false` for a default `403` response + * - a `Response` to return as-is (e.g. `401` with a body) + */ + authorize?: ( + id: string, + request: Request, + ) => boolean | Response | Promise +} + +/** + * Map the persisted job status to the client-facing resume-snapshot status. + * An `interrupted` job surfaces as `error` — the client has no live run to + * resume, and an interrupted generation produced no usable result. + */ +function snapshotStatus( + status: GenerationJobRecord['status'], +): 'running' | 'complete' | 'error' { + switch (status) { + case 'running': + return 'running' + case 'complete': + return 'complete' + case 'error': + case 'interrupted': + return 'error' + } +} + +function jobToSnapshot( + job: GenerationJobRecord, +): NonNullable { + const status = snapshotStatus(job.status) + return { + schemaVersion: 1, + resumeState: + status === 'running' + ? { + runId: job.jobId, + ...(job.threadId !== undefined ? { threadId: job.threadId } : {}), + } + : null, + status, + ...(job.result !== undefined ? { result: job.result } : {}), + ...(job.error !== undefined ? { error: job.error } : {}), + ...(job.activity !== undefined ? { activity: job.activity } : {}), + } +} + +function jsonResponse(body: ReconstructedGeneration): Response { + return new Response(JSON.stringify(body), { + headers: { + 'content-type': 'application/json', + 'cache-control': 'no-store', + }, + }) +} + +/** + * Build the JSON `Response` a server-authoritative client hydrates a generation + * from on load. Reads a `?jobId=` (preferred) or `?threadId=` from the request + * query and returns `{ resumeSnapshot, activeRun }` + * ({@link ReconstructedGeneration}): + * + * - Resolves the job by `jobId` via `stores.jobs.get`, else the latest job + * linked to `threadId` via the optional `stores.jobs.findLatestForThread`. + * - `resumeSnapshot` — the job mapped to a client snapshot (status, result, + * error, activity, and a `resumeState` cursor while still running), or `null`. + * - `activeRun` — `{ runId }` when the job is still generating, else `null`. + * + * Requires `stores.jobs`. Returns `{ resumeSnapshot: null, activeRun: null }` + * when no id is supplied or no matching job exists, so the caller never has to + * special-case a first load. + * + * This helper does **not** enforce tenancy by itself. Pass + * {@link ReconstructGenerationOptions.authorize} (or wrap the call in your own + * session gate) before exposing it on a public route. + * + * ```ts + * export async function GET(request: Request) { + * return reconstructGeneration(persistence, request, { + * authorize: async (id, req) => { + * const userId = await getSessionUserId(req) + * return userId != null && (await userOwnsThread(userId, id)) + * }, + * }) + * } + * ``` + */ +export async function reconstructGeneration( + persistence: AIPersistence, + request: Request, + options?: ReconstructGenerationOptions, +): Promise { + validateReconstructGenerationStores(persistence) + const jobStore = persistence.stores.jobs + if (!jobStore) { + // validateReconstructGenerationStores already throws; this narrows for TS. + throw new Error('reconstructGeneration requires stores.jobs.') + } + + const params = new URL(request.url).searchParams + const jobParam = options?.jobParam ?? 'jobId' + const threadParam = options?.param ?? 'threadId' + const jobId = params.get(jobParam) ?? '' + const threadId = params.get(threadParam) ?? '' + + const id = jobId || threadId + if (!id) { + return jsonResponse({ resumeSnapshot: null, activeRun: null }) + } + + if (options?.authorize) { + const decision = await options.authorize(id, request) + if (decision instanceof Response) { + return decision + } + if (!decision) { + return new Response(JSON.stringify({ error: 'Forbidden' }), { + status: 403, + headers: { + 'content-type': 'application/json', + 'cache-control': 'no-store', + }, + }) + } + } + + const job = jobId + ? await jobStore.get(jobId) + : ((await jobStore.findLatestForThread?.(threadId)) ?? null) + + if (!job) { + return jsonResponse({ resumeSnapshot: null, activeRun: null }) + } + + return jsonResponse({ + resumeSnapshot: jobToSnapshot(job), + activeRun: job.status === 'running' ? { runId: job.jobId } : null, + }) +} diff --git a/packages/ai-persistence/src/retrieve.ts b/packages/ai-persistence/src/retrieve.ts new file mode 100644 index 000000000..0984e3ab8 --- /dev/null +++ b/packages/ai-persistence/src/retrieve.ts @@ -0,0 +1,45 @@ +import type { AIPersistence, ArtifactRecord, BlobObject } from './types' + +/** + * The blob-store key a generation artifact's bytes are stored under. + * `withGenerationPersistence` writes bytes to this key; `retrieveBlob` reads + * from it. Keep the two in lockstep by using this helper on both sides. + */ +export function artifactBlobKey( + ref: Pick, +): string { + return `artifacts/${ref.runId}/${ref.artifactId}` +} + +/** + * Look up a persisted generation artifact's metadata by id. Returns `null` when + * the persistence has no `artifacts` store or no record matches — so a serve + * handler can map that straight to a 404. + */ +export async function retrieveArtifact( + persistence: AIPersistence, + artifactId: string, +): Promise { + const record = await persistence.stores.artifacts?.get(artifactId) + return record ?? null +} + +/** + * Look up a persisted generation artifact's stored bytes. Pass an `artifactId` + * (resolved to its record first) or an already-loaded {@link ArtifactRecord} + * (no second metadata lookup). Returns `null` when the artifact, its record, or + * its blob is missing, or the stores are not configured. + */ +export async function retrieveBlob( + persistence: AIPersistence, + artifact: string | ArtifactRecord, +): Promise { + const record = + typeof artifact === 'string' + ? await retrieveArtifact(persistence, artifact) + : artifact + if (!record) return null + + const blob = await persistence.stores.blobs?.get(artifactBlobKey(record)) + return blob ?? null +} diff --git a/packages/ai-persistence/src/types.ts b/packages/ai-persistence/src/types.ts index a1fd7f488..1f3f95b9c 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, + PersistedArtifactRef, + Scope, + 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: @@ -31,9 +36,10 @@ export type { Scope } // // TIMESTAMP CONVENTION // -------------------- -// Store *records* (`RunRecord`, `InterruptRecord`) speak **epoch -// milliseconds** (`number`), the native unit for SQL/`BIGINT` columns and -// `Date.now()`. Wire/result references that leave the persistence layer speak +// Store *records* (`RunRecord`, `InterruptRecord`, `ArtifactRecord`, +// `BlobRecord`) speak **epoch milliseconds** (`number`), the native unit for +// SQL/`BIGINT` columns and `Date.now()`. Wire/result references that leave the +// persistence layer (e.g. core's `PersistedArtifactRef.createdAt`) speak // **ISO-8601 strings**. The middleware performs the number→ISO conversion at // the boundary; do not mix the two on a single field. @@ -140,6 +146,94 @@ export interface RunStore { findActiveRun?: (threadId: string) => Promise } +export type GenerationJobStatus = + | 'running' + | 'complete' + | 'error' + | 'interrupted' + +/** + * A single generation job (one `generateImage` / `generateVideo` / … call). + * + * Its primary identity is `jobId` (the run/request id the activity mints) — a + * generation has no conversation, so `threadId` is only an OPTIONAL *link* to a + * chat when the product wants one (e.g. to correlate a generation with the + * thread that triggered it). Do not key generation state on the chat + * {@link RunStore} / {@link Scope.threadId}. + * + * `result` holds terminal result METADATA (ids, model, urls, a video jobId), + * never the media bytes — those live in a {@link BlobStore}. `artifacts` are the + * durable {@link PersistedArtifactRef}s, present only when byte storage is on. + * + * @property startedAt - Epoch ms when the job was first created. + * @property finishedAt - Epoch ms when the job reached a terminal status. + */ +export interface GenerationJobRecord { + jobId: string + /** Optional link to the chat conversation that triggered this generation. */ + threadId?: string + /** `'image' | 'audio' | 'tts' | 'video' | 'transcription'`. */ + activity: string + provider: string + model: string + status: GenerationJobStatus + startedAt: number + finishedAt?: number + error?: { message: string; code?: string } + /** Terminal result metadata (ids, model, urls). Never the media bytes. */ + result?: unknown + /** Durable artifact references, when an artifacts + blobs backend is used. */ + artifacts?: Array + usage?: TokenUsage +} + +/** + * Durable store for generation job records — the generation counterpart to + * {@link RunStore}, keyed by `jobId` rather than a conversation `threadId`. + */ +export interface GenerationJobStore { + /** + * Create a job record, or return the existing one if `jobId` is already + * present (resume). + * + * INVARIANT (idempotency): a second call for a `jobId` returns the existing + * record unchanged; `startedAt`/`activity`/`provider`/`model`/`threadId` are + * not mutated. `status` defaults to `'running'` on first creation. + */ + createOrResume: ( + input: Pick< + GenerationJobRecord, + 'jobId' | 'activity' | 'provider' | 'model' | 'startedAt' + > & { threadId?: string; status?: GenerationJobStatus }, + ) => Promise + /** + * Patch a job record's mutable fields. + * + * INVARIANT: patching a `jobId` that does not exist is a **no-op** — it must + * not throw and must not create a record. + */ + update: ( + jobId: string, + patch: Partial< + Pick< + GenerationJobRecord, + 'status' | 'finishedAt' | 'error' | 'result' | 'artifacts' | 'usage' + > + >, + ) => Promise + /** Return the job record for `jobId`, or `null` if none exists. */ + get: (jobId: string) => Promise + /** + * The most recent job linked to `threadId`, or `null`. OPTIONAL — callers + * feature-detect it (`store.findLatestForThread?.(threadId)`). Lets a + * server-authoritative client hydrate the last generation for a thread by the + * stable thread id, without handling a job id. + */ + findLatestForThread?: ( + threadId: string, + ) => Promise +} + /** Lifecycle status of a human-in-the-loop interrupt. */ export type InterruptStatus = 'pending' | 'resolved' | 'cancelled' @@ -276,6 +370,143 @@ export function defineInterruptStore(store: InterruptStore): InterruptStore { export function defineMetadataStore(store: MetadataStore): MetadataStore { return store } +/** Type a {@link GenerationJobStore} implementation inline. */ +export function defineGenerationJobStore( + store: GenerationJobStore, +): GenerationJobStore { + return store +} +/** Type an {@link ArtifactStore} implementation inline. */ +export function defineArtifactStore(store: ArtifactStore): ArtifactStore { + return store +} +/** Type a {@link BlobStore} implementation inline. */ +export function defineBlobStore(store: BlobStore): BlobStore { + return store +} + +/** + * Metadata row describing a persisted artifact (generated media, tool output). + * + * The bytes themselves live in a {@link BlobStore}; this record holds the + * descriptive metadata and an optional `externalUrl` for reference-only + * backends. + * + * @property createdAt - Epoch ms. (Core's wire-facing `PersistedArtifactRef` + * exposes the same instant as an ISO string; see the timestamp convention.) + */ +export interface ArtifactRecord { + artifactId: string + runId: string + threadId: string + name: string + mimeType: string + size: number + externalUrl?: string + createdAt: number +} + +/** Durable store for artifact metadata records. */ +export interface ArtifactStore { + /** Insert or overwrite the artifact metadata record. */ + save: (record: ArtifactRecord) => Promise + /** Return the artifact for `artifactId`, or `null` if none exists. */ + get: (artifactId: string) => Promise + /** All artifacts for a run. Returns `[]` when the run has none. */ + list: (runId: string) => Promise> + /** OPTIONAL: delete a single artifact by id. */ + delete?: (artifactId: string) => Promise + /** OPTIONAL: delete every artifact belonging to `runId`. */ + deleteForRun?: (runId: string) => Promise +} + +/** + * Accepted body shapes for {@link BlobStore.put}. `ArrayBufferView` already + * covers `Uint8Array` and every other typed-array/`DataView`, so no separate + * `Uint8Array` member is needed. + */ +export type BlobBody = + | ReadableStream + | ArrayBuffer + | ArrayBufferView + | string + | Blob + +/** + * Metadata for a stored blob. + * + * @property size - Byte length, when known. + * @property createdAt - Epoch ms first written. + * @property updatedAt - Epoch ms last overwritten. + */ +export interface BlobRecord { + key: string + size?: number + etag?: string + contentType?: string + customMetadata?: Record + createdAt?: number + updatedAt?: number +} + +/** A stored blob's metadata plus lazy accessors for its bytes. */ +export interface BlobObject extends BlobRecord { + arrayBuffer: () => Promise + text: () => Promise + body?: ReadableStream +} + +/** + * One page of a {@link BlobStore.list} scan. + * + * @property cursor - Opaque continuation token; present only when `truncated`. + * @property truncated - `true` when more objects match beyond this page. + */ +export interface BlobListPage { + objects: Array + cursor?: string + truncated?: boolean +} + +export interface BlobPutOptions { + contentType?: string + customMetadata?: Record +} + +export interface BlobListOptions { + prefix?: string + cursor?: string + limit?: number +} + +/** Durable object/blob store (byte-storing or reference-only backends). */ +export interface BlobStore { + /** Insert or overwrite the object at `key`, returning its metadata. */ + put: ( + key: string, + body: BlobBody, + options?: BlobPutOptions, + ) => Promise + /** Return the object at `key` (metadata + byte accessors), or `null`. */ + get: (key: string) => Promise + /** Return only the metadata for `key`, or `null`. */ + head: (key: string) => Promise + /** Remove the object at `key`. A no-op if absent. */ + delete: (key: string) => Promise + /** + * List objects, optionally filtered by `prefix`, in ascending key order. + * + * CURSOR SEMANTICS: `prefix` matches literally and case-sensitively (SQL + * backends must escape LIKE metacharacters, so `run_` matches only the exact + * bytes `run_`, not `_` as a wildcard). When `limit` is given and more keys + * match, the page is `truncated: true` with a `cursor`; passing that `cursor` + * back returns the strictly-following keys (keys `> cursor`). Cursor ordering + * is the same byte ordering as the sort, so paging visits every key exactly + * once with no gaps or repeats. `limit: 0` yields an empty, untruncated page + * with no cursor. + */ + list: (options?: BlobListOptions) => Promise +} /** * Sparse bag of **state** store keys — composition / validation only. @@ -293,6 +524,9 @@ export interface AIPersistenceStores { runs?: RunStore interrupts?: InterruptStore metadata?: MetadataStore + jobs?: GenerationJobStore + artifacts?: ArtifactStore + blobs?: BlobStore } /** @@ -463,8 +697,11 @@ export type ComposedAIPersistenceStores< const storeKeys = [ 'messages', 'runs', + 'jobs', 'interrupts', 'metadata', + 'artifacts', + 'blobs', ] satisfies Array const storeKeySet = new Set(storeKeys) @@ -499,15 +736,24 @@ export function validateChatPersistenceStores( } /** - * Generation middleware entrypoint rule: `runs` is required (run lifecycle is - * the only generation state this middleware tracks). + * Generation middleware entrypoint rule: `jobs` is required (the generation job + * lifecycle is keyed on `jobId`, not a chat conversation `threadId`). When + * artifact persistence is used, `artifacts` and `blobs` must be provided + * together. */ export function validateGenerationPersistenceStores( persistence: AIPersistence, ): void { validatePersistenceStoreKeys(persistence) - if (!persistence.stores.runs) { - throw new Error('Generation persistence requires stores.runs.') + const hasArtifacts = persistence.stores.artifacts !== undefined + const hasBlobs = persistence.stores.blobs !== undefined + if (hasArtifacts !== hasBlobs) { + throw new Error( + 'Generation artifact persistence requires both stores.artifacts and stores.blobs.', + ) + } + if (!persistence.stores.jobs) { + throw new Error('Generation persistence requires stores.jobs.') } } @@ -523,6 +769,21 @@ export function validateReconstructChatStores( } } +/** + * Server hydrate entrypoint rule for generation: `jobs` is required. The job + * store resolves the latest generation for a thread (or a specific job id), so + * a server-authoritative client can hydrate the last generation's status, + * result, and artifact refs on load. + */ +export function validateReconstructGenerationStores( + persistence: AIPersistence, +): void { + validatePersistenceStoreKeys(persistence) + if (!persistence.stores.jobs) { + throw new Error('reconstructGeneration requires stores.jobs.') + } +} + export function defineAIPersistence( persistence: AIPersistence>, ): AIPersistence { diff --git a/packages/ai-persistence/tests/error-abort.test.ts b/packages/ai-persistence/tests/error-abort.test.ts index 091382666..ac572924a 100644 --- a/packages/ai-persistence/tests/error-abort.test.ts +++ b/packages/ai-persistence/tests/error-abort.test.ts @@ -245,9 +245,10 @@ function imageAdapterThatThrows(thrown: unknown): ImageAdapter { } } -// A generation activity's only identity is `requestId` (auto-generated), so the -// integration tests capture it via a probe middleware, and the direct-drive -// tests set `requestId` to the pre-created run's id. +// A generation job is keyed on `jobId` (`ctx.runId ?? ctx.requestId`). With no +// runId supplied the auto-generated `requestId` is the jobId, so the integration +// tests capture it via a probe middleware, and the direct-drive tests set +// `requestId` to the pre-created job's id. function generationContext(requestId: string): GenerationMiddlewareContext { return { requestId, @@ -261,7 +262,7 @@ function generationContext(requestId: string): GenerationMiddlewareContext { } describe('generation persistence error/abort hooks', () => { - it('marks the run failed when generation throws', async () => { + it('marks the job errored when generation throws', async () => { const persistence = memoryPersistence() let requestId = '' @@ -280,12 +281,12 @@ describe('generation persistence error/abort hooks', () => { }), ).rejects.toThrow('image boom') - const run = await persistence.stores.runs!.get(requestId) - expect(run?.status).toBe('failed') - expect(run?.error).toBe('image boom') + const job = await persistence.stores.jobs.get(requestId) + expect(job?.status).toBe('error') + expect(job?.error).toEqual({ message: 'image boom' }) }) - it('coerces a non-Error generation failure into the run error string', async () => { + it('coerces a non-Error generation failure into the job error message', async () => { const persistence = memoryPersistence() let requestId = '' @@ -304,18 +305,20 @@ describe('generation persistence error/abort hooks', () => { }), ).rejects.toBeDefined() - const run = await persistence.stores.runs!.get(requestId) - expect(run?.status).toBe('failed') - expect(run?.error).toBe('image string failure') + const job = await persistence.stores.jobs.get(requestId) + expect(job?.status).toBe('error') + expect(job?.error).toEqual({ message: 'image string failure' }) }) - it('marks the run interrupted on generation abort', async () => { + it('marks the job interrupted on generation abort', async () => { const persistence = memoryPersistence() const middleware = withGenerationPersistence(persistence) - await persistence.stores.runs!.createOrResume({ - runId: 'req-abort', - threadId: 'req-abort', + await persistence.stores.jobs.createOrResume({ + jobId: 'req-abort', + activity: 'image', + provider: 'test', + model: 'test-model', startedAt: 1, }) @@ -327,17 +330,19 @@ describe('generation persistence error/abort hooks', () => { } await middleware.onAbort?.(generationContext('req-abort'), abortInfo) - expect((await persistence.stores.runs!.get('req-abort'))?.status).toBe( + expect((await persistence.stores.jobs.get('req-abort'))?.status).toBe( 'interrupted', ) }) - it('coerces a non-Error into the run error string via the onError handler', async () => { + it('coerces a non-Error into the job error message via the onError handler', async () => { const persistence = memoryPersistence() const middleware = withGenerationPersistence(persistence) - await persistence.stores.runs!.createOrResume({ - runId: 'req-err', - threadId: 'req-err', + await persistence.stores.jobs.createOrResume({ + jobId: 'req-err', + activity: 'image', + provider: 'test', + model: 'test-model', startedAt: 1, }) const errorInfo: GenerationErrorInfo = { @@ -346,8 +351,8 @@ describe('generation persistence error/abort hooks', () => { } await middleware.onError?.(generationContext('req-err'), errorInfo) - const run = await persistence.stores.runs!.get('req-err') - expect(run?.status).toBe('failed') - expect(run?.error).toBe('[object Object]') + const job = await persistence.stores.jobs.get('req-err') + expect(job?.status).toBe('error') + expect(job?.error).toEqual({ message: '[object Object]' }) }) }) diff --git a/packages/ai-persistence/tests/generation-artifacts.test.ts b/packages/ai-persistence/tests/generation-artifacts.test.ts new file mode 100644 index 000000000..98c8e2736 --- /dev/null +++ b/packages/ai-persistence/tests/generation-artifacts.test.ts @@ -0,0 +1,605 @@ +import { describe, expect, it, vi } from 'vitest' +import { + EventType, + generateAudio, + generateImage, + generateTranscription, +} from '@tanstack/ai' +import { composePersistence, defineAIPersistence } from '../src/types' +import { memoryPersistence } from '../src/memory' +import { withGenerationPersistence } from '../src/middleware' +import { retrieveArtifact, retrieveBlob } from '../src/retrieve' +import type { + GenerationArtifactDescriptor, + GenerationArtifactExtractionInput, + GenerationArtifactNameInput, + AIPersistence, +} from '../src' +import type { + AudioAdapter, + AudioGenerationResult, + ImageAdapter, + PersistedArtifactRef, + StreamChunk, + TranscriptionAdapter, + TranscriptionResult, +} from '@tanstack/ai' + +void (undefined as unknown as GenerationArtifactDescriptor) +void (undefined as unknown as GenerationArtifactExtractionInput) +void (undefined as unknown as GenerationArtifactNameInput) + +type AudioGenerateOptions = Parameters[0] & { + threadId?: string + runId?: string + replay?: unknown +} + +type TranscriptionGenerateOptions = Parameters< + typeof generateTranscription +>[0] & { + threadId?: string + runId?: string +} + +const imageAdapterTypes: ImageAdapter['~types'] = { + providerOptions: {}, + modelProviderOptionsByName: {}, + modelSizeByName: {}, + modelInputModalitiesByName: {}, +} + +const audioAdapterTypes: AudioAdapter['~types'] = { + providerOptions: {}, +} + +const transcriptionAdapterTypes: TranscriptionAdapter['~types'] = { + providerOptions: {}, +} + +async function collect(stream: AsyncIterable) { + const chunks: Array = [] + for await (const chunk of stream) chunks.push(chunk) + return chunks +} + +function imageAdapter(): ImageAdapter { + return { + kind: 'image', + name: 'test-image-provider', + model: 'test-image-model', + '~types': imageAdapterTypes, + generateImages: vi.fn(async () => ({ + id: 'image-result', + model: 'test-image-model', + images: [{ b64Json: 'b3V0cHV0LWltYWdl' }], + })), + } +} + +function audioAdapter(): AudioAdapter { + return { + kind: 'audio', + name: 'test-audio-provider', + model: 'test-audio-model', + '~types': audioAdapterTypes, + generateAudio: vi.fn(async () => ({ + id: 'audio-result', + model: 'test-audio-model', + audio: { + b64Json: 'b3V0cHV0LWF1ZGlv', + contentType: 'audio/wav', + duration: 1, + }, + })), + } +} + +function transcriptionAdapter(): TranscriptionAdapter { + return { + kind: 'transcription', + name: 'test-transcription-provider', + model: 'test-transcription-model', + '~types': transcriptionAdapterTypes, + transcribe: vi.fn(async () => ({ + id: 'transcription-result', + model: 'test-transcription-model', + text: 'hello world', + language: 'en', + segments: [{ id: 0, start: 0, end: 1, text: 'hello world' }], + })), + } +} + +describe('withGenerationPersistence generation artifacts', () => { + it('persists built-in image output artifacts and attaches refs', async () => { + const persistence = memoryPersistence() + + const result = await generateImage({ + adapter: imageAdapter(), + prompt: 'make an image', + threadId: 'thread-image', + runId: 'run-image', + middleware: [withGenerationPersistence(persistence)], + }) + + expect(result.artifacts).toHaveLength(1) + expect(result.artifacts?.[0]).toMatchObject({ + role: 'output', + threadId: 'thread-image', + runId: 'run-image', + mimeType: 'image/png', + size: 12, + source: { + activity: 'image', + path: 'images.0', + provider: 'test-image-provider', + model: 'test-image-model', + mediaType: 'image', + }, + }) + + const record = await persistence.stores.artifacts!.get( + result.artifacts![0]!.artifactId, + ) + expect(record).toMatchObject({ + runId: 'run-image', + threadId: 'thread-image', + mimeType: 'image/png', + size: 12, + }) + const blob = await persistence.stores.blobs!.get( + `artifacts/run-image/${result.artifacts![0]!.artifactId}`, + ) + await expect(blob?.text()).resolves.toBe('output-image') + }) + + it('stamps a durable url on refs and rewrites the live result media (artifactUrl)', async () => { + const persistence = memoryPersistence() + + const result = await generateImage({ + adapter: imageAdapter(), + prompt: 'make an image', + threadId: 'thread-url', + runId: 'run-url', + middleware: [ + withGenerationPersistence(persistence, { + artifactUrl: (ref) => `/artifacts/${ref.artifactId}`, + }), + ], + }) + + const ref = result.artifacts?.[0] + expect(ref).toBeDefined() + const expectedUrl = `/artifacts/${ref!.artifactId}` + // The ref carries the durable serve URL... + expect(ref!.url).toBe(expectedUrl) + // ...and the live result's media points at it too (durable everywhere). + expect(result.images[0]?.url).toBe(expectedUrl) + }) + + it('retrieveArtifact / retrieveBlob fetch a persisted artifact and its bytes', async () => { + const persistence = memoryPersistence() + + const result = await generateImage({ + adapter: imageAdapter(), + prompt: 'make an image', + threadId: 'thread-retrieve', + runId: 'run-retrieve', + middleware: [withGenerationPersistence(persistence)], + }) + const artifactId = result.artifacts![0]!.artifactId + + const record = await retrieveArtifact(persistence, artifactId) + expect(record).toMatchObject({ + runId: 'run-retrieve', + mimeType: 'image/png', + }) + + // By id (resolves the record first) and by an already-loaded record. + await expect( + (await retrieveBlob(persistence, artifactId))?.text(), + ).resolves.toBe('output-image') + await expect( + (await retrieveBlob(persistence, record!))?.text(), + ).resolves.toBe('output-image') + + // Unknown id resolves to null on both. + expect(await retrieveArtifact(persistence, 'missing')).toBeNull() + expect(await retrieveBlob(persistence, 'missing')).toBeNull() + }) + + it('persists non-image media outputs', async () => { + const persistence = memoryPersistence() + + const result = (await generateAudio({ + adapter: audioAdapter(), + prompt: 'make audio', + threadId: 'thread-audio', + runId: 'run-audio', + middleware: [withGenerationPersistence(persistence)], + } as AudioGenerateOptions)) as AudioGenerationResult + + expect(result.artifacts).toHaveLength(1) + expect(result.artifacts?.[0]).toMatchObject({ + role: 'output', + mimeType: 'audio/wav', + size: 12, + source: { + activity: 'audio', + path: 'audio', + mediaType: 'audio', + }, + }) + }) + + it('persists media inputs and includes input refs on the result', async () => { + const persistence = memoryPersistence() + + const result = await generateImage({ + adapter: imageAdapter(), + prompt: [ + { type: 'text', content: 'edit this' }, + { + type: 'image', + source: { + type: 'data', + value: 'aW5wdXQtaW1hZ2U=', + mimeType: 'image/png', + }, + }, + ], + threadId: 'thread-input', + runId: 'run-input', + middleware: [withGenerationPersistence(persistence)], + }) + + expect(result.artifacts?.map((artifact) => artifact.role)).toEqual([ + 'input', + 'output', + ]) + const input = result.artifacts?.[0] + expect(input).toMatchObject({ + role: 'input', + mimeType: 'image/png', + size: 11, + source: { path: 'prompt.images.0', mediaType: 'image' }, + }) + }) + + it('allows job tracking without artifact stores', () => { + const full = memoryPersistence() + const persistence = defineAIPersistence({ + stores: { + jobs: full.stores.jobs, + }, + }) + + expect(() => withGenerationPersistence(persistence)).not.toThrow() + }) + + it('throws when the job store is missing', () => { + const full = memoryPersistence() + const persistence: AIPersistence = defineAIPersistence({ + stores: { + runs: full.stores.runs, + }, + }) + + expect(() => withGenerationPersistence(persistence)).toThrow( + /Generation persistence requires stores\.jobs/i, + ) + }) + + it('records a job that transitions running -> complete with result + artifacts', async () => { + const persistence = memoryPersistence() + + const result = await generateImage({ + adapter: imageAdapter(), + prompt: 'make an image', + threadId: 'thread-job', + runId: 'run-job', + middleware: [withGenerationPersistence(persistence)], + }) + + const job = await persistence.stores.jobs.get('run-job') + expect(job).toMatchObject({ + jobId: 'run-job', + threadId: 'thread-job', + activity: 'image', + provider: 'test-image-provider', + model: 'test-image-model', + status: 'complete', + }) + expect(job?.finishedAt).toEqual(expect.any(Number)) + // Terminal result metadata is captured on the job (never the media bytes). + expect(job?.result).toBeDefined() + // Persisted artifact refs land on the job too. + expect(job?.artifacts).toHaveLength(1) + expect(job?.artifacts?.[0]?.artifactId).toBe( + result.artifacts![0]!.artifactId, + ) + }) + + it('links the job to a thread and finds the latest for that thread', async () => { + const persistence = memoryPersistence() + + await generateImage({ + adapter: imageAdapter(), + prompt: 'make an image', + threadId: 'thread-latest', + runId: 'run-latest-1', + middleware: [withGenerationPersistence(persistence)], + }) + + const latest = + await persistence.stores.jobs.findLatestForThread!('thread-latest') + expect(latest?.jobId).toBe('run-latest-1') + expect(latest?.status).toBe('complete') + }) + + it('records an error job when generation throws', async () => { + const persistence = memoryPersistence() + const adapter = imageAdapter() + adapter.generateImages = vi.fn(async () => { + throw new Error('boom') + }) + + await expect( + generateImage({ + adapter, + prompt: 'make an image', + threadId: 'thread-error', + runId: 'run-error', + middleware: [withGenerationPersistence(persistence)], + }), + ).rejects.toThrow('boom') + + const job = await persistence.stores.jobs.get('run-error') + expect(job).toMatchObject({ + jobId: 'run-error', + status: 'error', + error: { message: 'boom' }, + }) + }) + + it('uses custom artifact extraction instead of built-in extraction', async () => { + const persistence = memoryPersistence() + + const result = await generateImage({ + adapter: imageAdapter(), + prompt: [ + { type: 'text', content: 'edit this' }, + { + type: 'image', + source: { + type: 'data', + value: 'aW5wdXQtaW1hZ2U=', + mimeType: 'image/png', + }, + }, + ], + threadId: 'thread-custom', + runId: 'run-custom', + middleware: [ + withGenerationPersistence(persistence, { + extractArtifacts: () => [ + { + role: 'output', + path: 'custom', + mediaType: 'json', + mimeType: 'application/json', + json: { ok: true }, + name: 'custom.json', + }, + ], + }), + ], + }) + + expect(result.artifacts).toHaveLength(1) + expect(result.artifacts?.[0]).toMatchObject({ + name: 'custom.json', + mimeType: 'application/json', + source: { path: 'custom', mediaType: 'json' }, + }) + }) + + it('does not leak data URL bytes into artifact refs', async () => { + const persistence = memoryPersistence() + const dataUrl = 'data:image/png;base64,ZGF0YS11cmwtYnl0ZXM=' + + const result = await generateImage({ + adapter: imageAdapter(), + prompt: 'make an image', + threadId: 'thread-data-url', + runId: 'run-data-url', + middleware: [ + withGenerationPersistence(persistence, { + extractArtifacts: () => [ + { + role: 'input', + path: 'prompt.images.0', + mediaType: 'image', + url: dataUrl, + }, + { + role: 'output', + path: 'images.0', + mediaType: 'image', + url: dataUrl, + }, + ], + }), + ], + }) + + expect(result.artifacts).toHaveLength(2) + expect(result.artifacts?.map((artifact) => artifact.externalUrl)).toEqual([ + undefined, + undefined, + ]) + expect(JSON.stringify(result.artifacts)).not.toContain(dataUrl) + + const [input, output] = result.artifacts! + await expect( + persistence.stores.blobs + ?.get(`artifacts/run-data-url/${input!.artifactId}`) + .then((blob) => blob?.text()), + ).resolves.toBe('data-url-bytes') + await expect( + persistence.stores.blobs + ?.get(`artifacts/run-data-url/${output!.artifactId}`) + .then((blob) => blob?.text()), + ).resolves.toBe('data-url-bytes') + }) + + it('uses nameArtifact overrides', async () => { + const persistence = memoryPersistence() + + const result = (await generateAudio({ + adapter: audioAdapter(), + prompt: 'make audio', + threadId: 'thread-name', + runId: 'run-name', + middleware: [ + withGenerationPersistence(persistence, { + nameArtifact: ({ descriptor, index }) => + `${descriptor.role}-${descriptor.mediaType}-${index}.bin`, + }), + ], + } as AudioGenerateOptions)) as AudioGenerationResult + + expect(result.artifacts?.[0]?.name).toBe('output-audio-0.bin') + }) + + it('emits generation:artifacts before generation:result with persisted refs', async () => { + const persistence = memoryPersistence() + + const chunks = await collect( + generateImage, true>({ + adapter: imageAdapter(), + prompt: 'make an image', + stream: true, + threadId: 'thread-stream', + runId: 'run-stream', + middleware: [withGenerationPersistence(persistence)], + }), + ) + + const customEvents = chunks.filter( + (chunk) => chunk.type === EventType.CUSTOM, + ) + expect(customEvents.map((chunk) => chunk.name)).toEqual([ + 'generation:artifacts', + 'generation:result', + ]) + expect(customEvents[0]?.value).toEqual( + (customEvents[1]?.value as { artifacts?: unknown }).artifacts, + ) + }) + + it('uses the same fallback run and thread ids for streamed events and persisted artifact refs', async () => { + const persistence = memoryPersistence() + + const chunks = await collect( + generateImage, true>({ + adapter: imageAdapter(), + prompt: 'make an image', + stream: true, + middleware: [withGenerationPersistence(persistence)], + }), + ) + + const started = chunks.find((chunk) => chunk.type === EventType.RUN_STARTED) + const result = chunks.find( + (chunk) => + chunk.type === EventType.CUSTOM && chunk.name === 'generation:result', + ) + const artifact = ( + result as unknown as + | { value?: { artifacts?: Array } } + | undefined + )?.value?.artifacts?.[0] + + expect(started).toMatchObject({ + runId: expect.any(String), + threadId: expect.any(String), + }) + expect(artifact).toMatchObject({ + runId: started?.runId, + threadId: started?.threadId, + }) + await expect( + persistence.stores.artifacts!.list(started!.runId!), + ).resolves.toHaveLength(1) + }) + + it('does not persist generation artifacts when artifact stores are removed', async () => { + const full = memoryPersistence() + const put = vi.spyOn(full.stores.blobs, 'put') + const save = vi.spyOn(full.stores.artifacts, 'save') + const persistence = composePersistence(full, { + overrides: { artifacts: false, blobs: false }, + }) + + const result = await generateImage({ + adapter: imageAdapter(), + prompt: 'make an image', + threadId: 'thread-messages-only', + runId: 'run-messages-only', + middleware: [withGenerationPersistence(persistence)], + }) + + expect(result.artifacts).toBeUndefined() + expect(put).not.toHaveBeenCalled() + expect(save).not.toHaveBeenCalled() + }) + + it('fails early when artifact persistence is enabled without a paired blob store', () => { + const full = memoryPersistence() + const persistence: AIPersistence = defineAIPersistence({ + stores: { + artifacts: full.stores.artifacts, + }, + }) + + expect(() => withGenerationPersistence(persistence)).toThrow( + /artifact persistence requires both stores\.artifacts and stores\.blobs/i, + ) + }) + + it('persists transcription structured JSON output', async () => { + const persistence = memoryPersistence() + + const result = (await generateTranscription({ + adapter: transcriptionAdapter(), + audio: 'aW5wdXQtYXVkaW8=', + responseFormat: 'verbose_json', + threadId: 'thread-transcription', + runId: 'run-transcription', + middleware: [withGenerationPersistence(persistence)], + } as TranscriptionGenerateOptions)) as TranscriptionResult + + expect(result.artifacts?.map((artifact) => artifact.role)).toEqual([ + 'input', + 'output', + ]) + const structured = result.artifacts?.find( + (artifact) => artifact.source.mediaType === 'json', + ) as PersistedArtifactRef | undefined + expect(structured).toMatchObject({ + role: 'output', + mimeType: 'application/json', + source: { + activity: 'transcription', + path: 'transcription', + mediaType: 'json', + }, + }) + const blob = await persistence.stores.blobs!.get( + `artifacts/run-transcription/${structured!.artifactId}`, + ) + await expect(blob?.text()).resolves.toContain('"segments"') + }) +}) diff --git a/packages/ai-persistence/tests/memory.test.ts b/packages/ai-persistence/tests/memory.test.ts index 95820b9a8..30d7bc780 100644 --- a/packages/ai-persistence/tests/memory.test.ts +++ b/packages/ai-persistence/tests/memory.test.ts @@ -13,7 +13,10 @@ describe('memoryPersistence', () => { it('exposes the complete state store set (no locks)', () => { expect(Object.keys(memoryPersistence().stores).sort()).toEqual([ + 'artifacts', + 'blobs', 'interrupts', + 'jobs', 'messages', 'metadata', 'runs', diff --git a/packages/ai-persistence/tests/persistence-fixtures.ts b/packages/ai-persistence/tests/persistence-fixtures.ts index f86ec7ce5..81069c253 100644 --- a/packages/ai-persistence/tests/persistence-fixtures.ts +++ b/packages/ai-persistence/tests/persistence-fixtures.ts @@ -1,4 +1,6 @@ import type { + GenerationJobRecord, + GenerationJobStore, InterruptStore, MessageStore, MetadataStore, @@ -42,6 +44,33 @@ export function createRunStore(): RunStore { } } +export function createGenerationJobStore(): GenerationJobStore { + const jobs = new Map() + return { + createOrResume: (input) => { + const existing = jobs.get(input.jobId) + if (existing) return Promise.resolve(existing) + const record: GenerationJobRecord = { + jobId: input.jobId, + activity: input.activity, + provider: input.provider, + model: input.model, + status: input.status ?? 'running', + startedAt: input.startedAt, + ...(input.threadId !== undefined ? { threadId: input.threadId } : {}), + } + jobs.set(record.jobId, record) + return Promise.resolve(record) + }, + update: (jobId, patch) => { + const existing = jobs.get(jobId) + if (existing) jobs.set(jobId, { ...existing, ...patch }) + return Promise.resolve() + }, + get: (jobId) => Promise.resolve(jobs.get(jobId) ?? null), + } +} + export function createInterruptStore(): InterruptStore { return { create: () => Promise.resolve(), diff --git a/packages/ai-persistence/tests/persistence-types.test-d.ts b/packages/ai-persistence/tests/persistence-types.test-d.ts index 60f040141..6c30066d8 100644 --- a/packages/ai-persistence/tests/persistence-types.test-d.ts +++ b/packages/ai-persistence/tests/persistence-types.test-d.ts @@ -14,9 +14,12 @@ import { InMemoryLockStore, withLocks } from '@tanstack/ai/locks' import type { LockStore } from '@tanstack/ai/locks' import type { AIPersistence, + ArtifactStore, + BlobStore, ChatPersistence, ChatTranscriptPersistence, ChatTranscriptStores, + GenerationJobStore, InterruptStore, MessageStore, MetadataStore, @@ -33,6 +36,7 @@ declare const replacementRuns: RunStore & { } declare const interrupts: InterruptStore declare const metadata: MetadataStore +declare const jobs: GenerationJobStore declare const locks: LockStore const messagesOnly = defineAIPersistence({ stores: { messages } }) @@ -123,8 +127,18 @@ const uncertainInherited = composePersistence(base, { }) expectTypeOf(uncertainInherited.stores.messages).toEqualTypeOf() -// Named shapes -expectTypeOf(memoryPersistence()).toEqualTypeOf() +// Named shapes — memoryPersistence now includes generation stores too +expectTypeOf(memoryPersistence()).toEqualTypeOf< + AIPersistence<{ + messages: MessageStore + runs: RunStore + jobs: GenerationJobStore + interrupts: InterruptStore + metadata: MetadataStore + artifacts: ArtifactStore + blobs: BlobStore + }> +>() const transcript: ChatTranscriptPersistence = messagesOnly void transcript declare const fullChat: ChatPersistence @@ -138,10 +152,12 @@ withPersistence(defineAIPersistence({ stores: { runs } })) // @ts-expect-error a known interrupt store requires a known run store withPersistence(defineAIPersistence({ stores: { interrupts, messages } })) -// Generation requires runs -withGenerationPersistence(defineAIPersistence({ stores: { runs } })) -// @ts-expect-error generation persistence requires runs +// Generation requires jobs +withGenerationPersistence(defineAIPersistence({ stores: { jobs } })) +// @ts-expect-error generation persistence requires jobs withGenerationPersistence(messagesOnly) +// @ts-expect-error a runs store alone does not satisfy generation persistence +withGenerationPersistence(defineAIPersistence({ stores: { runs } })) const chatWithRemovedRuns = composePersistence(base, { overrides: { runs: false }, diff --git a/packages/ai-persistence/tests/persistence-validation.test.ts b/packages/ai-persistence/tests/persistence-validation.test.ts index 7b53bfddc..e504ed911 100644 --- a/packages/ai-persistence/tests/persistence-validation.test.ts +++ b/packages/ai-persistence/tests/persistence-validation.test.ts @@ -4,6 +4,7 @@ import { reconstructChat } from '../src/reconstruct' import { defineAIPersistence } from '../src/types' import type { ChatTranscriptPersistence } from '../src/types' import { + createGenerationJobStore, createInterruptStore, createMessageStore, createRunStore, @@ -60,7 +61,7 @@ describe('persistence store dependency validation', () => { expect(() => withPersistence(persistence)).not.toThrow() }) - it('rejects generation persistence without runs', () => { + it('rejects generation persistence without jobs', () => { const persistence = defineAIPersistence({ stores: { messages: createMessageStore() }, }) @@ -69,12 +70,12 @@ describe('persistence store dependency validation', () => { withGenerationPersistence( persistence as Parameters[0], ), - ).toThrow(/requires stores\.runs/i) + ).toThrow(/requires stores\.jobs/i) }) - it('allows generation persistence with runs', () => { + it('allows generation persistence with jobs', () => { const persistence = defineAIPersistence({ - stores: { runs: createRunStore() }, + stores: { jobs: createGenerationJobStore() }, }) expect(() => withGenerationPersistence(persistence)).not.toThrow() diff --git a/packages/ai-persistence/tests/reconstruct-generation.test.ts b/packages/ai-persistence/tests/reconstruct-generation.test.ts new file mode 100644 index 000000000..4f923a579 --- /dev/null +++ b/packages/ai-persistence/tests/reconstruct-generation.test.ts @@ -0,0 +1,138 @@ +import { describe, expect, it } from 'vitest' +import { memoryPersistence } from '../src/memory' +import { reconstructGeneration } from '../src/reconstruct-generation' +import type { ReconstructedGeneration } from '../src/reconstruct-generation' + +async function body(response: Response): Promise { + return (await response.json()) as ReconstructedGeneration +} + +describe('reconstructGeneration', () => { + it('maps a completed job to a resume snapshot', async () => { + const persistence = memoryPersistence() + await persistence.stores.jobs.createOrResume({ + jobId: 'job-done', + threadId: 'thread-1', + activity: 'image', + provider: 'test-image-provider', + model: 'test-image-model', + startedAt: 1000, + }) + await persistence.stores.jobs.update('job-done', { + status: 'complete', + finishedAt: 2000, + result: { id: 'image-result' }, + }) + + const response = await reconstructGeneration( + persistence, + new Request('http://example.test/api/generation?threadId=thread-1'), + ) + expect(response.status).toBe(200) + expect(response.headers.get('cache-control')).toBe('no-store') + + const parsed = await body(response) + expect(parsed).toEqual({ + resumeSnapshot: { + schemaVersion: 1, + resumeState: null, + status: 'complete', + result: { id: 'image-result' }, + activity: 'image', + }, + activeRun: null, + }) + }) + + it('reports an active run and resume state for a running job', async () => { + const persistence = memoryPersistence() + await persistence.stores.jobs.createOrResume({ + jobId: 'job-live', + threadId: 'thread-live', + activity: 'video', + provider: 'test-video-provider', + model: 'test-video-model', + startedAt: 5000, + }) + + // Resolve by jobId directly. + const parsed = await body( + await reconstructGeneration( + persistence, + new Request('http://example.test/api/generation?jobId=job-live'), + ), + ) + expect(parsed.activeRun).toEqual({ runId: 'job-live' }) + expect(parsed.resumeSnapshot).toMatchObject({ + status: 'running', + resumeState: { runId: 'job-live', threadId: 'thread-live' }, + activity: 'video', + }) + }) + + it('surfaces an interrupted job as error status', async () => { + const persistence = memoryPersistence() + await persistence.stores.jobs.createOrResume({ + jobId: 'job-int', + threadId: 'thread-int', + activity: 'audio', + provider: 'p', + model: 'm', + startedAt: 1, + }) + await persistence.stores.jobs.update('job-int', { + status: 'interrupted', + finishedAt: 2, + }) + + const parsed = await body( + await reconstructGeneration( + persistence, + new Request('http://example.test/api/generation?jobId=job-int'), + ), + ) + expect(parsed.resumeSnapshot?.status).toBe('error') + expect(parsed.resumeSnapshot?.resumeState).toBeNull() + expect(parsed.activeRun).toBeNull() + }) + + it('returns nulls when the id is missing or the thread is unknown', async () => { + const persistence = memoryPersistence() + + const missing = await body( + await reconstructGeneration( + persistence, + new Request('http://example.test/api/generation'), + ), + ) + expect(missing).toEqual({ resumeSnapshot: null, activeRun: null }) + + const unknown = await body( + await reconstructGeneration( + persistence, + new Request('http://example.test/api/generation?threadId=nope'), + ), + ) + expect(unknown).toEqual({ resumeSnapshot: null, activeRun: null }) + }) + + it('returns 403 when authorize returns false', async () => { + const persistence = memoryPersistence() + await persistence.stores.jobs.createOrResume({ + jobId: 'job-secret', + threadId: 'thread-secret', + activity: 'image', + provider: 'p', + model: 'm', + startedAt: 1, + }) + + const response = await reconstructGeneration( + persistence, + new Request('http://example.test/api/generation?jobId=job-secret'), + { authorize: () => false }, + ) + expect(response.status).toBe(403) + expect(await response.json()).toEqual({ error: 'Forbidden' }) + }) +}) diff --git a/packages/ai-react/src/index.ts b/packages/ai-react/src/index.ts index 75bff7c3c..47db566a2 100644 --- a/packages/ai-react/src/index.ts +++ b/packages/ai-react/src/index.ts @@ -107,4 +107,10 @@ export { type VideoGenerateInput, type VideoGenerateResult, type VideoStatusInfo, + type GenerationPersistence, + type GenerationResumeSnapshot, + type GenerationResumeState, + type GenerationResumeStatus, + type GenerationPendingArtifact, } from '@tanstack/ai-client' +export type { PersistedArtifactRef } from '@tanstack/ai/client' diff --git a/packages/ai-react/src/use-generate-audio.ts b/packages/ai-react/src/use-generate-audio.ts index a5f8223c0..af976c48a 100644 --- a/packages/ai-react/src/use-generate-audio.ts +++ b/packages/ai-react/src/use-generate-audio.ts @@ -1,4 +1,5 @@ import { useGeneration } from './use-generation' +import { reconstructAudioResult } from '@tanstack/ai-client' import type { AudioGenerationResult, StreamChunk } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -6,6 +7,9 @@ import type { ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, } from '@tanstack/ai-client' @@ -25,6 +29,21 @@ export interface UseGenerateAudioOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** + * How this generation persists across reloads. + * - Omit / `false`: ephemeral, in-memory only. + * - `true`: server-driven — on mount the client hydrates the last generation + * for its `threadId` from the server (needs a connection with a + * `hydrateGeneration` handler) and repaints it; it never auto-starts a run. + * - a storage adapter: client-driven — the lightweight snapshot is cached under + * `generation:` as a run streams and read back on mount. Media bytes are + * never stored. + */ + persistence?: boolean | GenerationPersistence + /** Thread id for this generation, stable across reloads. Server-driven mode (`persistence: true`) hydrates the last generation under this key. Falls back to `id`. */ + threadId?: string + /** Explicit resume-snapshot seed for apps that manage storage themselves; skips automatic hydration from `persistence`. Later run events merge into it. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when audio is generated. Can optionally return a transformed value. * @@ -61,6 +80,8 @@ export interface UseGenerateAudioReturn { stop: () => void /** Clear result, error, and return to idle */ reset: () => void + /** Identity of the in-flight run while one is streaming, or null after it ends */ + resumeState: GenerationResumeState | null } /** @@ -106,19 +127,15 @@ export function useGenerateAudio( hookName: 'useGenerateAudio', outputKind: 'audio' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - useGeneration({ - ...options, - devtools, - }) + const generation = useGeneration< + AudioGenerateInput, + AudioGenerationResult, + TTransformed + >({ + ...options, + devtools, + reconstructResult: reconstructAudioResult, + }) - return { - generate: generate as (input: AudioGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + return generation } diff --git a/packages/ai-react/src/use-generate-image.ts b/packages/ai-react/src/use-generate-image.ts index 078a792d1..693ff2f7c 100644 --- a/packages/ai-react/src/use-generate-image.ts +++ b/packages/ai-react/src/use-generate-image.ts @@ -1,10 +1,14 @@ import { useGeneration } from './use-generation' +import { reconstructImageResult } from '@tanstack/ai-client' import type { ImageGenerationResult, StreamChunk } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, ImageGenerateInput, InferGenerationOutputFromReturn, } from '@tanstack/ai-client' @@ -25,6 +29,21 @@ export interface UseGenerateImageOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** + * How this generation persists across reloads. + * - Omit / `false`: ephemeral, in-memory only. + * - `true`: server-driven — on mount the client hydrates the last generation + * for its `threadId` from the server (needs a connection with a + * `hydrateGeneration` handler) and repaints it; it never auto-starts a run. + * - a storage adapter: client-driven — the lightweight snapshot is cached under + * `generation:` as a run streams and read back on mount. Media bytes are + * never stored. + */ + persistence?: boolean | GenerationPersistence + /** Thread id for this generation, stable across reloads. Server-driven mode (`persistence: true`) hydrates the last generation under this key. Falls back to `id`. */ + threadId?: string + /** Explicit resume-snapshot seed for apps that manage storage themselves; skips automatic hydration from `persistence`. Later run events merge into it. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when images are generated. Can optionally return a transformed value. * @@ -61,6 +80,8 @@ export interface UseGenerateImageReturn { stop: () => void /** Clear result, error, and return to idle */ reset: () => void + /** Identity of the in-flight run while one is streaming, or null after it ends */ + resumeState: GenerationResumeState | null } /** @@ -108,19 +129,15 @@ export function useGenerateImage( hookName: 'useGenerateImage', outputKind: 'image' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - useGeneration({ - ...options, - devtools, - }) + const generation = useGeneration< + ImageGenerateInput, + ImageGenerationResult, + TTransformed + >({ + ...options, + devtools, + reconstructResult: reconstructImageResult, + }) - return { - generate: generate as (input: ImageGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + return generation } diff --git a/packages/ai-react/src/use-generate-speech.ts b/packages/ai-react/src/use-generate-speech.ts index b22199ff5..3afe373b0 100644 --- a/packages/ai-react/src/use-generate-speech.ts +++ b/packages/ai-react/src/use-generate-speech.ts @@ -5,6 +5,9 @@ import type { ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, SpeechGenerateInput, } from '@tanstack/ai-client' @@ -25,6 +28,21 @@ export interface UseGenerateSpeechOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** + * How this generation persists across reloads. + * - Omit / `false`: ephemeral, in-memory only. + * - `true`: server-driven — on mount the client hydrates the last generation + * for its `threadId` from the server (needs a connection with a + * `hydrateGeneration` handler) and repaints it; it never auto-starts a run. + * - a storage adapter: client-driven — the lightweight snapshot is cached under + * `generation:` as a run streams and read back on mount. Media bytes are + * never stored. + */ + persistence?: boolean | GenerationPersistence + /** Thread id for this generation, stable across reloads. Server-driven mode (`persistence: true`) hydrates the last generation under this key. Falls back to `id`. */ + threadId?: string + /** Explicit resume-snapshot seed for apps that manage storage themselves; skips automatic hydration from `persistence`. Later run events merge into it. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when speech is generated. Can optionally return a transformed value. * @@ -61,6 +79,8 @@ export interface UseGenerateSpeechReturn { stop: () => void /** Clear result, error, and return to idle */ reset: () => void + /** Identity of the in-flight run while one is streaming, or null after it ends */ + resumeState: GenerationResumeState | null } /** @@ -102,19 +122,14 @@ export function useGenerateSpeech( hookName: 'useGenerateSpeech', outputKind: 'audio' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - useGeneration({ - ...options, - devtools, - }) + const generation = useGeneration< + SpeechGenerateInput, + TTSResult, + TTransformed + >({ + ...options, + devtools, + }) - return { - generate: generate as (input: SpeechGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + return generation } diff --git a/packages/ai-react/src/use-generate-video.ts b/packages/ai-react/src/use-generate-video.ts index 3aece73c5..688b37a3a 100644 --- a/packages/ai-react/src/use-generate-video.ts +++ b/packages/ai-react/src/use-generate-video.ts @@ -7,6 +7,9 @@ import type { ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, VideoGenerateInput, VideoGenerateResult, @@ -27,6 +30,21 @@ export interface UseGenerateVideoOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** + * How this generation persists across reloads. + * - Omit / `false`: ephemeral, in-memory only. + * - `true`: server-driven — on mount the client hydrates the last generation + * for its `threadId` from the server (needs a connection with a + * `hydrateGeneration` handler) and repaints it; it never auto-starts a run. + * - a storage adapter: client-driven — the lightweight snapshot is cached under + * `generation:` as a run streams and read back on mount. Media bytes are + * never stored. + */ + persistence?: boolean | GenerationPersistence + /** Thread id for this generation, stable across reloads. Server-driven mode (`persistence: true`) hydrates the last generation under this key. Falls back to `id`. */ + threadId?: string + /** Explicit resume-snapshot seed for apps that manage storage themselves; skips automatic hydration from `persistence`. Later run events merge into it. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when video generation completes. Can optionally return a transformed value. * @@ -71,6 +89,8 @@ export interface UseGenerateVideoReturn { stop: () => void /** Clear all state and return to idle */ reset: () => void + /** Identity of the in-flight run while one is streaming, or null after it ends */ + resumeState: GenerationResumeState | null } /** @@ -127,9 +147,13 @@ export function useGenerateVideo( const [isLoading, setIsLoading] = useState(false) const [error, setError] = useState(undefined) const [status, setStatus] = useState('idle') + const [resumeState, setResumeState] = useState( + options.initialResumeSnapshot?.resumeState ?? null, + ) const optionsRef = useRef(options) optionsRef.current = options + const disposedRef = useRef(false) const client = useMemo(() => { const opts = optionsRef.current @@ -142,6 +166,11 @@ export function useGenerateVideo( const baseOptions = { id: clientId, body: opts.body, + ...(opts.persistence !== undefined && { persistence: opts.persistence }), + ...(opts.threadId !== undefined && { threadId: opts.threadId }), + ...(opts.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: opts.initialResumeSnapshot, + }), devtoolsBridgeFactory: createVideoDevtoolsBridge, devtools: { ...opts.devtools, @@ -157,26 +186,41 @@ export function useGenerateVideo( result: VideoGenerateResult, ) => TOutput | null | void, onError: (e: Error) => { - optionsRef.current.onError?.(e) + if (!disposedRef.current) optionsRef.current.onError?.(e) }, onProgress: (p: number, m?: string) => { - optionsRef.current.onProgress?.(p, m) + if (!disposedRef.current) optionsRef.current.onProgress?.(p, m) }, onChunk: (c: StreamChunk) => { - optionsRef.current.onChunk?.(c) + if (!disposedRef.current) optionsRef.current.onChunk?.(c) }, onJobCreated: (id: string) => { - optionsRef.current.onJobCreated?.(id) + if (!disposedRef.current) optionsRef.current.onJobCreated?.(id) }, onStatusUpdate: (s: VideoStatusInfo) => { - optionsRef.current.onStatusUpdate?.(s) + if (!disposedRef.current) optionsRef.current.onStatusUpdate?.(s) + }, + onResultChange: (r: TOutput | null) => { + if (!disposedRef.current) setResult(r) + }, + onLoadingChange: (l: boolean) => { + if (!disposedRef.current) setIsLoading(l) + }, + onErrorChange: (e: Error | undefined) => { + if (!disposedRef.current) setError(e) + }, + onStatusChange: (s: GenerationClientState) => { + if (!disposedRef.current) setStatus(s) + }, + onJobIdChange: (id: string | null) => { + if (!disposedRef.current) setJobId(id) + }, + onVideoStatusChange: (s: VideoStatusInfo | null) => { + if (!disposedRef.current) setVideoStatus(s) + }, + onResumeStateChange: (rs: GenerationResumeState | null) => { + if (!disposedRef.current) setResumeState(rs) }, - onResultChange: setResult, - onLoadingChange: setIsLoading, - onErrorChange: setError, - onStatusChange: setStatus, - onJobIdChange: setJobId, - onVideoStatusChange: setVideoStatus, } if (opts.connection) { @@ -206,11 +250,15 @@ export function useGenerateVideo( }) }, [client, options.body]) - // Cleanup on unmount + // Mount devtools and clean up on unmount. Generation runs are never + // auto-started on mount — persisted state is only displayed. Mounting + // revives the client after a StrictMode dispose → remount replay. useEffect(() => { + disposedRef.current = false client.mountDevtools() return () => { + disposedRef.current = true client.dispose() } }, [client]) @@ -240,5 +288,6 @@ export function useGenerateVideo( status, stop, reset, + resumeState, } } diff --git a/packages/ai-react/src/use-generation.ts b/packages/ai-react/src/use-generation.ts index 6e5557d9b..0d00bbae6 100644 --- a/packages/ai-react/src/use-generation.ts +++ b/packages/ai-react/src/use-generation.ts @@ -8,6 +8,10 @@ import type { GenerationClientOptions, GenerationClientState, GenerationFetcher, + GenerationPersistence, + GenerationRestoredResult, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, } from '@tanstack/ai-client' @@ -31,6 +35,21 @@ export interface UseGenerationOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** + * How this generation persists across reloads. + * - Omit / `false`: ephemeral, in-memory only. + * - `true`: server-driven — on mount the client hydrates the last generation + * for its `threadId` from the server (needs a connection with a + * `hydrateGeneration` handler) and repaints it; it never auto-starts a run. + * - a storage adapter: client-driven — the lightweight snapshot is cached under + * `generation:` as a run streams and read back on mount. Media bytes are + * never stored. + */ + persistence?: boolean | GenerationPersistence + /** Thread id for this generation, stable across reloads. Server-driven mode (`persistence: true`) hydrates the last generation under this key. Falls back to `id`. */ + threadId?: string + /** Explicit resume-snapshot seed for apps that manage storage themselves; skips automatic hydration from `persistence`. Later run events merge into it. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when a result is received. Can optionally return a transformed value. * @@ -45,16 +64,26 @@ export interface UseGenerationOptions { onProgress?: (progress: number, message?: string) => void /** Callback for each stream chunk (connect-based adapter mode only) */ onChunk?: (chunk: StreamChunk) => void + /** + * @internal Rebuild a typed result from a restored snapshot, injected by each + * specialized hook (image / audio / transcription / summarize). Forwarded to + * the client so a client-store / server-hydrate restore repaints `result`. + */ + reconstructResult?: (restored: GenerationRestoredResult) => TResult | null } /** * Return type for the useGeneration hook. * * @template TOutput - The output type (after optional transform) + * @template TInput - The input type accepted by `generate` (defaults to any object) */ -export interface UseGenerationReturn { +export interface UseGenerationReturn< + TOutput, + TInput extends Record = Record, +> { /** Trigger a generation request */ - generate: (input: Record) => Promise + generate: (input: TInput) => Promise /** The generation result, or null if not yet generated */ result: TOutput | null /** Whether a generation is currently in progress */ @@ -67,6 +96,8 @@ export interface UseGenerationReturn { stop: () => void /** Clear result, error, and return to idle */ reset: () => void + /** Identity of the in-flight run while one is streaming, or null after it ends */ + resumeState: GenerationResumeState | null } /** @@ -102,7 +133,10 @@ export function useGeneration< options: Omit, 'onResult'> & { onResult?: (result: TResult) => TTransformed }, -): UseGenerationReturn> { +): UseGenerationReturn< + InferGenerationOutputFromReturn, + TInput +> { type TOutput = InferGenerationOutputFromReturn const hookId = useId() const clientId = options.id || hookId @@ -111,9 +145,13 @@ export function useGeneration< const [isLoading, setIsLoading] = useState(false) const [error, setError] = useState(undefined) const [status, setStatus] = useState('idle') + const [resumeState, setResumeState] = useState( + options.initialResumeSnapshot?.resumeState ?? null, + ) const optionsRef = useRef(options) optionsRef.current = options + const disposedRef = useRef(false) const client = useMemo(() => { const opts = optionsRef.current @@ -125,6 +163,14 @@ export function useGeneration< const clientOptions: GenerationClientOptions = { id: clientId, body: opts.body, + ...(opts.persistence !== undefined && { persistence: opts.persistence }), + ...(opts.threadId !== undefined && { threadId: opts.threadId }), + ...(opts.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: opts.initialResumeSnapshot, + }), + ...(opts.reconstructResult + ? { reconstructResult: opts.reconstructResult } + : {}), devtoolsBridgeFactory: createGenerationDevtoolsBridge, devtools: { hookName: 'useGeneration', @@ -138,18 +184,29 @@ export function useGeneration< result: TResult, ) => TOutput | null | void, onError: (e: Error) => { - optionsRef.current.onError?.(e) + if (!disposedRef.current) optionsRef.current.onError?.(e) }, onProgress: (p: number, m?: string) => { - optionsRef.current.onProgress?.(p, m) + if (!disposedRef.current) optionsRef.current.onProgress?.(p, m) }, onChunk: (c: StreamChunk) => { - optionsRef.current.onChunk?.(c) + if (!disposedRef.current) optionsRef.current.onChunk?.(c) + }, + onResultChange: (r) => { + if (!disposedRef.current) setResult(r) + }, + onLoadingChange: (l) => { + if (!disposedRef.current) setIsLoading(l) + }, + onErrorChange: (e) => { + if (!disposedRef.current) setError(e) + }, + onStatusChange: (s) => { + if (!disposedRef.current) setStatus(s) + }, + onResumeStateChange: (rs) => { + if (!disposedRef.current) setResumeState(rs) }, - onResultChange: setResult, - onLoadingChange: setIsLoading, - onErrorChange: setError, - onStatusChange: setStatus, } if (opts.connection) { @@ -179,11 +236,15 @@ export function useGeneration< }) }, [client, options.body]) - // Cleanup on unmount + // Mount devtools and clean up on unmount. Generation runs are never + // auto-started on mount — persisted state is only displayed. Mounting + // revives the client after a StrictMode dispose → remount replay. useEffect(() => { + disposedRef.current = false client.mountDevtools() return () => { + disposedRef.current = true client.dispose() } }, [client]) @@ -204,12 +265,13 @@ export function useGeneration< }, [client]) return { - generate: generate as (input: Record) => Promise, + generate, result, isLoading, error, status, stop, reset, + resumeState, } } diff --git a/packages/ai-react/src/use-summarize.ts b/packages/ai-react/src/use-summarize.ts index 9de8eeb71..03703fa32 100644 --- a/packages/ai-react/src/use-summarize.ts +++ b/packages/ai-react/src/use-summarize.ts @@ -1,10 +1,14 @@ import { useGeneration } from './use-generation' +import { reconstructSummarizeResult } from '@tanstack/ai-client' import type { StreamChunk, SummarizationResult } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, SummarizeGenerateInput, } from '@tanstack/ai-client' @@ -25,6 +29,21 @@ export interface UseSummarizeOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** + * How this generation persists across reloads. + * - Omit / `false`: ephemeral, in-memory only. + * - `true`: server-driven — on mount the client hydrates the last generation + * for its `threadId` from the server (needs a connection with a + * `hydrateGeneration` handler) and repaints it; it never auto-starts a run. + * - a storage adapter: client-driven — the lightweight snapshot is cached under + * `generation:` as a run streams and read back on mount. Media bytes are + * never stored. + */ + persistence?: boolean | GenerationPersistence + /** Thread id for this generation, stable across reloads. Server-driven mode (`persistence: true`) hydrates the last generation under this key. Falls back to `id`. */ + threadId?: string + /** Explicit resume-snapshot seed for apps that manage storage themselves; skips automatic hydration from `persistence`. Later run events merge into it. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when summarization is complete. Can optionally return a transformed value. * @@ -61,6 +80,8 @@ export interface UseSummarizeReturn { stop: () => void /** Clear result, error, and return to idle */ reset: () => void + /** Identity of the in-flight run while one is streaming, or null after it ends */ + resumeState: GenerationResumeState | null } /** @@ -105,19 +126,15 @@ export function useSummarize( hookName: 'useSummarize', outputKind: 'text' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - useGeneration({ - ...options, - devtools, - }) + const generation = useGeneration< + SummarizeGenerateInput, + SummarizationResult, + TTransformed + >({ + ...options, + devtools, + reconstructResult: reconstructSummarizeResult, + }) - return { - generate: generate as (input: SummarizeGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + return generation } diff --git a/packages/ai-react/src/use-transcription.ts b/packages/ai-react/src/use-transcription.ts index d0e01df8c..316d50871 100644 --- a/packages/ai-react/src/use-transcription.ts +++ b/packages/ai-react/src/use-transcription.ts @@ -1,10 +1,14 @@ import { useGeneration } from './use-generation' +import { reconstructTranscriptionResult } from '@tanstack/ai-client' import type { StreamChunk, TranscriptionResult } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, TranscriptionGenerateInput, } from '@tanstack/ai-client' @@ -25,6 +29,21 @@ export interface UseTranscriptionOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** + * How this generation persists across reloads. + * - Omit / `false`: ephemeral, in-memory only. + * - `true`: server-driven — on mount the client hydrates the last generation + * for its `threadId` from the server (needs a connection with a + * `hydrateGeneration` handler) and repaints it; it never auto-starts a run. + * - a storage adapter: client-driven — the lightweight snapshot is cached under + * `generation:` as a run streams and read back on mount. Media bytes are + * never stored. + */ + persistence?: boolean | GenerationPersistence + /** Thread id for this generation, stable across reloads. Server-driven mode (`persistence: true`) hydrates the last generation under this key. Falls back to `id`. */ + threadId?: string + /** Explicit resume-snapshot seed for apps that manage storage themselves; skips automatic hydration from `persistence`. Later run events merge into it. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when transcription is complete. Can optionally return a transformed value. * @@ -61,6 +80,8 @@ export interface UseTranscriptionReturn { stop: () => void /** Clear result, error, and return to idle */ reset: () => void + /** Identity of the in-flight run while one is streaming, or null after it ends */ + resumeState: GenerationResumeState | null } /** @@ -110,20 +131,11 @@ export function useTranscription( hookName: 'useTranscription', outputKind: 'text' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - useGeneration< - TranscriptionGenerateInput, - TranscriptionResult, - TTransformed - >({ ...options, devtools }) + const generation = useGeneration< + TranscriptionGenerateInput, + TranscriptionResult, + TTransformed + >({ ...options, devtools, reconstructResult: reconstructTranscriptionResult }) - return { - generate: generate as (input: TranscriptionGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + return generation } diff --git a/packages/ai-react/tests/use-generation.test.ts b/packages/ai-react/tests/use-generation.test.ts index f3540b520..dc6693602 100644 --- a/packages/ai-react/tests/use-generation.test.ts +++ b/packages/ai-react/tests/use-generation.test.ts @@ -1,3 +1,4 @@ +import { StrictMode } from 'react' import { renderHook, waitFor, act } from '@testing-library/react' import { describe, expect, expectTypeOf, it, vi } from 'vitest' import { useGeneration } from '../src/use-generation' @@ -8,8 +9,19 @@ import { useTranscription } from '../src/use-transcription' import { useSummarize } from '../src/use-summarize' import { useGenerateVideo } from '../src/use-generate-video' import { createMockConnectionAdapter } from './test-utils' -import type { StreamChunk, TTSResult, TranscriptionResult } from '@tanstack/ai' +import type { + PersistedArtifactRef, + StreamChunk, + TTSResult, + TranscriptionResult, +} from '@tanstack/ai' import { EventType } from '@tanstack/ai' +import type { + ConnectConnectionAdapter, + GenerationResumeSnapshot, + GenerationPersistence, + RunAgentInputContext, +} from '@tanstack/ai-client' // Helper to create generation stream chunks function createGenerationChunks(result: unknown): Array { @@ -71,6 +83,87 @@ function createVideoChunks(jobId: string, url: string): Array { ] } +const videoResumeSnapshot: GenerationResumeSnapshot = { + resumeState: { + threadId: 'thread-resume', + runId: 'run-resume', + }, + status: 'running', +} + +const replayedVideoArtifact: PersistedArtifactRef = { + role: 'output', + artifactId: 'artifact-video-1', + threadId: 'thread-resume', + runId: 'run-resume', + name: 'video.mp4', + mimeType: 'video/mp4', + size: 1234, + createdAt: '2026-07-06T00:00:00.000Z', + externalUrl: 'https://example.com/video.mp4', + source: { + activity: 'video', + path: 'runs/run-resume/video.mp4', + provider: 'test', + model: 'test-video', + mediaType: 'video', + jobId: 'job-replay', + expiresAt: '2026-07-07T00:00:00.000Z', + }, +} + +function createReplayVideoChunks(): Array { + return [ + { + type: EventType.RUN_STARTED, + runId: 'run-resume', + threadId: 'thread-resume', + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: 'generation:result', + value: { + jobId: 'job-replay', + status: 'completed', + url: 'https://example.com/video.mp4', + artifacts: [replayedVideoArtifact], + }, + timestamp: Date.now(), + }, + { + type: EventType.RUN_FINISHED, + runId: 'run-resume', + threadId: 'thread-resume', + timestamp: Date.now(), + }, + ] +} + +function createRunContextCaptureAdapter(chunks: Array): { + adapter: ConnectConnectionAdapter + connect: ReturnType + runContexts: Array +} { + const runContexts: Array = [] + const connect = vi.fn() + const adapter: ConnectConnectionAdapter = { + async *connect(_messages, _data, _signal, runContext) { + connect(runContext) + runContexts.push(runContext) + for (const chunk of chunks) { + yield chunk + } + }, + } + return { adapter, connect, runContexts } +} + +async function flushPromises(): Promise { + await Promise.resolve() + await new Promise((resolve) => setTimeout(resolve, 0)) +} + // Helper to create error stream chunks. // NOTE: The AG-UI spec for RUN_ERROR carries `message` directly on the event // (not nested under `error`). We emit BOTH shapes here because GenerationClient @@ -276,6 +369,191 @@ describe('useGeneration', () => { // Resolve the promise after unmount — should not cause errors resolvePromise!({ id: '1' }) }) + + it('does not auto-fire a generation on mount from a persisted running snapshot', async () => { + // Regression guard for the removed generation resume surface: mounting a + // generation hook that has server persistence and a persisted `running` + // snapshot must NOT start a fresh empty-prompt generation (previously + // `maybeAutoResume()` -> `resume()` -> `generate({})` fired here). + const { adapter, connect } = createRunContextCaptureAdapter( + createGenerationChunks({ id: '1' }), + ) + const persistence: GenerationPersistence = { + getItem: vi.fn(() => ({ + resumeState: { threadId: 'thread-resume', runId: 'run-resume' }, + status: 'running' as const, + })), + setItem: vi.fn(), + removeItem: vi.fn(), + } + + const { result } = renderHook(() => + useGeneration({ + id: 'no-auto-fire', + connection: adapter, + persistence: persistence, + initialResumeSnapshot: { + resumeState: { threadId: 'thread-resume', runId: 'run-resume' }, + status: 'running', + }, + }), + ) + + await act(async () => { + await flushPromises() + }) + + expect(connect).not.toHaveBeenCalled() + // The explicit initialResumeSnapshot seed takes precedence over storage, + // so automatic hydration (getItem) is skipped entirely. + expect(persistence.getItem).not.toHaveBeenCalled() + expect(result.current.isLoading).toBe(false) + expect(result.current.status).toBe('idle') + // The persisted snapshot is still exposed as display state. + expect(result.current.resumeState).toEqual({ + threadId: 'thread-resume', + runId: 'run-resume', + }) + }) + }) + + describe('persistence', () => { + it('repaints status from a stored complete snapshot on mount without starting a run', async () => { + const { adapter, connect } = createRunContextCaptureAdapter( + createGenerationChunks({ id: '1' }), + ) + const persistence: GenerationPersistence = { + getItem: vi.fn(() => ({ + resumeState: null, + status: 'complete' as const, + result: { id: 'result-1', model: 'image-model' }, + })), + setItem: vi.fn(), + removeItem: vi.fn(), + } + + const { result } = renderHook(() => + useGeneration({ + id: 'hydrated', + connection: adapter, + persistence: persistence, + }), + ) + + // A completed snapshot repaints the normal `status` field to `success`. + await waitFor(() => { + expect(result.current.status).toBe('success') + }) + expect(persistence.getItem).toHaveBeenCalledWith('generation:hydrated') + expect(connect).not.toHaveBeenCalled() + // The base hook injects no reconstructResult, so `result` stays null. + expect(result.current.result).toBeNull() + expect(result.current.resumeState).toBeNull() + }) + + it('repaints resumeState from a stored running snapshot on mount', async () => { + const { adapter, connect } = createRunContextCaptureAdapter( + createGenerationChunks({ id: '1' }), + ) + const persistence: GenerationPersistence = { + getItem: vi.fn(() => ({ + resumeState: { threadId: 'thread-resume', runId: 'run-resume' }, + status: 'running' as const, + })), + setItem: vi.fn(), + removeItem: vi.fn(), + } + + const { result } = renderHook(() => + useGeneration({ + id: 'running-hydrate', + connection: adapter, + persistence: persistence, + }), + ) + + await waitFor(() => { + expect(result.current.resumeState).toEqual({ + threadId: 'thread-resume', + runId: 'run-resume', + }) + }) + // A restored running run presents as `generating` but never auto-tails. + expect(result.current.status).toBe('generating') + expect(result.current.isLoading).toBe(false) + expect(connect).not.toHaveBeenCalled() + }) + + it('server-driven (persistence: true) hydrates from the connection by threadId without a local store', async () => { + const { adapter, connect } = createRunContextCaptureAdapter( + createGenerationChunks({ id: '1' }), + ) + const hydrateGeneration = vi.fn(async () => ({ + resumeSnapshot: { + schemaVersion: 1 as const, + resumeState: null, + status: 'complete' as const, + result: { id: 'server-result', model: 'image-model' }, + }, + activeRun: null, + })) + + const { result } = renderHook(() => + useGeneration({ + threadId: 'thread-server', + connection: { ...adapter, hydrateGeneration }, + persistence: true, + }), + ) + + // The server snapshot repaints `status` to `success`. + await waitFor(() => { + expect(result.current.status).toBe('success') + }) + expect(hydrateGeneration).toHaveBeenCalledWith('thread-server') + expect(connect).not.toHaveBeenCalled() + }) + + it('generates normally under React StrictMode (dispose → remount replay)', async () => { + const { adapter, connect } = createRunContextCaptureAdapter( + createGenerationChunks({ id: 'strict-1' }), + ) + + const { result } = renderHook( + () => useGeneration({ connection: adapter }), + { wrapper: StrictMode }, + ) + + await act(async () => { + await result.current.generate({ prompt: 'test' }) + }) + + expect(connect).toHaveBeenCalledTimes(1) + await waitFor(() => { + expect(result.current.status).toBe('success') + expect(result.current.result).toEqual({ id: 'strict-1' }) + }) + }) + + it('clears resumeState once a streamed run finishes', async () => { + const adapter = createMockConnectionAdapter({ + chunks: createReplayVideoChunks(), + }) + + const { result } = renderHook(() => + useGeneration({ connection: adapter }), + ) + + await act(async () => { + await result.current.generate({ prompt: 'replay' }) + }) + + await waitFor(() => { + expect(result.current.status).toBe('success') + }) + // Once the run ends the in-flight identity is gone. + expect(result.current.resumeState).toBeNull() + }) }) }) @@ -362,6 +640,64 @@ describe('useGenerateImage', () => { expect(typeof result.current.stop).toBe('function') expect(typeof result.current.reset).toBe('function') }) + + it('restores a completed image result from a durable artifact url', async () => { + // useGenerateImage injects `reconstructImageResult`, so a restored complete + // snapshot repaints `result` with the durable serve url — as if the run had + // just finished. + const artifact: PersistedArtifactRef = { + role: 'output', + artifactId: 'artifact-image-1', + threadId: 'thread-img', + runId: 'run-img', + name: 'image.png', + mimeType: 'image/png', + size: 2048, + createdAt: '2026-07-06T00:00:00.000Z', + url: '/api/artifacts/artifact-image-1', + source: { + activity: 'image', + path: 'runs/run-img/image.png', + provider: 'test', + model: 'test-image', + mediaType: 'image', + }, + } + const persistence: GenerationPersistence = { + getItem: vi.fn(() => ({ + resumeState: null, + status: 'complete' as const, + activity: 'image' as const, + result: { + id: 'img-restored', + model: 'test-image', + artifacts: [artifact], + }, + })), + setItem: vi.fn(), + removeItem: vi.fn(), + } + const adapter = createMockConnectionAdapter() + + const { result } = renderHook(() => + useGenerateImage({ + id: 'img-hydrate', + connection: adapter, + persistence: persistence, + }), + ) + + await waitFor(() => { + expect(result.current.status).toBe('success') + }) + expect(result.current.result).toEqual({ + id: 'img-restored', + model: 'test-image', + images: [{ url: '/api/artifacts/artifact-image-1' }], + artifacts: [artifact], + }) + expect(result.current.resumeState).toBeNull() + }) }) describe('useGenerateSpeech', () => { @@ -535,7 +871,7 @@ describe('useSummarize', () => { const mockResult = { id: 'sum-1', summary: 'A brief summary', - model: 'gpt-4', + model: 'gpt-5.5', usage: { promptTokens: 100, completionTokens: 20, totalTokens: 120 }, } @@ -554,7 +890,7 @@ describe('useSummarize', () => { }) it('should summarize text using connection', async () => { - const mockResult = { summary: 'A brief summary', model: 'gpt-4' } + const mockResult = { summary: 'A brief summary', model: 'gpt-5.5' } const chunks = createGenerationChunks(mockResult) const adapter = createMockConnectionAdapter({ chunks }) @@ -681,6 +1017,38 @@ describe('useGenerateVideo', () => { expect(result.current.videoStatus).toBeNull() expect(result.current.status).toBe('idle') }) + + it('does not auto-fire a video generation on mount from a persisted running snapshot', async () => { + // Regression guard for the removed generation resume surface (video). + const { adapter, connect } = createRunContextCaptureAdapter( + createReplayVideoChunks(), + ) + const persistence: GenerationPersistence = { + getItem: vi.fn(() => videoResumeSnapshot), + setItem: vi.fn(), + removeItem: vi.fn(), + } + + const { result } = renderHook(() => + useGenerateVideo({ + id: 'video-no-auto-fire', + connection: adapter, + persistence: persistence, + initialResumeSnapshot: videoResumeSnapshot, + }), + ) + + await act(async () => { + await flushPromises() + }) + + expect(connect).not.toHaveBeenCalled() + expect(persistence.getItem).not.toHaveBeenCalled() + expect(result.current.isLoading).toBe(false) + expect(result.current.status).toBe('idle') + // The seeded in-flight identity is exposed as read-only `resumeState`. + expect(result.current.resumeState).toEqual(videoResumeSnapshot.resumeState) + }) }) describe('onResult transform', () => { diff --git a/packages/ai-solid/src/index.ts b/packages/ai-solid/src/index.ts index 86b39b403..ae43ec04c 100644 --- a/packages/ai-solid/src/index.ts +++ b/packages/ai-solid/src/index.ts @@ -94,4 +94,10 @@ export { type VideoGenerateInput, type VideoGenerateResult, type VideoStatusInfo, + type GenerationPersistence, + type GenerationResumeSnapshot, + type GenerationResumeState, + type GenerationResumeStatus, + type GenerationPendingArtifact, } from '@tanstack/ai-client' +export type { PersistedArtifactRef } from '@tanstack/ai/client' diff --git a/packages/ai-solid/src/use-generate-audio.ts b/packages/ai-solid/src/use-generate-audio.ts index 94bfac297..efafa8943 100644 --- a/packages/ai-solid/src/use-generate-audio.ts +++ b/packages/ai-solid/src/use-generate-audio.ts @@ -1,4 +1,9 @@ import { useGeneration } from './use-generation' +import { reconstructAudioResult } from '@tanstack/ai-client' +import type { + UseGenerationOptions, + UseGenerationReturn, +} from './use-generation' import type { AudioGenerationResult, StreamChunk } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -15,7 +20,12 @@ import type { Accessor } from 'solid-js' * * @template TOutput - The transformed output type (defaults to AudioGenerationResult) */ -export interface UseGenerateAudioOptions { +export interface UseGenerateAudioOptions< + TOutput = AudioGenerationResult, +> extends Pick< + UseGenerationOptions, + 'persistence' | 'threadId' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for audio generation */ @@ -47,7 +57,9 @@ export interface UseGenerateAudioOptions { * * @template TOutput - The transformed output type (defaults to AudioGenerationResult) */ -export interface UseGenerateAudioReturn { +export interface UseGenerateAudioReturn< + TOutput = AudioGenerationResult, +> extends Omit, 'generate'> { /** Trigger audio generation */ generate: (input: AudioGenerateInput) => Promise /** The generation result containing audio, or null */ @@ -58,10 +70,6 @@ export interface UseGenerateAudioReturn { error: Accessor /** Current state of the generation */ status: Accessor - /** Abort the current generation */ - stop: () => void - /** Clear result, error, and return to idle */ - reset: () => void } /** @@ -101,19 +109,15 @@ export function useGenerateAudio( hookName: 'useGenerateAudio', outputKind: 'audio' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - useGeneration({ - ...options, - devtools, - }) + const generation = useGeneration< + AudioGenerateInput, + AudioGenerationResult, + TTransformed + >({ + ...options, + devtools, + reconstructResult: reconstructAudioResult, + }) - return { - generate: generate as (input: AudioGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + return generation } diff --git a/packages/ai-solid/src/use-generate-image.ts b/packages/ai-solid/src/use-generate-image.ts index b88902163..40f226604 100644 --- a/packages/ai-solid/src/use-generate-image.ts +++ b/packages/ai-solid/src/use-generate-image.ts @@ -1,4 +1,9 @@ import { useGeneration } from './use-generation' +import { reconstructImageResult } from '@tanstack/ai-client' +import type { + UseGenerationOptions, + UseGenerationReturn, +} from './use-generation' import type { ImageGenerationResult, StreamChunk } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -15,7 +20,12 @@ import type { Accessor } from 'solid-js' * * @template TOutput - The transformed output type (defaults to ImageGenerationResult) */ -export interface UseGenerateImageOptions { +export interface UseGenerateImageOptions< + TOutput = ImageGenerationResult, +> extends Pick< + UseGenerationOptions, + 'persistence' | 'threadId' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for image generation */ @@ -47,7 +57,9 @@ export interface UseGenerateImageOptions { * * @template TOutput - The transformed output type (defaults to ImageGenerationResult) */ -export interface UseGenerateImageReturn { +export interface UseGenerateImageReturn< + TOutput = ImageGenerationResult, +> extends Omit, 'generate'> { /** Trigger image generation */ generate: (input: ImageGenerateInput) => Promise /** The generation result containing images, or null */ @@ -58,10 +70,6 @@ export interface UseGenerateImageReturn { error: Accessor /** Current state of the generation */ status: Accessor - /** Abort the current generation */ - stop: () => void - /** Clear result, error, and return to idle */ - reset: () => void } /** @@ -109,19 +117,15 @@ export function useGenerateImage( hookName: 'useGenerateImage', outputKind: 'image' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - useGeneration({ - ...options, - devtools, - }) + const generation = useGeneration< + ImageGenerateInput, + ImageGenerationResult, + TTransformed + >({ + ...options, + devtools, + reconstructResult: reconstructImageResult, + }) - return { - generate: generate as (input: ImageGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + return generation } diff --git a/packages/ai-solid/src/use-generate-speech.ts b/packages/ai-solid/src/use-generate-speech.ts index b682e71a4..0daf9e1b1 100644 --- a/packages/ai-solid/src/use-generate-speech.ts +++ b/packages/ai-solid/src/use-generate-speech.ts @@ -1,4 +1,8 @@ import { useGeneration } from './use-generation' +import type { + UseGenerationOptions, + UseGenerationReturn, +} from './use-generation' import type { StreamChunk, TTSResult } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -15,7 +19,10 @@ import type { Accessor } from 'solid-js' * * @template TOutput - The transformed output type (defaults to TTSResult) */ -export interface UseGenerateSpeechOptions { +export interface UseGenerateSpeechOptions extends Pick< + UseGenerationOptions, + 'persistence' | 'threadId' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for speech generation */ @@ -47,7 +54,10 @@ export interface UseGenerateSpeechOptions { * * @template TOutput - The transformed output type (defaults to TTSResult) */ -export interface UseGenerateSpeechReturn { +export interface UseGenerateSpeechReturn extends Omit< + UseGenerationReturn, + 'generate' +> { /** Trigger speech generation */ generate: (input: SpeechGenerateInput) => Promise /** The TTS result containing audio data, or null */ @@ -58,10 +68,6 @@ export interface UseGenerateSpeechReturn { error: Accessor /** Current state of the generation */ status: Accessor - /** Abort the current generation */ - stop: () => void - /** Clear result, error, and return to idle */ - reset: () => void } /** @@ -103,19 +109,14 @@ export function useGenerateSpeech( hookName: 'useGenerateSpeech', outputKind: 'audio' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - useGeneration({ - ...options, - devtools, - }) + const generation = useGeneration< + SpeechGenerateInput, + TTSResult, + TTransformed + >({ + ...options, + devtools, + }) - return { - generate: generate as (input: SpeechGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + return generation } diff --git a/packages/ai-solid/src/use-generate-video.ts b/packages/ai-solid/src/use-generate-video.ts index 4fd68b5de..6fd991046 100644 --- a/packages/ai-solid/src/use-generate-video.ts +++ b/packages/ai-solid/src/use-generate-video.ts @@ -2,11 +2,11 @@ import { VideoGenerationClient } from '@tanstack/ai-client' import { createVideoDevtoolsBridge } from '@tanstack/ai-client/devtools' import { createEffect, - createMemo, createSignal, createUniqueId, onCleanup, onMount, + untrack, } from 'solid-js' import type { StreamChunk } from '@tanstack/ai' import type { @@ -14,6 +14,9 @@ import type { ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, VideoGenerateInput, VideoGenerateResult, @@ -37,6 +40,21 @@ export interface UseGenerateVideoOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** + * How this generation persists across reloads. + * - Omit / `false`: ephemeral, in-memory only. + * - `true`: server-driven — on mount the client hydrates the last generation + * for its `threadId` from the server (needs a connection with a + * `hydrateGeneration` handler) and repaints it; it never auto-starts a run. + * - a storage adapter: client-driven — the lightweight snapshot is cached under + * `generation:` as a run streams and read back on mount. Media bytes are + * never stored. + */ + persistence?: boolean | GenerationPersistence + /** Thread id for this generation, stable across reloads. Server-driven mode (`persistence: true`) hydrates the last generation under this key. Falls back to `id`. */ + threadId?: string + /** Explicit resume-snapshot seed for apps that manage storage themselves; skips automatic hydration from `persistence`. Later run events merge into it. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when video generation completes. Can optionally return a transformed value. * @@ -81,6 +99,8 @@ export interface UseGenerateVideoReturn { stop: () => void /** Clear all state and return to idle */ reset: () => void + /** Identity of the in-flight run while one is streaming, or null after it ends */ + resumeState: Accessor } /** @@ -139,13 +159,30 @@ export function useGenerateVideo( const [isLoading, setIsLoading] = createSignal(false) const [error, setError] = createSignal(undefined) const [status, setStatus] = createSignal('idle') + const [resumeState, setResumeState] = + createSignal( + options.initialResumeSnapshot?.resumeState ?? null, + ) + let disposed = false - const client = createMemo(() => { + // Built once. `untrack` keeps the option reads below from subscribing + // construction to `options.persistence` / `options.initialResumeSnapshot` / + // `options.devtools` / `options.body`: a re-run would build a second client + // and orphan the first (only the live one is disposed on cleanup). Later + // `options.body` changes are pushed through `updateOptions` instead. + const client = untrack((): VideoGenerationClient => { // Conditional spread on `body`: VideoGenerationClientOptions.body // is a strict optional; EOPT forbids passing `T | undefined`. const baseOptions = { id: clientId, body: options.body, + ...(options.persistence !== undefined && { + persistence: options.persistence, + }), + ...(options.threadId !== undefined && { threadId: options.threadId }), + ...(options.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: options.initialResumeSnapshot, + }), devtoolsBridgeFactory: createVideoDevtoolsBridge, devtools: { ...options.devtools, @@ -159,17 +196,42 @@ export function useGenerateVideo( onResult: ((r: VideoGenerateResult) => options.onResult?.(r)) as ( result: VideoGenerateResult, ) => TOutput | null | void, - onError: (e: Error) => options.onError?.(e), - onProgress: (p: number, m?: string) => options.onProgress?.(p, m), - onChunk: (c: StreamChunk) => options.onChunk?.(c), - onJobCreated: (id: string) => options.onJobCreated?.(id), - onStatusUpdate: (s: VideoStatusInfo) => options.onStatusUpdate?.(s), - onResultChange: setResult, - onLoadingChange: setIsLoading, - onErrorChange: setError, - onStatusChange: setStatus, - onJobIdChange: setJobId, - onVideoStatusChange: setVideoStatus, + onError: (e: Error) => { + if (!disposed) options.onError?.(e) + }, + onProgress: (p: number, m?: string) => { + if (!disposed) options.onProgress?.(p, m) + }, + onChunk: (c: StreamChunk) => { + if (!disposed) options.onChunk?.(c) + }, + onJobCreated: (id: string) => { + if (!disposed) options.onJobCreated?.(id) + }, + onStatusUpdate: (s: VideoStatusInfo) => { + if (!disposed) options.onStatusUpdate?.(s) + }, + onResultChange: (r: TOutput | null) => { + if (!disposed) setResult(() => r) + }, + onLoadingChange: (l: boolean) => { + if (!disposed) setIsLoading(l) + }, + onErrorChange: (e: Error | undefined) => { + if (!disposed) setError(e) + }, + onStatusChange: (s: GenerationClientState) => { + if (!disposed) setStatus(s) + }, + onJobIdChange: (id: string | null) => { + if (!disposed) setJobId(id) + }, + onVideoStatusChange: (s: VideoStatusInfo | null) => { + if (!disposed) setVideoStatus(s) + }, + onResumeStateChange: (rs: GenerationResumeState | null) => { + if (!disposed) setResumeState(rs) + }, } if (options.connection) { @@ -189,35 +251,38 @@ export function useGenerateVideo( throw new Error( 'useGenerateVideo requires either a connection or fetcher option', ) - }, [clientId]) + }) // Sync body changes without recreating client createEffect(() => { const currentBody = options.body - client().updateOptions({ + client.updateOptions({ ...(currentBody !== undefined && { body: currentBody }), }) }) + // Mount devtools only. Generation runs are never auto-started on mount — a + // persisted snapshot is hydrated for display, never replayed. onMount(() => { - client().mountDevtools() + client.mountDevtools() }) // Cleanup on unmount: stop any in-flight requests and unregister devtools onCleanup(() => { - client().dispose() + disposed = true + client.dispose() }) const generate = async (input: VideoGenerateInput) => { - await client().generate(input) + await client.generate(input) } const stop = () => { - client().stop() + client.stop() } const reset = () => { - client().reset() + client.reset() } return { @@ -230,5 +295,6 @@ export function useGenerateVideo( status, stop, reset, + resumeState, } } diff --git a/packages/ai-solid/src/use-generation.ts b/packages/ai-solid/src/use-generation.ts index 2db94a579..1d37e1e94 100644 --- a/packages/ai-solid/src/use-generation.ts +++ b/packages/ai-solid/src/use-generation.ts @@ -2,11 +2,11 @@ import { GenerationClient } from '@tanstack/ai-client' import { createGenerationDevtoolsBridge } from '@tanstack/ai-client/devtools' import { createEffect, - createMemo, createSignal, createUniqueId, onCleanup, onMount, + untrack, } from 'solid-js' import type { StreamChunk } from '@tanstack/ai' import type { @@ -15,6 +15,10 @@ import type { GenerationClientOptions, GenerationClientState, GenerationFetcher, + GenerationPersistence, + GenerationRestoredResult, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, } from '@tanstack/ai-client' import type { Accessor } from 'solid-js' @@ -39,6 +43,21 @@ export interface UseGenerationOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** + * How this generation persists across reloads. + * - Omit / `false`: ephemeral, in-memory only. + * - `true`: server-driven — on mount the client hydrates the last generation + * for its `threadId` from the server (needs a connection with a + * `hydrateGeneration` handler) and repaints it; it never auto-starts a run. + * - a storage adapter: client-driven — the lightweight snapshot is cached under + * `generation:` as a run streams and read back on mount. Media bytes are + * never stored. + */ + persistence?: boolean | GenerationPersistence + /** Thread id for this generation, stable across reloads. Server-driven mode (`persistence: true`) hydrates the last generation under this key. Falls back to `id`. */ + threadId?: string + /** Explicit resume-snapshot seed for apps that manage storage themselves; skips automatic hydration from `persistence`. Later run events merge into it. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when a result is received. Can optionally return a transformed value. * @@ -53,16 +72,26 @@ export interface UseGenerationOptions { onProgress?: (progress: number, message?: string) => void /** Callback for each stream chunk (connect-based adapter mode only) */ onChunk?: (chunk: StreamChunk) => void + /** + * @internal Rebuild a typed result from a restored snapshot, injected by each + * specialized hook (image / audio / transcription / summarize). Forwarded to + * the client so a client-store / server-hydrate restore repaints `result`. + */ + reconstructResult?: (restored: GenerationRestoredResult) => TResult | null } /** * Return type for the useGeneration hook. * * @template TOutput - The output type (possibly transformed from the raw result) + * @template TInput - The input type accepted by `generate` (defaults to any object) */ -export interface UseGenerationReturn { +export interface UseGenerationReturn< + TOutput, + TInput extends Record = Record, +> { /** Trigger a generation request */ - generate: (input: Record) => Promise + generate: (input: TInput) => Promise /** The generation result, or null if not yet generated */ result: Accessor /** Whether a generation is currently in progress */ @@ -75,6 +104,8 @@ export interface UseGenerationReturn { stop: () => void /** Clear result, error, and return to idle */ reset: () => void + /** Identity of the in-flight run while one is streaming, or null after it ends */ + resumeState: Accessor } /** @@ -111,7 +142,10 @@ export function useGeneration< options: Omit, 'onResult'> & { onResult?: (result: TResult) => TTransformed }, -): UseGenerationReturn> { +): UseGenerationReturn< + InferGenerationOutputFromReturn, + TInput +> { type TOutput = InferGenerationOutputFromReturn const hookId = createUniqueId() const clientId = options.id || hookId @@ -120,14 +154,34 @@ export function useGeneration< const [isLoading, setIsLoading] = createSignal(false) const [error, setError] = createSignal(undefined) const [status, setStatus] = createSignal('idle') + const [resumeState, setResumeState] = + createSignal( + options.initialResumeSnapshot?.resumeState ?? null, + ) + let disposed = false - const client = createMemo(() => { + // Built once. `untrack` keeps the option reads below from subscribing + // construction to `options.persistence` / `options.initialResumeSnapshot` / + // `options.devtools` / `options.body`: a re-run would build a second client + // and orphan the first (only the live one is disposed on cleanup). Later + // `options.body` changes are pushed through `updateOptions` instead. + const client = untrack((): GenerationClient => { // Conditional spread on `body`: `GenerationClientOptions.body` is a // strict optional (`body?: Record`) and EOPT forbids // assigning the source `T | undefined` directly. const clientOptions: GenerationClientOptions = { id: clientId, body: options.body, + ...(options.persistence !== undefined && { + persistence: options.persistence, + }), + ...(options.threadId !== undefined && { threadId: options.threadId }), + ...(options.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: options.initialResumeSnapshot, + }), + ...(options.reconstructResult + ? { reconstructResult: options.reconstructResult } + : {}), devtoolsBridgeFactory: createGenerationDevtoolsBridge, devtools: { ...options.devtools, @@ -140,13 +194,30 @@ export function useGeneration< onResult: ((r: TResult) => options.onResult?.(r)) as ( result: TResult, ) => TOutput | null | void, - onError: (e: Error) => options.onError?.(e), - onProgress: (p: number, m?: string) => options.onProgress?.(p, m), - onChunk: (c: StreamChunk) => options.onChunk?.(c), - onResultChange: setResult, - onLoadingChange: setIsLoading, - onErrorChange: setError, - onStatusChange: setStatus, + onError: (e: Error) => { + if (!disposed) options.onError?.(e) + }, + onProgress: (p: number, m?: string) => { + if (!disposed) options.onProgress?.(p, m) + }, + onChunk: (c: StreamChunk) => { + if (!disposed) options.onChunk?.(c) + }, + onResultChange: (r) => { + if (!disposed) setResult(() => r) + }, + onLoadingChange: (l) => { + if (!disposed) setIsLoading(l) + }, + onErrorChange: (e) => { + if (!disposed) setError(e) + }, + onStatusChange: (s) => { + if (!disposed) setStatus(s) + }, + onResumeStateChange: (rs) => { + if (!disposed) setResumeState(rs) + }, } if (options.connection) { @@ -166,44 +237,48 @@ export function useGeneration< throw new Error( 'useGeneration requires either a connection or fetcher option', ) - }, [clientId]) + }) // Sync body changes without recreating client createEffect(() => { const currentBody = options.body - client().updateOptions({ + client.updateOptions({ ...(currentBody !== undefined && { body: currentBody }), }) }) + // Mount devtools only. Generation runs are never auto-started on mount — a + // persisted snapshot is hydrated for display, never replayed. onMount(() => { - client().mountDevtools() + client.mountDevtools() }) // Cleanup on unmount: stop any in-flight requests and unregister devtools onCleanup(() => { - client().dispose() + disposed = true + client.dispose() }) const generate = async (input: TInput) => { - await client().generate(input) + await client.generate(input) } const stop = () => { - client().stop() + client.stop() } const reset = () => { - client().reset() + client.reset() } return { - generate: generate as (input: Record) => Promise, + generate, result, isLoading, error, status, stop, reset, + resumeState, } } diff --git a/packages/ai-solid/src/use-summarize.ts b/packages/ai-solid/src/use-summarize.ts index b6a9eefc6..59380ea8b 100644 --- a/packages/ai-solid/src/use-summarize.ts +++ b/packages/ai-solid/src/use-summarize.ts @@ -1,4 +1,9 @@ import { useGeneration } from './use-generation' +import { reconstructSummarizeResult } from '@tanstack/ai-client' +import type { + UseGenerationOptions, + UseGenerationReturn, +} from './use-generation' import type { StreamChunk, SummarizationResult } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -15,7 +20,12 @@ import type { Accessor } from 'solid-js' * * @template TOutput - The transformed output type (defaults to SummarizationResult) */ -export interface UseSummarizeOptions { +export interface UseSummarizeOptions< + TOutput = SummarizationResult, +> extends Pick< + UseGenerationOptions, + 'persistence' | 'threadId' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for summarization */ @@ -47,7 +57,10 @@ export interface UseSummarizeOptions { * * @template TOutput - The transformed output type (defaults to SummarizationResult) */ -export interface UseSummarizeReturn { +export interface UseSummarizeReturn extends Omit< + UseGenerationReturn, + 'generate' +> { /** Trigger summarization */ generate: (input: SummarizeGenerateInput) => Promise /** The summarization result, or null */ @@ -58,10 +71,6 @@ export interface UseSummarizeReturn { error: Accessor /** Current state of the generation */ status: Accessor - /** Abort the current summarization */ - stop: () => void - /** Clear result, error, and return to idle */ - reset: () => void } /** @@ -106,19 +115,15 @@ export function useSummarize( hookName: 'useSummarize', outputKind: 'text' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - useGeneration({ - ...options, - devtools, - }) + const generation = useGeneration< + SummarizeGenerateInput, + SummarizationResult, + TTransformed + >({ + ...options, + devtools, + reconstructResult: reconstructSummarizeResult, + }) - return { - generate: generate as (input: SummarizeGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + return generation } diff --git a/packages/ai-solid/src/use-transcription.ts b/packages/ai-solid/src/use-transcription.ts index e3e19ee7d..00b139633 100644 --- a/packages/ai-solid/src/use-transcription.ts +++ b/packages/ai-solid/src/use-transcription.ts @@ -1,4 +1,9 @@ import { useGeneration } from './use-generation' +import { reconstructTranscriptionResult } from '@tanstack/ai-client' +import type { + UseGenerationOptions, + UseGenerationReturn, +} from './use-generation' import type { StreamChunk, TranscriptionResult } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -15,7 +20,16 @@ import type { Accessor } from 'solid-js' * * @template TOutput - The transformed output type (defaults to TranscriptionResult) */ -export interface UseTranscriptionOptions { +export interface UseTranscriptionOptions< + TOutput = TranscriptionResult, +> extends Pick< + UseGenerationOptions< + TranscriptionGenerateInput, + TranscriptionResult, + TOutput + >, + 'persistence' | 'threadId' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for transcription */ @@ -47,7 +61,9 @@ export interface UseTranscriptionOptions { * * @template TOutput - The transformed output type (defaults to TranscriptionResult) */ -export interface UseTranscriptionReturn { +export interface UseTranscriptionReturn< + TOutput = TranscriptionResult, +> extends Omit, 'generate'> { /** Trigger transcription */ generate: (input: TranscriptionGenerateInput) => Promise /** The transcription result, or null */ @@ -58,10 +74,6 @@ export interface UseTranscriptionReturn { error: Accessor /** Current state of the generation */ status: Accessor - /** Abort the current transcription */ - stop: () => void - /** Clear result, error, and return to idle */ - reset: () => void } /** @@ -112,20 +124,11 @@ export function useTranscription( hookName: 'useTranscription', outputKind: 'text' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - useGeneration< - TranscriptionGenerateInput, - TranscriptionResult, - TTransformed - >({ ...options, devtools }) + const generation = useGeneration< + TranscriptionGenerateInput, + TranscriptionResult, + TTransformed + >({ ...options, devtools, reconstructResult: reconstructTranscriptionResult }) - return { - generate: generate as (input: TranscriptionGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + return generation } diff --git a/packages/ai-solid/tests/use-generation.test.ts b/packages/ai-solid/tests/use-generation.test.ts index 562268aa0..7b95ad439 100644 --- a/packages/ai-solid/tests/use-generation.test.ts +++ b/packages/ai-solid/tests/use-generation.test.ts @@ -8,8 +8,19 @@ import { useTranscription } from '../src/use-transcription' import { useSummarize } from '../src/use-summarize' import { useGenerateVideo } from '../src/use-generate-video' import { createMockConnectionAdapter } from './test-utils' -import type { StreamChunk, TTSResult, TranscriptionResult } from '@tanstack/ai' +import type { + PersistedArtifactRef, + StreamChunk, + TTSResult, + TranscriptionResult, +} from '@tanstack/ai' import { EventType } from '@tanstack/ai' +import type { + ConnectConnectionAdapter, + GenerationResumeSnapshot, + GenerationPersistence, + RunAgentInputContext, +} from '@tanstack/ai-client' // Helper to create generation stream chunks function createGenerationChunks(result: unknown): Array { @@ -71,6 +82,60 @@ function createVideoChunks(jobId: string, url: string): Array { ] } +const videoResumeSnapshot: GenerationResumeSnapshot = { + resumeState: { + threadId: 'thread-resume', + runId: 'run-resume', + }, + status: 'running', +} + +function createReplayVideoChunks(): Array { + return [ + { + type: EventType.RUN_STARTED, + runId: 'run-resume', + threadId: 'thread-resume', + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: 'generation:result', + value: { + jobId: 'job-replay', + status: 'completed', + url: 'https://example.com/video.mp4', + }, + timestamp: Date.now(), + }, + { + type: EventType.RUN_FINISHED, + runId: 'run-resume', + threadId: 'thread-resume', + timestamp: Date.now(), + }, + ] +} + +function createRunContextCaptureAdapter(chunks: Array): { + adapter: ConnectConnectionAdapter + connect: ReturnType + runContexts: Array +} { + const runContexts: Array = [] + const connect = vi.fn() + const adapter: ConnectConnectionAdapter = { + async *connect(_messages, _data, _signal, runContext) { + connect(runContext) + runContexts.push(runContext) + for (const chunk of chunks) { + yield chunk + } + }, + } + return { adapter, connect, runContexts } +} + // Helper to create error stream chunks. // NOTE: The AG-UI spec for RUN_ERROR carries `message` directly on the event // (not nested under `error`). We emit BOTH shapes here because GenerationClient @@ -809,7 +874,7 @@ describe('useSummarize', () => { const mockResult = { id: 'sum-1', summary: 'A brief summary', - model: 'gpt-4', + model: 'gpt-5.5', usage: { promptTokens: 100, completionTokens: 20, totalTokens: 120 }, } const onResult = vi.fn() @@ -851,7 +916,7 @@ describe('useSummarize', () => { describe('connection mode', () => { it('should summarize text using connection', async () => { - const mockResult = { summary: 'A brief summary', model: 'gpt-4' } + const mockResult = { summary: 'A brief summary', model: 'gpt-5.5' } const chunks = createGenerationChunks(mockResult) const adapter = createMockConnectionAdapter({ chunks }) @@ -869,7 +934,7 @@ describe('useSummarize', () => { const mockResult = { id: 'sum-1', summary: 'A brief summary', - model: 'gpt-4', + model: 'gpt-5.5', usage: { promptTokens: 100, completionTokens: 20, totalTokens: 120 }, } @@ -1074,6 +1139,39 @@ describe('useGenerateVideo', () => { expect(result.isLoading()).toBe(false) expect(result.status()).toBe('idle') }) + + it('does not auto-fire a video generation on mount from a persisted running snapshot', async () => { + // Regression guard for the removed generation resume surface (video). + const { adapter, connect } = createRunContextCaptureAdapter( + createReplayVideoChunks(), + ) + const persistence: GenerationPersistence = { + getItem: vi.fn(() => videoResumeSnapshot), + setItem: vi.fn(), + removeItem: vi.fn(), + } + + const { result } = renderHook(() => + useGenerateVideo({ + id: 'video-no-auto-fire', + connection: adapter, + persistence: persistence, + initialResumeSnapshot: videoResumeSnapshot, + }), + ) + + await Promise.resolve() + await new Promise((resolve) => setTimeout(resolve, 0)) + + expect(connect).not.toHaveBeenCalled() + // An explicit `initialResumeSnapshot` seed wins over storage, so the + // client skips hydration entirely here. + expect(persistence.getItem).not.toHaveBeenCalled() + expect(result.isLoading()).toBe(false) + expect(result.status()).toBe('idle') + // The seeded in-flight identity is exposed as read-only `resumeState`. + expect(result.resumeState()).toEqual(videoResumeSnapshot.resumeState) + }) }) describe('error handling', () => { @@ -1090,6 +1188,307 @@ describe('useGenerateVideo', () => { }) }) +describe('resume snapshot persistence', () => { + // Map-backed stand-in for a real storage adapter, so a test can both assert + // on the calls and read back what actually landed in storage. + function createMapPersistence( + seed: Array<[string, GenerationResumeSnapshot]> = [], + ): { + persistence: GenerationPersistence + store: Map + } { + const store = new Map(seed) + return { + store, + persistence: { + getItem: vi.fn((key: string) => store.get(key) ?? null), + setItem: vi.fn((key: string, value: GenerationResumeSnapshot) => { + store.set(key, value) + }), + removeItem: vi.fn((key: string) => { + store.delete(key) + }), + }, + } + } + + // Hydration and the persistence write queue are async; give both a turn. + const flush = async () => { + await Promise.resolve() + await new Promise((resolve) => setTimeout(resolve, 0)) + } + + it('does not auto-fire a generation on mount from a persisted running snapshot', async () => { + // Regression guard for the removed generation resume surface: mounting a + // generation hook that has server persistence and a persisted `running` + // snapshot must NOT start a fresh empty-prompt generation. + const { adapter, connect } = createRunContextCaptureAdapter( + createGenerationChunks({ id: '1' }), + ) + const persistence: GenerationPersistence = { + getItem: vi.fn(() => ({ + resumeState: { threadId: 'thread-resume', runId: 'run-resume' }, + status: 'running' as const, + })), + setItem: vi.fn(), + removeItem: vi.fn(), + } + + const { result } = renderHook(() => + useGeneration({ + id: 'no-auto-fire', + connection: adapter, + persistence, + initialResumeSnapshot: { + resumeState: { threadId: 'thread-resume', runId: 'run-resume' }, + status: 'running', + }, + }), + ) + + await flush() + + expect(connect).not.toHaveBeenCalled() + // The explicit initialResumeSnapshot seed takes precedence over storage, + // so automatic hydration (getItem) is skipped entirely. + expect(persistence.getItem).not.toHaveBeenCalled() + expect(result.isLoading()).toBe(false) + expect(result.status()).toBe('idle') + // The persisted snapshot is still exposed as read-only display state. + expect(result.resumeState()).toEqual({ + threadId: 'thread-resume', + runId: 'run-resume', + }) + }) + + it('repaints status from a stored complete snapshot on mount without starting a run', async () => { + const { adapter, connect } = createRunContextCaptureAdapter( + createGenerationChunks({ id: '1' }), + ) + const persistence: GenerationPersistence = { + getItem: vi.fn(() => ({ + resumeState: null, + status: 'complete' as const, + result: { id: 'result-1', model: 'image-model' }, + })), + setItem: vi.fn(), + removeItem: vi.fn(), + } + + const { result } = renderHook(() => + useGeneration({ + id: 'hydrated', + connection: adapter, + persistence, + }), + ) + + await flush() + + // A completed snapshot repaints the normal `status` field to `success`. + expect(result.status()).toBe('success') + expect(persistence.getItem).toHaveBeenCalledWith('generation:hydrated') + expect(connect).not.toHaveBeenCalled() + // The base hook injects no reconstructResult, so `result` stays null. + expect(result.result()).toBeNull() + expect(result.resumeState()).toBeNull() + }) + + it('repaints resumeState from a stored running snapshot on mount', async () => { + const { adapter, connect } = createRunContextCaptureAdapter( + createGenerationChunks({ id: '1' }), + ) + const persistence: GenerationPersistence = { + getItem: vi.fn(() => ({ + resumeState: { threadId: 'thread-resume', runId: 'run-resume' }, + status: 'running' as const, + })), + setItem: vi.fn(), + removeItem: vi.fn(), + } + + const { result } = renderHook(() => + useGeneration({ + id: 'running-hydrate', + connection: adapter, + persistence, + }), + ) + + await flush() + + expect(result.resumeState()).toEqual({ + threadId: 'thread-resume', + runId: 'run-resume', + }) + // A restored running run presents as `generating` but never auto-tails. + expect(result.status()).toBe('generating') + expect(result.isLoading()).toBe(false) + expect(connect).not.toHaveBeenCalled() + }) + + it('server-driven (persistence: true) hydrates from the connection by threadId without a local store', async () => { + const { adapter, connect } = createRunContextCaptureAdapter( + createGenerationChunks({ id: '1' }), + ) + const hydrateGeneration = vi.fn(async () => ({ + resumeSnapshot: { + schemaVersion: 1 as const, + resumeState: null, + status: 'complete' as const, + result: { id: 'server-result', model: 'image-model' }, + }, + activeRun: null, + })) + + const { result } = renderHook(() => + useGeneration({ + threadId: 'thread-server', + connection: { ...adapter, hydrateGeneration }, + persistence: true, + }), + ) + + await flush() + + // The server snapshot repaints `status` to `success`. + expect(result.status()).toBe('success') + expect(hydrateGeneration).toHaveBeenCalledWith('thread-server') + expect(connect).not.toHaveBeenCalled() + }) + + it('clears resumeState once a streamed run finishes', async () => { + const adapter = createMockConnectionAdapter({ + chunks: createReplayVideoChunks(), + }) + + const { result } = renderHook(() => useGeneration({ connection: adapter })) + + await result.generate({ prompt: 'replay' }) + + expect(result.status()).toBe('success') + // Once the run ends the in-flight identity is gone. + expect(result.resumeState()).toBeNull() + }) + + it('restores a completed image result from a durable artifact url', async () => { + // useGenerateImage injects `reconstructImageResult`, so a restored complete + // snapshot repaints `result` with the durable serve url — as if the run had + // just finished. + const artifact: PersistedArtifactRef = { + role: 'output', + artifactId: 'artifact-image-1', + threadId: 'thread-img', + runId: 'run-img', + name: 'image.png', + mimeType: 'image/png', + size: 2048, + createdAt: '2026-07-06T00:00:00.000Z', + url: '/api/artifacts/artifact-image-1', + source: { + activity: 'image', + path: 'runs/run-img/image.png', + provider: 'test', + model: 'test-image', + mediaType: 'image', + }, + } + const persistence: GenerationPersistence = { + getItem: vi.fn(() => ({ + resumeState: null, + status: 'complete' as const, + activity: 'image' as const, + result: { + id: 'img-restored', + model: 'test-image', + artifacts: [artifact], + }, + })), + setItem: vi.fn(), + removeItem: vi.fn(), + } + const adapter = createMockConnectionAdapter() + + const { result } = renderHook(() => + useGenerateImage({ + id: 'img-hydrate', + connection: adapter, + persistence, + }), + ) + + await flush() + + expect(result.status()).toBe('success') + expect(result.result()).toEqual({ + id: 'img-restored', + model: 'test-image', + images: [{ url: '/api/artifacts/artifact-image-1' }], + artifacts: [artifact], + }) + expect(result.resumeState()).toBeNull() + }) + + it('repaints resumeState for useGenerateVideo from a stored running snapshot', async () => { + const stored: GenerationResumeSnapshot = { + schemaVersion: 1, + resumeState: { threadId: 'thread-stored', runId: 'run-stored' }, + status: 'running', + } + const { persistence } = createMapPersistence([ + ['generation:video-hydrate-me', stored], + ]) + const adapter = createMockConnectionAdapter() + + const { result } = renderHook(() => + useGenerateVideo({ + id: 'video-hydrate-me', + connection: adapter, + persistence, + }), + ) + + await flush() + + expect(persistence.getItem).toHaveBeenCalledWith( + 'generation:video-hydrate-me', + ) + expect(result.resumeState()).toEqual(stored.resumeState) + // A restored running run presents as `generating` but never auto-tails. + expect(result.status()).toBe('generating') + expect(result.isLoading()).toBe(false) + }) + + it('clears resumeState and removes the persisted record on reset', async () => { + const { persistence, store } = createMapPersistence() + + const { result } = renderHook(() => + useGeneration({ + id: 'reset-me', + fetcher: async () => ({ id: '1' }), + persistence, + }), + ) + + await result.generate({ prompt: 'test' }) + await flush() + + expect(persistence.setItem).toHaveBeenCalledWith( + 'generation:reset-me', + expect.objectContaining({ status: 'complete' }), + ) + + result.reset() + + expect(result.resumeState()).toBeNull() + + await flush() + + expect(persistence.removeItem).toHaveBeenCalledWith('generation:reset-me') + expect(store.has('generation:reset-me')).toBe(false) + }) +}) + describe('onResult transform', () => { it('should transform result when onResult returns a value (fetcher)', async () => { // Inference (issue #848): `onResult`'s parameter is contextually typed from diff --git a/packages/ai-svelte/src/create-generate-audio.svelte.ts b/packages/ai-svelte/src/create-generate-audio.svelte.ts index 5838ee6f2..dd9da4c18 100644 --- a/packages/ai-svelte/src/create-generate-audio.svelte.ts +++ b/packages/ai-svelte/src/create-generate-audio.svelte.ts @@ -1,4 +1,9 @@ import { createGeneration } from './create-generation.svelte' +import { reconstructAudioResult } from '@tanstack/ai-client' +import type { + CreateGenerationOptions, + CreateGenerationReturn, +} from './create-generation.svelte' import type { AudioGenerationResult, StreamChunk } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -14,7 +19,12 @@ import type { * * @template TOutput - The output type after optional transform (defaults to AudioGenerationResult) */ -export interface CreateGenerateAudioOptions { +export interface CreateGenerateAudioOptions< + TOutput = AudioGenerationResult, +> extends Pick< + CreateGenerationOptions, + 'persistence' | 'threadId' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for audio generation */ @@ -46,7 +56,9 @@ export interface CreateGenerateAudioOptions { * * @template TOutput - The output type (after optional transform) */ -export interface CreateGenerateAudioReturn { +export interface CreateGenerateAudioReturn< + TOutput = AudioGenerationResult, +> extends Omit, 'generate'> { /** The generation result containing audio, or null */ readonly result: TOutput | null /** Whether generation is in progress */ @@ -57,12 +69,6 @@ export interface CreateGenerateAudioReturn { readonly status: GenerationClientState /** Trigger audio generation */ generate: (input: AudioGenerateInput) => Promise - /** Abort the current generation */ - stop: () => void - /** Clear result, error, and return to idle */ - reset: () => void - /** Update additional body parameters */ - updateBody: (body: Record) => void } /** @@ -111,6 +117,7 @@ export function createGenerateAudio( >({ ...options, devtools, + reconstructResult: reconstructAudioResult, }) return { @@ -126,9 +133,13 @@ export function createGenerateAudio( get status() { return gen.status }, - generate: gen.generate as (input: AudioGenerateInput) => Promise, + generate: gen.generate, stop: gen.stop, reset: gen.reset, updateBody: gen.updateBody, + dispose: gen.dispose, + get resumeState() { + return gen.resumeState + }, } } diff --git a/packages/ai-svelte/src/create-generate-image.svelte.ts b/packages/ai-svelte/src/create-generate-image.svelte.ts index 0909ddca2..684e5c366 100644 --- a/packages/ai-svelte/src/create-generate-image.svelte.ts +++ b/packages/ai-svelte/src/create-generate-image.svelte.ts @@ -1,4 +1,9 @@ import { createGeneration } from './create-generation.svelte' +import { reconstructImageResult } from '@tanstack/ai-client' +import type { + CreateGenerationOptions, + CreateGenerationReturn, +} from './create-generation.svelte' import type { ImageGenerationResult, StreamChunk } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -14,7 +19,12 @@ import type { * * @template TOutput - The output type after optional transform (defaults to ImageGenerationResult) */ -export interface CreateGenerateImageOptions { +export interface CreateGenerateImageOptions< + TOutput = ImageGenerationResult, +> extends Pick< + CreateGenerationOptions, + 'persistence' | 'threadId' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for image generation */ @@ -46,7 +56,9 @@ export interface CreateGenerateImageOptions { * * @template TOutput - The output type (after optional transform) */ -export interface CreateGenerateImageReturn { +export interface CreateGenerateImageReturn< + TOutput = ImageGenerationResult, +> extends Omit, 'generate'> { /** The generation result containing images, or null */ readonly result: TOutput | null /** Whether generation is in progress */ @@ -57,12 +69,6 @@ export interface CreateGenerateImageReturn { readonly status: GenerationClientState /** Trigger image generation */ generate: (input: ImageGenerateInput) => Promise - /** Abort the current generation */ - stop: () => void - /** Clear result, error, and return to idle */ - reset: () => void - /** Update additional body parameters */ - updateBody: (body: Record) => void } /** @@ -120,6 +126,7 @@ export function createGenerateImage( >({ ...options, devtools, + reconstructResult: reconstructImageResult, }) return { @@ -135,9 +142,13 @@ export function createGenerateImage( get status() { return gen.status }, - generate: gen.generate as (input: ImageGenerateInput) => Promise, + generate: gen.generate, stop: gen.stop, reset: gen.reset, updateBody: gen.updateBody, + dispose: gen.dispose, + get resumeState() { + return gen.resumeState + }, } } diff --git a/packages/ai-svelte/src/create-generate-speech.svelte.ts b/packages/ai-svelte/src/create-generate-speech.svelte.ts index 20ccafbf9..e41300655 100644 --- a/packages/ai-svelte/src/create-generate-speech.svelte.ts +++ b/packages/ai-svelte/src/create-generate-speech.svelte.ts @@ -1,4 +1,8 @@ import { createGeneration } from './create-generation.svelte' +import type { + CreateGenerationOptions, + CreateGenerationReturn, +} from './create-generation.svelte' import type { StreamChunk, TTSResult } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -14,7 +18,10 @@ import type { * * @template TOutput - The output type after optional transform (defaults to TTSResult) */ -export interface CreateGenerateSpeechOptions { +export interface CreateGenerateSpeechOptions extends Pick< + CreateGenerationOptions, + 'persistence' | 'threadId' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for speech generation */ @@ -46,7 +53,10 @@ export interface CreateGenerateSpeechOptions { * * @template TOutput - The output type (after optional transform) */ -export interface CreateGenerateSpeechReturn { +export interface CreateGenerateSpeechReturn extends Omit< + CreateGenerationReturn, + 'generate' +> { /** The TTS result containing audio data, or null */ readonly result: TOutput | null /** Whether generation is in progress */ @@ -57,12 +67,6 @@ export interface CreateGenerateSpeechReturn { readonly status: GenerationClientState /** Trigger speech generation */ generate: (input: SpeechGenerateInput) => Promise - /** Abort the current generation */ - stop: () => void - /** Clear result, error, and return to idle */ - reset: () => void - /** Update additional body parameters */ - updateBody: (body: Record) => void } /** @@ -122,9 +126,13 @@ export function createGenerateSpeech( get status() { return gen.status }, - generate: gen.generate as (input: SpeechGenerateInput) => Promise, + generate: gen.generate, stop: gen.stop, reset: gen.reset, updateBody: gen.updateBody, + dispose: gen.dispose, + get resumeState() { + return gen.resumeState + }, } } diff --git a/packages/ai-svelte/src/create-generate-video.svelte.ts b/packages/ai-svelte/src/create-generate-video.svelte.ts index 05f2d787b..d39d8207c 100644 --- a/packages/ai-svelte/src/create-generate-video.svelte.ts +++ b/packages/ai-svelte/src/create-generate-video.svelte.ts @@ -6,6 +6,9 @@ import type { ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, VideoGenerateInput, VideoGenerateResult, @@ -28,6 +31,21 @@ export interface CreateGenerateVideoOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** + * How this generation persists across reloads. + * - Omit / `false`: ephemeral, in-memory only. + * - `true`: server-driven — on mount the client hydrates the last generation + * for its `threadId` from the server (needs a connection with a + * `hydrateGeneration` handler) and repaints it; it never auto-starts a run. + * - a storage adapter: client-driven — the lightweight snapshot is cached under + * `generation:` as a run streams and read back on mount. Media bytes are + * never stored. + */ + persistence?: boolean | GenerationPersistence + /** Thread id for this generation, stable across reloads. Server-driven mode (`persistence: true`) hydrates the last generation under this key. Falls back to `id`. */ + threadId?: string + /** Explicit resume-snapshot seed for apps that manage storage themselves; skips automatic hydration from `persistence`. Later run events merge into it. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when video generation completes. Can optionally return a transformed value. * @@ -76,6 +94,8 @@ export interface CreateGenerateVideoReturn { dispose: () => void /** Update additional body parameters */ updateBody: (body: Record) => void + /** Identity of the in-flight run while one is streaming, or null after it ends */ + readonly resumeState: GenerationResumeState | null } /** @@ -133,6 +153,10 @@ export function createGenerateVideo( let isLoading = $state(false) let error = $state(undefined) let status = $state('idle') + let resumeState = $state( + options.initialResumeSnapshot?.resumeState ?? null, + ) + let disposed = false // `body` uses a conditional spread because `VideoGenerationClientOptions.body` // is declared `body?: Record` (absent vs. present) under @@ -141,6 +165,13 @@ export function createGenerateVideo( const baseOptions = { id: clientId, body: options.body, + ...(options.persistence !== undefined && { + persistence: options.persistence, + }), + ...(options.threadId !== undefined && { threadId: options.threadId }), + ...(options.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: options.initialResumeSnapshot, + }), devtoolsBridgeFactory: createVideoDevtoolsBridge, devtools: { ...options.devtools, @@ -154,29 +185,49 @@ export function createGenerateVideo( onResult: ((r: VideoGenerateResult) => options.onResult?.(r)) as ( result: VideoGenerateResult, ) => TOutput | null | void, - onError: (e: Error) => options.onError?.(e), - onProgress: (p: number, m?: string) => options.onProgress?.(p, m), - onChunk: (c: StreamChunk) => options.onChunk?.(c), - onJobCreated: (id: string) => options.onJobCreated?.(id), - onStatusUpdate: (s: VideoStatusInfo) => options.onStatusUpdate?.(s), + onError: (e: Error) => { + if (!disposed) options.onError?.(e) + }, + onProgress: (p: number, m?: string) => { + if (!disposed) options.onProgress?.(p, m) + }, + onChunk: (c: StreamChunk) => { + if (!disposed) options.onChunk?.(c) + }, + onJobCreated: (id: string) => { + if (!disposed) options.onJobCreated?.(id) + }, + onStatusUpdate: (s: VideoStatusInfo) => { + if (!disposed) options.onStatusUpdate?.(s) + }, onResultChange: (r: TOutput | null) => { + if (disposed) return result = r }, onLoadingChange: (l: boolean) => { + if (disposed) return isLoading = l }, onErrorChange: (e: Error | undefined) => { + if (disposed) return error = e }, onStatusChange: (s: GenerationClientState) => { + if (disposed) return status = s }, onJobIdChange: (id: string | null) => { + if (disposed) return jobId = id }, onVideoStatusChange: (s: VideoStatusInfo | null) => { + if (disposed) return videoStatus = s }, + onResumeStateChange: (rs: GenerationResumeState | null) => { + if (disposed) return + resumeState = rs + }, } let client: VideoGenerationClient @@ -197,6 +248,8 @@ export function createGenerateVideo( ) } + // Mount devtools only. Generation runs are never auto-started on setup — + // persisted state is read-only for display. client.mountDevtools() // Note: Cleanup is handled by calling dispose() directly when needed. @@ -205,6 +258,12 @@ export function createGenerateVideo( // Users should call video.dispose() in their component's cleanup if needed. const generate = async (input: VideoGenerateInput) => { + // Svelte has no remount effect to revive a disposed client (the other + // frameworks revive via mountDevtools() in their mount effects), so an + // explicit generate() after dispose() is the Svelte revive path: bring + // the client and the reactive bindings back together. + disposed = false + client.mountDevtools() await client.generate(input) } @@ -217,6 +276,7 @@ export function createGenerateVideo( } const dispose = () => { + disposed = true client.dispose() } @@ -248,5 +308,8 @@ export function createGenerateVideo( reset, dispose, updateBody, + get resumeState() { + return resumeState + }, } } diff --git a/packages/ai-svelte/src/create-generation.svelte.ts b/packages/ai-svelte/src/create-generation.svelte.ts index a731ac398..0e61eccc4 100644 --- a/packages/ai-svelte/src/create-generation.svelte.ts +++ b/packages/ai-svelte/src/create-generation.svelte.ts @@ -7,6 +7,10 @@ import type { GenerationClientOptions, GenerationClientState, GenerationFetcher, + GenerationPersistence, + GenerationRestoredResult, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, } from '@tanstack/ai-client' @@ -30,6 +34,21 @@ export interface CreateGenerationOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** + * How this generation persists across reloads. + * - Omit / `false`: ephemeral, in-memory only. + * - `true`: server-driven — on mount the client hydrates the last generation + * for its `threadId` from the server (needs a connection with a + * `hydrateGeneration` handler) and repaints it; it never auto-starts a run. + * - a storage adapter: client-driven — the lightweight snapshot is cached under + * `generation:` as a run streams and read back on mount. Media bytes are + * never stored. + */ + persistence?: boolean | GenerationPersistence + /** Thread id for this generation, stable across reloads. Server-driven mode (`persistence: true`) hydrates the last generation under this key. Falls back to `id`. */ + threadId?: string + /** Explicit resume-snapshot seed for apps that manage storage themselves; skips automatic hydration from `persistence`. Later run events merge into it. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when a result is received. Can optionally return a transformed value. * @@ -44,14 +63,24 @@ export interface CreateGenerationOptions { onProgress?: (progress: number, message?: string) => void /** Callback for each stream chunk (connect-based adapter mode only) */ onChunk?: (chunk: StreamChunk) => void + /** + * @internal Rebuild a typed result from a restored snapshot, injected by each + * specialized function (image / audio / transcription / summarize). Forwarded + * to the client so a client-store / server-hydrate restore repaints `result`. + */ + reconstructResult?: (restored: GenerationRestoredResult) => TResult | null } /** * Return type for the createGeneration function. * * @template TOutput - The output type (after optional transform) + * @template TInput - The input type accepted by `generate` (defaults to any object) */ -export interface CreateGenerationReturn { +export interface CreateGenerationReturn< + TOutput, + TInput extends Record = Record, +> { /** The generation result, or null if not yet generated */ readonly result: TOutput | null /** Whether a generation is currently in progress */ @@ -61,7 +90,7 @@ export interface CreateGenerationReturn { /** Current state of the generation client */ readonly status: GenerationClientState /** Trigger a generation request */ - generate: (input: Record) => Promise + generate: (input: TInput) => Promise /** Abort the current generation */ stop: () => void /** Clear result, error, and return to idle */ @@ -70,6 +99,8 @@ export interface CreateGenerationReturn { dispose: () => void /** Update additional body parameters */ updateBody: (body: Record) => void + /** Identity of the in-flight run while one is streaming, or null after it ends */ + readonly resumeState: GenerationResumeState | null } /** @@ -118,7 +149,8 @@ export function createGeneration< onResult?: (result: TResult) => TTransformed }, ): CreateGenerationReturn< - InferGenerationOutputFromReturn + InferGenerationOutputFromReturn, + TInput > { type TOutput = InferGenerationOutputFromReturn const clientId = @@ -129,6 +161,10 @@ export function createGeneration< let isLoading = $state(false) let error = $state(undefined) let status = $state('idle') + let resumeState = $state( + options.initialResumeSnapshot?.resumeState ?? null, + ) + let disposed = false // `body` uses a conditional spread because `GenerationClientOptions.body` // is declared `body?: Record` (absent vs. present) under @@ -138,6 +174,16 @@ export function createGeneration< const clientOptions: GenerationClientOptions = { id: clientId, body: options.body, + ...(options.persistence !== undefined && { + persistence: options.persistence, + }), + ...(options.threadId !== undefined && { threadId: options.threadId }), + ...(options.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: options.initialResumeSnapshot, + }), + ...(options.reconstructResult + ? { reconstructResult: options.reconstructResult } + : {}), devtoolsBridgeFactory: createGenerationDevtoolsBridge, devtools: { ...options.devtools, @@ -150,21 +196,35 @@ export function createGeneration< onResult: ((r: TResult) => options.onResult?.(r)) as ( result: TResult, ) => TOutput | null | void, - onError: (e: Error) => options.onError?.(e), - onProgress: (p: number, m?: string) => options.onProgress?.(p, m), - onChunk: (c: StreamChunk) => options.onChunk?.(c), + onError: (e: Error) => { + if (!disposed) options.onError?.(e) + }, + onProgress: (p: number, m?: string) => { + if (!disposed) options.onProgress?.(p, m) + }, + onChunk: (c: StreamChunk) => { + if (!disposed) options.onChunk?.(c) + }, onResultChange: (r: TOutput | null) => { + if (disposed) return result = r }, onLoadingChange: (l: boolean) => { + if (disposed) return isLoading = l }, onErrorChange: (e: Error | undefined) => { + if (disposed) return error = e }, onStatusChange: (s: GenerationClientState) => { + if (disposed) return status = s }, + onResumeStateChange: (rs: GenerationResumeState | null) => { + if (disposed) return + resumeState = rs + }, } let client: GenerationClient @@ -185,6 +245,8 @@ export function createGeneration< ) } + // Mount devtools only. Generation runs are never auto-started on setup — + // persisted state is read-only for display. client.mountDevtools() // Note: Cleanup is handled by calling dispose() directly when needed. @@ -193,6 +255,12 @@ export function createGeneration< // Users should call gen.dispose() in their component's cleanup if needed. const generate = async (input: TInput) => { + // Svelte has no remount effect to revive a disposed client (the other + // frameworks revive via mountDevtools() in their mount effects), so an + // explicit generate() after dispose() is the Svelte revive path: bring + // the client and the reactive bindings back together. + disposed = false + client.mountDevtools() await client.generate(input) } @@ -205,6 +273,7 @@ export function createGeneration< } const dispose = () => { + disposed = true client.dispose() } @@ -225,10 +294,13 @@ export function createGeneration< get status() { return status }, - generate: generate as (input: Record) => Promise, + generate, stop, reset, dispose, updateBody, + get resumeState() { + return resumeState + }, } } diff --git a/packages/ai-svelte/src/create-summarize.svelte.ts b/packages/ai-svelte/src/create-summarize.svelte.ts index 0009eee15..dde442282 100644 --- a/packages/ai-svelte/src/create-summarize.svelte.ts +++ b/packages/ai-svelte/src/create-summarize.svelte.ts @@ -1,4 +1,9 @@ import { createGeneration } from './create-generation.svelte' +import { reconstructSummarizeResult } from '@tanstack/ai-client' +import type { + CreateGenerationOptions, + CreateGenerationReturn, +} from './create-generation.svelte' import type { StreamChunk, SummarizationResult } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -14,7 +19,12 @@ import type { * * @template TOutput - The output type after optional transform (defaults to SummarizationResult) */ -export interface CreateSummarizeOptions { +export interface CreateSummarizeOptions< + TOutput = SummarizationResult, +> extends Pick< + CreateGenerationOptions, + 'persistence' | 'threadId' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for summarization */ @@ -46,7 +56,9 @@ export interface CreateSummarizeOptions { * * @template TOutput - The output type (after optional transform) */ -export interface CreateSummarizeReturn { +export interface CreateSummarizeReturn< + TOutput = SummarizationResult, +> extends Omit, 'generate'> { /** The summarization result, or null */ readonly result: TOutput | null /** Whether summarization is in progress */ @@ -57,12 +69,6 @@ export interface CreateSummarizeReturn { readonly status: GenerationClientState /** Trigger summarization */ generate: (input: SummarizeGenerateInput) => Promise - /** Abort the current summarization */ - stop: () => void - /** Clear result, error, and return to idle */ - reset: () => void - /** Update additional body parameters */ - updateBody: (body: Record) => void } /** @@ -115,6 +121,7 @@ export function createSummarize( >({ ...options, devtools, + reconstructResult: reconstructSummarizeResult, }) return { @@ -130,9 +137,13 @@ export function createSummarize( get status() { return gen.status }, - generate: gen.generate as (input: SummarizeGenerateInput) => Promise, + generate: gen.generate, stop: gen.stop, reset: gen.reset, updateBody: gen.updateBody, + dispose: gen.dispose, + get resumeState() { + return gen.resumeState + }, } } diff --git a/packages/ai-svelte/src/create-transcription.svelte.ts b/packages/ai-svelte/src/create-transcription.svelte.ts index b6583cb98..9403eb7a9 100644 --- a/packages/ai-svelte/src/create-transcription.svelte.ts +++ b/packages/ai-svelte/src/create-transcription.svelte.ts @@ -1,4 +1,9 @@ import { createGeneration } from './create-generation.svelte' +import { reconstructTranscriptionResult } from '@tanstack/ai-client' +import type { + CreateGenerationOptions, + CreateGenerationReturn, +} from './create-generation.svelte' import type { StreamChunk, TranscriptionResult } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -14,7 +19,16 @@ import type { * * @template TOutput - The output type after optional transform (defaults to TranscriptionResult) */ -export interface CreateTranscriptionOptions { +export interface CreateTranscriptionOptions< + TOutput = TranscriptionResult, +> extends Pick< + CreateGenerationOptions< + TranscriptionGenerateInput, + TranscriptionResult, + TOutput + >, + 'persistence' | 'threadId' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for transcription */ @@ -46,7 +60,9 @@ export interface CreateTranscriptionOptions { * * @template TOutput - The output type (after optional transform) */ -export interface CreateTranscriptionReturn { +export interface CreateTranscriptionReturn< + TOutput = TranscriptionResult, +> extends Omit, 'generate'> { /** The transcription result, or null */ readonly result: TOutput | null /** Whether transcription is in progress */ @@ -57,12 +73,6 @@ export interface CreateTranscriptionReturn { readonly status: GenerationClientState /** Trigger transcription */ generate: (input: TranscriptionGenerateInput) => Promise - /** Abort the current transcription */ - stop: () => void - /** Clear result, error, and return to idle */ - reset: () => void - /** Update additional body parameters */ - updateBody: (body: Record) => void } /** @@ -120,6 +130,7 @@ export function createTranscription( >({ ...options, devtools, + reconstructResult: reconstructTranscriptionResult, }) return { @@ -135,11 +146,13 @@ export function createTranscription( get status() { return gen.status }, - generate: gen.generate as ( - input: TranscriptionGenerateInput, - ) => Promise, + generate: gen.generate, stop: gen.stop, reset: gen.reset, updateBody: gen.updateBody, + dispose: gen.dispose, + get resumeState() { + return gen.resumeState + }, } } diff --git a/packages/ai-svelte/src/index.ts b/packages/ai-svelte/src/index.ts index efe617ebe..8ac12363d 100644 --- a/packages/ai-svelte/src/index.ts +++ b/packages/ai-svelte/src/index.ts @@ -98,4 +98,10 @@ export { type VideoGenerateInput, type VideoGenerateResult, type VideoStatusInfo, + type GenerationPersistence, + type GenerationResumeSnapshot, + type GenerationResumeState, + type GenerationResumeStatus, + type GenerationPendingArtifact, } from '@tanstack/ai-client' +export type { PersistedArtifactRef } from '@tanstack/ai/client' diff --git a/packages/ai-svelte/tests/create-generation.test.ts b/packages/ai-svelte/tests/create-generation.test.ts index ca2761c9f..7216daa74 100644 --- a/packages/ai-svelte/tests/create-generation.test.ts +++ b/packages/ai-svelte/tests/create-generation.test.ts @@ -7,7 +7,16 @@ import { createSummarize } from '../src/create-summarize.svelte' import { createGenerateVideo } from '../src/create-generate-video.svelte' import { createMockConnectionAdapter } from './test-utils' import { EventType, type StreamChunk } from '@tanstack/ai' -import type { TTSResult, TranscriptionResult } from '@tanstack/ai' +import type { + PersistedArtifactRef, + TTSResult, + TranscriptionResult, +} from '@tanstack/ai' +import type { + ConnectConnectionAdapter, + GenerationResumeSnapshot, + RunAgentInputContext, +} from '@tanstack/ai-client' // Helper to create generation stream chunks function createGenerationChunks(result: unknown): Array { @@ -71,6 +80,59 @@ function createVideoChunks(jobId: string, url: string): Array { ] } +const videoResumeSnapshot: GenerationResumeSnapshot = { + resumeState: { + threadId: 'thread-resume', + runId: 'run-resume', + }, + status: 'running', +} + +function createRunContextCaptureAdapter(chunks: Array): { + adapter: ConnectConnectionAdapter + connect: ReturnType + runContexts: Array +} { + const runContexts: Array = [] + const connect = vi.fn() + const adapter: ConnectConnectionAdapter = { + async *connect(_messages, _data, _signal, runContext) { + connect(runContext) + runContexts.push(runContext) + for (const chunk of chunks) { + yield chunk + } + }, + } + return { adapter, connect, runContexts } +} + +/** + * Storage adapter backed by a plain Map, seeded with raw (untyped) records the + * way real storage hands them back — the client re-validates whatever it reads. + */ +function createMapPersistence(seed: Record = {}) { + const store = new Map(Object.entries(seed)) + return { + store, + getItem: vi.fn(async (key: string) => { + return store.get(key) as GenerationResumeSnapshot | undefined + }), + setItem: vi.fn(async (key: string, value: GenerationResumeSnapshot) => { + store.set(key, value) + }), + removeItem: vi.fn(async (key: string) => { + store.delete(key) + }), + } +} + +// Snapshot hydration and removal both run through promise queues detached from +// the caller, so tests wait a macrotask for them to settle. +function flushAsync(): Promise { + return new Promise((resolve) => setTimeout(resolve, 0)) +} + // Helper to create error stream chunks function createErrorChunks(message: string): Array { return [ @@ -186,6 +248,120 @@ describe('createGeneration', () => { expect(gen.status).toBe('error') expect(gen.error?.message).toBe('Generation failed') }) + + it('does not auto-fire a generation on setup from a persisted running snapshot', async () => { + // Regression guard for the removed generation resume surface. + const snapshot: GenerationResumeSnapshot = { + resumeState: { threadId: 'thread-resume', runId: 'run-resume' }, + status: 'running', + } + const { adapter, connect } = createRunContextCaptureAdapter([]) + const getItem = vi.fn(() => snapshot) + const gen = createGeneration({ + id: 'no-auto-fire', + connection: adapter, + persistence: { getItem, setItem: vi.fn(), removeItem: vi.fn() }, + initialResumeSnapshot: snapshot, + }) + + await Promise.resolve() + + expect(connect).not.toHaveBeenCalled() + // An explicit `initialResumeSnapshot` seed takes precedence over + // storage, so the client skips its hydration read entirely. + expect(getItem).not.toHaveBeenCalled() + expect(gen.isLoading).toBe(false) + expect(gen.status).toBe('idle') + // The persisted snapshot remains exposed as read-only state. + expect(gen.resumeState).toEqual(snapshot.resumeState) + }) + }) + + describe('resume snapshot persistence', () => { + it('repaints resumeState from a stored running snapshot on creation', async () => { + const persistence = createMapPersistence({ + 'generation:hydrate-me': { + schemaVersion: 1, + resumeState: { threadId: 'thread-stored', runId: 'run-stored' }, + status: 'running', + }, + }) + + const gen = createGeneration({ + id: 'hydrate-me', + connection: createMockConnectionAdapter(), + persistence, + }) + + // Hydration is async — the storage read is awaited off the constructor. + expect(gen.resumeState).toBeNull() + await flushAsync() + + expect(persistence.getItem).toHaveBeenCalledTimes(1) + expect(persistence.getItem).toHaveBeenCalledWith('generation:hydrate-me') + // A restored running run presents as `generating` but never auto-tails. + expect(gen.resumeState).toEqual({ + threadId: 'thread-stored', + runId: 'run-stored', + }) + expect(gen.status).toBe('generating') + expect(gen.isLoading).toBe(false) + }) + + it('repaints status from a stored complete snapshot on creation without starting a run', async () => { + const persistence = createMapPersistence({ + 'generation:complete-me': { + schemaVersion: 1, + resumeState: null, + status: 'complete', + result: { id: 'result-1', model: 'image-model' }, + }, + }) + + const gen = createGeneration({ + id: 'complete-me', + connection: createMockConnectionAdapter(), + persistence, + }) + + await flushAsync() + + // A completed snapshot repaints the normal `status` field to `success`. + expect(gen.status).toBe('success') + // The base function injects no reconstructResult, so `result` stays null. + expect(gen.result).toBeNull() + expect(gen.resumeState).toBeNull() + }) + + it('clears resumeState and removes the persisted record on reset', async () => { + const persistence = createMapPersistence({ + 'generation:reset-me': { + schemaVersion: 1, + resumeState: { threadId: 'thread-stored', runId: 'run-stored' }, + status: 'running', + }, + }) + + const gen = createGeneration({ + id: 'reset-me', + connection: createMockConnectionAdapter(), + persistence, + }) + + await flushAsync() + expect(gen.resumeState).toEqual({ + threadId: 'thread-stored', + runId: 'run-stored', + }) + + gen.reset() + + expect(gen.resumeState).toBeNull() + + await flushAsync() + expect(persistence.removeItem).toHaveBeenCalledWith('generation:reset-me') + expect(persistence.store.has('generation:reset-me')).toBe(false) + }) }) describe('stop and reset', () => { @@ -303,6 +479,62 @@ describe('createGenerateImage', () => { expect(typeof gen.stop).toBe('function') expect(typeof gen.reset).toBe('function') }) + + it('restores a completed image result from a durable artifact url', async () => { + // createGenerateImage injects `reconstructImageResult`, so a restored + // complete snapshot repaints `result` with the durable serve url — as if the + // run had just finished. + const artifact: PersistedArtifactRef = { + role: 'output', + artifactId: 'artifact-image-1', + threadId: 'thread-img', + runId: 'run-img', + name: 'image.png', + mimeType: 'image/png', + size: 2048, + createdAt: '2026-07-06T00:00:00.000Z', + url: '/api/artifacts/artifact-image-1', + source: { + activity: 'image', + path: 'runs/run-img/image.png', + provider: 'test', + model: 'test-image', + mediaType: 'image', + }, + } + const persistence = createMapPersistence({ + 'generation:img-hydrate': { + schemaVersion: 1, + resumeState: null, + status: 'complete', + activity: 'image', + result: { + id: 'img-restored', + model: 'test-image', + artifacts: [artifact], + }, + }, + }) + + const gen = createGenerateImage({ + id: 'img-hydrate', + connection: createMockConnectionAdapter(), + persistence, + }) + + await flushAsync() + + // The completed snapshot repaints `status` and rebuilds `result` from the + // durable serve url. + expect(gen.status).toBe('success') + expect(gen.result).toEqual({ + id: 'img-restored', + model: 'test-image', + images: [{ url: '/api/artifacts/artifact-image-1' }], + artifacts: [artifact], + }) + expect(gen.resumeState).toBeNull() + }) }) describe('createGenerateSpeech', () => { @@ -489,7 +721,7 @@ describe('createSummarize', () => { const mockResult = { id: 'sum-1', summary: 'A brief summary', - model: 'gpt-4', + model: 'gpt-5.5', usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, } @@ -504,7 +736,7 @@ describe('createSummarize', () => { }) it('should summarize text using connection', async () => { - const mockResult = { summary: 'A brief summary', model: 'gpt-4' } + const mockResult = { summary: 'A brief summary', model: 'gpt-5.5' } const chunks = createGenerationChunks(mockResult) const adapter = createMockConnectionAdapter({ chunks }) @@ -538,7 +770,7 @@ describe('createSummarize', () => { fetcher: async () => ({ id: 'sum-1', summary: 'A brief summary', - model: 'gpt-4', + model: 'gpt-5.5', usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, }), }) @@ -658,6 +890,27 @@ describe('createGenerateVideo', () => { expect(gen.status).toBe('idle') }) + it('does not auto-fire a video generation on setup from a persisted running snapshot', async () => { + // Regression guard for the removed generation resume surface (video). + const { adapter, connect } = createRunContextCaptureAdapter([]) + const getItem = vi.fn(() => videoResumeSnapshot) + const gen = createGenerateVideo({ + id: 'video-no-auto-fire', + connection: adapter, + persistence: { getItem, setItem: vi.fn(), removeItem: vi.fn() }, + initialResumeSnapshot: videoResumeSnapshot, + }) + + await Promise.resolve() + + expect(connect).not.toHaveBeenCalled() + expect(getItem).not.toHaveBeenCalled() + expect(gen.isLoading).toBe(false) + expect(gen.status).toBe('idle') + // The seeded in-flight identity is exposed as read-only `resumeState`. + expect(gen.resumeState).toEqual(videoResumeSnapshot.resumeState) + }) + it('should expose generate, stop, reset, and updateBody methods', () => { const adapter = createMockConnectionAdapter() const gen = createGenerateVideo({ connection: adapter }) diff --git a/packages/ai-utils/src/base64.ts b/packages/ai-utils/src/base64.ts index 58e8e9a61..bb7827078 100644 --- a/packages/ai-utils/src/base64.ts +++ b/packages/ai-utils/src/base64.ts @@ -55,6 +55,13 @@ export function arrayBufferToBase64(buffer: ArrayBuffer): string { throw new Error('No base64 encoder available in this environment.') } +/** + * Decode a base64 string into a `Uint8Array`. + */ +export function base64ToUint8Array(base64: string): Uint8Array { + return new Uint8Array(base64ToArrayBuffer(base64)) +} + /** * Decode a base64 string into an `ArrayBuffer`. */ diff --git a/packages/ai-utils/src/index.ts b/packages/ai-utils/src/index.ts index 843d5eb37..619338eee 100644 --- a/packages/ai-utils/src/index.ts +++ b/packages/ai-utils/src/index.ts @@ -2,4 +2,8 @@ export { generateId } from './id' export { getApiKeyFromEnv } from './env' export { transformNullsToUndefined, undoNullWidening } from './transforms' export type { NullWideningMap } from './transforms' -export { arrayBufferToBase64, base64ToArrayBuffer } from './base64' +export { + arrayBufferToBase64, + base64ToArrayBuffer, + base64ToUint8Array, +} from './base64' diff --git a/packages/ai-vue/src/index.ts b/packages/ai-vue/src/index.ts index 86b39b403..ae43ec04c 100644 --- a/packages/ai-vue/src/index.ts +++ b/packages/ai-vue/src/index.ts @@ -94,4 +94,10 @@ export { type VideoGenerateInput, type VideoGenerateResult, type VideoStatusInfo, + type GenerationPersistence, + type GenerationResumeSnapshot, + type GenerationResumeState, + type GenerationResumeStatus, + type GenerationPendingArtifact, } from '@tanstack/ai-client' +export type { PersistedArtifactRef } from '@tanstack/ai/client' diff --git a/packages/ai-vue/src/use-generate-audio.ts b/packages/ai-vue/src/use-generate-audio.ts index 91c8eb6a0..3a325f5f7 100644 --- a/packages/ai-vue/src/use-generate-audio.ts +++ b/packages/ai-vue/src/use-generate-audio.ts @@ -1,4 +1,9 @@ import { useGeneration } from './use-generation' +import { reconstructAudioResult } from '@tanstack/ai-client' +import type { + UseGenerationOptions, + UseGenerationReturn, +} from './use-generation' import type { AudioGenerationResult, StreamChunk } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -15,7 +20,12 @@ import type { DeepReadonly, ShallowRef } from 'vue' * * @template TOutput - The output type after optional transform (defaults to AudioGenerationResult) */ -export interface UseGenerateAudioOptions { +export interface UseGenerateAudioOptions< + TOutput = AudioGenerationResult, +> extends Pick< + UseGenerationOptions, + 'persistence' | 'threadId' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for audio generation */ @@ -47,7 +57,9 @@ export interface UseGenerateAudioOptions { * * @template TOutput - The output type (after optional transform) */ -export interface UseGenerateAudioReturn { +export interface UseGenerateAudioReturn< + TOutput = AudioGenerationResult, +> extends Omit, 'generate'> { /** Trigger audio generation */ generate: (input: AudioGenerateInput) => Promise /** The generation result containing audio, or null */ @@ -58,10 +70,6 @@ export interface UseGenerateAudioReturn { error: DeepReadonly> /** Current state of the generation */ status: DeepReadonly> - /** Abort the current generation */ - stop: () => void - /** Clear result, error, and return to idle */ - reset: () => void } /** @@ -101,19 +109,15 @@ export function useGenerateAudio( hookName: 'useGenerateAudio', outputKind: 'audio' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - useGeneration({ - ...options, - devtools, - }) + const generation = useGeneration< + AudioGenerateInput, + AudioGenerationResult, + TTransformed + >({ + ...options, + devtools, + reconstructResult: reconstructAudioResult, + }) - return { - generate: generate as (input: AudioGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + return generation } diff --git a/packages/ai-vue/src/use-generate-image.ts b/packages/ai-vue/src/use-generate-image.ts index f80cc2f5c..339f0666f 100644 --- a/packages/ai-vue/src/use-generate-image.ts +++ b/packages/ai-vue/src/use-generate-image.ts @@ -1,4 +1,9 @@ import { useGeneration } from './use-generation' +import { reconstructImageResult } from '@tanstack/ai-client' +import type { + UseGenerationOptions, + UseGenerationReturn, +} from './use-generation' import type { ImageGenerationResult, StreamChunk } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -15,7 +20,12 @@ import type { DeepReadonly, ShallowRef } from 'vue' * * @template TOutput - The output type after optional transform (defaults to ImageGenerationResult) */ -export interface UseGenerateImageOptions { +export interface UseGenerateImageOptions< + TOutput = ImageGenerationResult, +> extends Pick< + UseGenerationOptions, + 'persistence' | 'threadId' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for image generation */ @@ -47,7 +57,9 @@ export interface UseGenerateImageOptions { * * @template TOutput - The output type (after optional transform) */ -export interface UseGenerateImageReturn { +export interface UseGenerateImageReturn< + TOutput = ImageGenerationResult, +> extends Omit, 'generate'> { /** Trigger image generation */ generate: (input: ImageGenerateInput) => Promise /** The generation result containing images, or null */ @@ -58,10 +70,6 @@ export interface UseGenerateImageReturn { error: DeepReadonly> /** Current state of the generation */ status: DeepReadonly> - /** Abort the current generation */ - stop: () => void - /** Clear result, error, and return to idle */ - reset: () => void } /** @@ -111,19 +119,15 @@ export function useGenerateImage( hookName: 'useGenerateImage', outputKind: 'image' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - useGeneration({ - ...options, - devtools, - }) + const generation = useGeneration< + ImageGenerateInput, + ImageGenerationResult, + TTransformed + >({ + ...options, + devtools, + reconstructResult: reconstructImageResult, + }) - return { - generate: generate as (input: ImageGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + return generation } diff --git a/packages/ai-vue/src/use-generate-speech.ts b/packages/ai-vue/src/use-generate-speech.ts index 2302fc766..33734afd1 100644 --- a/packages/ai-vue/src/use-generate-speech.ts +++ b/packages/ai-vue/src/use-generate-speech.ts @@ -1,4 +1,8 @@ import { useGeneration } from './use-generation' +import type { + UseGenerationOptions, + UseGenerationReturn, +} from './use-generation' import type { StreamChunk, TTSResult } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -15,7 +19,10 @@ import type { DeepReadonly, ShallowRef } from 'vue' * * @template TOutput - The output type after optional transform (defaults to TTSResult) */ -export interface UseGenerateSpeechOptions { +export interface UseGenerateSpeechOptions extends Pick< + UseGenerationOptions, + 'persistence' | 'threadId' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for speech generation */ @@ -47,7 +54,10 @@ export interface UseGenerateSpeechOptions { * * @template TOutput - The output type (after optional transform) */ -export interface UseGenerateSpeechReturn { +export interface UseGenerateSpeechReturn extends Omit< + UseGenerationReturn, + 'generate' +> { /** Trigger speech generation */ generate: (input: SpeechGenerateInput) => Promise /** The TTS result containing audio data, or null */ @@ -58,10 +68,6 @@ export interface UseGenerateSpeechReturn { error: DeepReadonly> /** Current state of the generation */ status: DeepReadonly> - /** Abort the current generation */ - stop: () => void - /** Clear result, error, and return to idle */ - reset: () => void } /** @@ -105,19 +111,14 @@ export function useGenerateSpeech( hookName: 'useGenerateSpeech', outputKind: 'audio' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - useGeneration({ - ...options, - devtools, - }) + const generation = useGeneration< + SpeechGenerateInput, + TTSResult, + TTransformed + >({ + ...options, + devtools, + }) - return { - generate: generate as (input: SpeechGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + return generation } diff --git a/packages/ai-vue/src/use-generate-video.ts b/packages/ai-vue/src/use-generate-video.ts index 644558508..67e4c51de 100644 --- a/packages/ai-vue/src/use-generate-video.ts +++ b/packages/ai-vue/src/use-generate-video.ts @@ -14,6 +14,9 @@ import type { ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, VideoGenerateInput, VideoGenerateResult, @@ -37,6 +40,21 @@ export interface UseGenerateVideoOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** + * How this generation persists across reloads. + * - Omit / `false`: ephemeral, in-memory only. + * - `true`: server-driven — on mount the client hydrates the last generation + * for its `threadId` from the server (needs a connection with a + * `hydrateGeneration` handler) and repaints it; it never auto-starts a run. + * - a storage adapter: client-driven — the lightweight snapshot is cached under + * `generation:` as a run streams and read back on mount. Media bytes are + * never stored. + */ + persistence?: boolean | GenerationPersistence + /** Thread id for this generation, stable across reloads. Server-driven mode (`persistence: true`) hydrates the last generation under this key. Falls back to `id`. */ + threadId?: string + /** Explicit resume-snapshot seed for apps that manage storage themselves; skips automatic hydration from `persistence`. Later run events merge into it. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when video generation completes. Can optionally return a transformed value. * @@ -81,6 +99,8 @@ export interface UseGenerateVideoReturn { stop: () => void /** Clear all state and return to idle */ reset: () => void + /** Identity of the in-flight run while one is streaming, or null after it ends */ + resumeState: DeepReadonly> } /** @@ -137,6 +157,10 @@ export function useGenerateVideo( const isLoading = shallowRef(false) const error = shallowRef(undefined) const status = shallowRef('idle') + const resumeState = shallowRef( + options.initialResumeSnapshot?.resumeState ?? null, + ) + let disposed = false // Conditional spread on `body`: `VideoGenerationClientOptions.body` is a // strict optional and under EOPT we must omit the key when absent rather @@ -144,6 +168,13 @@ export function useGenerateVideo( const baseOptions = { id: clientId, body: options.body, + ...(options.persistence !== undefined && { + persistence: options.persistence, + }), + ...(options.threadId !== undefined && { threadId: options.threadId }), + ...(options.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: options.initialResumeSnapshot, + }), devtoolsBridgeFactory: createVideoDevtoolsBridge, devtools: { ...options.devtools, @@ -157,29 +188,49 @@ export function useGenerateVideo( onResult: ((r: VideoGenerateResult) => options.onResult?.(r)) as ( result: VideoGenerateResult, ) => TOutput | null | void, - onError: (e: Error) => options.onError?.(e), - onProgress: (p: number, m?: string) => options.onProgress?.(p, m), - onChunk: (c: StreamChunk) => options.onChunk?.(c), - onJobCreated: (id: string) => options.onJobCreated?.(id), - onStatusUpdate: (s: VideoStatusInfo) => options.onStatusUpdate?.(s), + onError: (e: Error) => { + if (!disposed) options.onError?.(e) + }, + onProgress: (p: number, m?: string) => { + if (!disposed) options.onProgress?.(p, m) + }, + onChunk: (c: StreamChunk) => { + if (!disposed) options.onChunk?.(c) + }, + onJobCreated: (id: string) => { + if (!disposed) options.onJobCreated?.(id) + }, + onStatusUpdate: (s: VideoStatusInfo) => { + if (!disposed) options.onStatusUpdate?.(s) + }, onResultChange: (r: TOutput | null) => { + if (disposed) return result.value = r }, onLoadingChange: (l: boolean) => { + if (disposed) return isLoading.value = l }, onErrorChange: (e: Error | undefined) => { + if (disposed) return error.value = e }, onStatusChange: (s: GenerationClientState) => { + if (disposed) return status.value = s }, onJobIdChange: (id: string | null) => { + if (disposed) return jobId.value = id }, onVideoStatusChange: (s: VideoStatusInfo | null) => { + if (disposed) return videoStatus.value = s }, + onResumeStateChange: (rs: GenerationResumeState | null) => { + if (disposed) return + resumeState.value = rs + }, } let client: VideoGenerationClient @@ -212,12 +263,15 @@ export function useGenerateVideo( }, ) + // Mount devtools only. Generation runs are never auto-started on mount — + // persisted state is read-only for display. onMounted(() => { client.mountDevtools() }) // Cleanup on scope dispose: stop any in-flight requests and unregister devtools onScopeDispose(() => { + disposed = true client.dispose() }) @@ -247,5 +301,6 @@ export function useGenerateVideo( status: readonly(status), stop, reset, + resumeState: readonly(resumeState), } } diff --git a/packages/ai-vue/src/use-generation.ts b/packages/ai-vue/src/use-generation.ts index 596f14e97..e15114862 100644 --- a/packages/ai-vue/src/use-generation.ts +++ b/packages/ai-vue/src/use-generation.ts @@ -15,6 +15,10 @@ import type { GenerationClientOptions, GenerationClientState, GenerationFetcher, + GenerationPersistence, + GenerationRestoredResult, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, } from '@tanstack/ai-client' import type { DeepReadonly, ShallowRef } from 'vue' @@ -39,6 +43,21 @@ export interface UseGenerationOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** + * How this generation persists across reloads. + * - Omit / `false`: ephemeral, in-memory only. + * - `true`: server-driven — on mount the client hydrates the last generation + * for its `threadId` from the server (needs a connection with a + * `hydrateGeneration` handler) and repaints it; it never auto-starts a run. + * - a storage adapter: client-driven — the lightweight snapshot is cached under + * `generation:` as a run streams and read back on mount. Media bytes are + * never stored. + */ + persistence?: boolean | GenerationPersistence + /** Thread id for this generation, stable across reloads. Server-driven mode (`persistence: true`) hydrates the last generation under this key. Falls back to `id`. */ + threadId?: string + /** Explicit resume-snapshot seed for apps that manage storage themselves; skips automatic hydration from `persistence`. Later run events merge into it. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when a result is received. Can optionally return a transformed value. * @@ -53,16 +72,27 @@ export interface UseGenerationOptions { onProgress?: (progress: number, message?: string) => void /** Callback for each stream chunk (connect-based adapter mode only) */ onChunk?: (chunk: StreamChunk) => void + /** + * @internal Rebuild a typed result from a restored snapshot, injected by each + * specialized composable (image / audio / transcription / summarize). + * Forwarded to the client so a client-store / server-hydrate restore repaints + * `result`. + */ + reconstructResult?: (restored: GenerationRestoredResult) => TResult | null } /** * Return type for the useGeneration hook. * * @template TOutput - The output type (after optional transform) + * @template TInput - The input type accepted by `generate` (defaults to any object) */ -export interface UseGenerationReturn { +export interface UseGenerationReturn< + TOutput, + TInput extends Record = Record, +> { /** Trigger a generation request */ - generate: (input: Record) => Promise + generate: (input: TInput) => Promise /** The generation result, or null if not yet generated */ result: DeepReadonly> /** Whether a generation is currently in progress */ @@ -75,6 +105,8 @@ export interface UseGenerationReturn { stop: () => void /** Clear result, error, and return to idle */ reset: () => void + /** Identity of the in-flight run while one is streaming, or null after it ends */ + resumeState: DeepReadonly> } /** @@ -113,7 +145,10 @@ export function useGeneration< options: Omit, 'onResult'> & { onResult?: (result: TResult) => TTransformed }, -): UseGenerationReturn> { +): UseGenerationReturn< + InferGenerationOutputFromReturn, + TInput +> { type TOutput = InferGenerationOutputFromReturn const hookId = useId() const clientId = options.id || hookId @@ -122,6 +157,10 @@ export function useGeneration< const isLoading = shallowRef(false) const error = shallowRef(undefined) const status = shallowRef('idle') + const resumeState = shallowRef( + options.initialResumeSnapshot?.resumeState ?? null, + ) + let disposed = false // Conditional spread on `body`: `GenerationClientOptions.body` is a strict // optional (`body?: Record`), and under EOPT we must omit the @@ -129,6 +168,16 @@ export function useGeneration< const clientOptions: GenerationClientOptions = { id: clientId, body: options.body, + ...(options.persistence !== undefined && { + persistence: options.persistence, + }), + ...(options.threadId !== undefined && { threadId: options.threadId }), + ...(options.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: options.initialResumeSnapshot, + }), + ...(options.reconstructResult + ? { reconstructResult: options.reconstructResult } + : {}), devtoolsBridgeFactory: createGenerationDevtoolsBridge, devtools: { ...options.devtools, @@ -141,21 +190,35 @@ export function useGeneration< onResult: ((r: TResult) => options.onResult?.(r)) as ( result: TResult, ) => TOutput | null | void, - onError: (e: Error) => options.onError?.(e), - onProgress: (p: number, m?: string) => options.onProgress?.(p, m), - onChunk: (c: StreamChunk) => options.onChunk?.(c), + onError: (e: Error) => { + if (!disposed) options.onError?.(e) + }, + onProgress: (p: number, m?: string) => { + if (!disposed) options.onProgress?.(p, m) + }, + onChunk: (c: StreamChunk) => { + if (!disposed) options.onChunk?.(c) + }, onResultChange: (r: TOutput | null) => { + if (disposed) return result.value = r }, onLoadingChange: (l: boolean) => { + if (disposed) return isLoading.value = l }, onErrorChange: (e: Error | undefined) => { + if (disposed) return error.value = e }, onStatusChange: (s: GenerationClientState) => { + if (disposed) return status.value = s }, + onResumeStateChange: (rs: GenerationResumeState | null) => { + if (disposed) return + resumeState.value = rs + }, } let client: GenerationClient @@ -188,12 +251,15 @@ export function useGeneration< }, ) + // Mount devtools only. Generation runs are never auto-started on mount — + // persisted state is read-only for display. onMounted(() => { client.mountDevtools() }) // Cleanup on scope dispose: stop any in-flight requests and unregister devtools onScopeDispose(() => { + disposed = true client.dispose() }) @@ -210,7 +276,7 @@ export function useGeneration< } return { - generate: generate as (input: Record) => Promise, + generate, // `readonly()` distributes `DeepReadonly`/`UnwrapNestedRefs` over the // `TOutput` conditional, which TS can't prove equal to the declared // `DeepReadonly>` while `TTransformed` is free. @@ -221,5 +287,6 @@ export function useGeneration< status: readonly(status), stop, reset, + resumeState: readonly(resumeState), } } diff --git a/packages/ai-vue/src/use-summarize.ts b/packages/ai-vue/src/use-summarize.ts index 7297c5277..d00a32e9b 100644 --- a/packages/ai-vue/src/use-summarize.ts +++ b/packages/ai-vue/src/use-summarize.ts @@ -1,4 +1,9 @@ import { useGeneration } from './use-generation' +import { reconstructSummarizeResult } from '@tanstack/ai-client' +import type { + UseGenerationOptions, + UseGenerationReturn, +} from './use-generation' import type { StreamChunk, SummarizationResult } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -15,7 +20,12 @@ import type { DeepReadonly, ShallowRef } from 'vue' * * @template TOutput - The output type after optional transform (defaults to SummarizationResult) */ -export interface UseSummarizeOptions { +export interface UseSummarizeOptions< + TOutput = SummarizationResult, +> extends Pick< + UseGenerationOptions, + 'persistence' | 'threadId' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for summarization */ @@ -47,7 +57,10 @@ export interface UseSummarizeOptions { * * @template TOutput - The output type (after optional transform) */ -export interface UseSummarizeReturn { +export interface UseSummarizeReturn extends Omit< + UseGenerationReturn, + 'generate' +> { /** Trigger summarization */ generate: (input: SummarizeGenerateInput) => Promise /** The summarization result, or null */ @@ -58,10 +71,6 @@ export interface UseSummarizeReturn { error: DeepReadonly> /** Current state of the generation */ status: DeepReadonly> - /** Abort the current summarization */ - stop: () => void - /** Clear result, error, and return to idle */ - reset: () => void } /** @@ -106,19 +115,15 @@ export function useSummarize( hookName: 'useSummarize', outputKind: 'text' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - useGeneration({ - ...options, - devtools, - }) + const generation = useGeneration< + SummarizeGenerateInput, + SummarizationResult, + TTransformed + >({ + ...options, + devtools, + reconstructResult: reconstructSummarizeResult, + }) - return { - generate: generate as (input: SummarizeGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + return generation } diff --git a/packages/ai-vue/src/use-transcription.ts b/packages/ai-vue/src/use-transcription.ts index 26e42156c..6bfae4c91 100644 --- a/packages/ai-vue/src/use-transcription.ts +++ b/packages/ai-vue/src/use-transcription.ts @@ -1,4 +1,9 @@ import { useGeneration } from './use-generation' +import { reconstructTranscriptionResult } from '@tanstack/ai-client' +import type { + UseGenerationOptions, + UseGenerationReturn, +} from './use-generation' import type { StreamChunk, TranscriptionResult } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -15,7 +20,16 @@ import type { DeepReadonly, ShallowRef } from 'vue' * * @template TOutput - The output type after optional transform (defaults to TranscriptionResult) */ -export interface UseTranscriptionOptions { +export interface UseTranscriptionOptions< + TOutput = TranscriptionResult, +> extends Pick< + UseGenerationOptions< + TranscriptionGenerateInput, + TranscriptionResult, + TOutput + >, + 'persistence' | 'threadId' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for transcription */ @@ -47,7 +61,9 @@ export interface UseTranscriptionOptions { * * @template TOutput - The output type (after optional transform) */ -export interface UseTranscriptionReturn { +export interface UseTranscriptionReturn< + TOutput = TranscriptionResult, +> extends Omit, 'generate'> { /** Trigger transcription */ generate: (input: TranscriptionGenerateInput) => Promise /** The transcription result, or null */ @@ -58,10 +74,6 @@ export interface UseTranscriptionReturn { error: DeepReadonly> /** Current state of the generation */ status: DeepReadonly> - /** Abort the current transcription */ - stop: () => void - /** Clear result, error, and return to idle */ - reset: () => void } /** @@ -111,20 +123,11 @@ export function useTranscription( hookName: 'useTranscription', outputKind: 'text' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - useGeneration< - TranscriptionGenerateInput, - TranscriptionResult, - TTransformed - >({ ...options, devtools }) + const generation = useGeneration< + TranscriptionGenerateInput, + TranscriptionResult, + TTransformed + >({ ...options, devtools, reconstructResult: reconstructTranscriptionResult }) - return { - generate: generate as (input: TranscriptionGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, - } + return generation } diff --git a/packages/ai-vue/tests/use-generation.test.ts b/packages/ai-vue/tests/use-generation.test.ts index 21b40eafe..dd30c72e0 100644 --- a/packages/ai-vue/tests/use-generation.test.ts +++ b/packages/ai-vue/tests/use-generation.test.ts @@ -9,13 +9,28 @@ import { useTranscription } from '../src/use-transcription' import { useSummarize } from '../src/use-summarize' import { useGenerateVideo } from '../src/use-generate-video' import { createMockConnectionAdapter } from './test-utils' -import type { StreamChunk, TTSResult, TranscriptionResult } from '@tanstack/ai' +import type { + PersistedArtifactRef, + StreamChunk, + TTSResult, + TranscriptionResult, +} from '@tanstack/ai' +import type { + ConnectConnectionAdapter, + GenerationResumeSnapshot, + RunAgentInputContext, +} from '@tanstack/ai-client' import type { DeepReadonly } from 'vue' // Helper to create generation stream chunks function createGenerationChunks(result: unknown): Array { return [ - { type: 'RUN_STARTED', runId: 'run-1', timestamp: Date.now() }, + { + type: 'RUN_STARTED', + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + }, { type: 'CUSTOM', name: 'generation:result', @@ -25,6 +40,7 @@ function createGenerationChunks(result: unknown): Array { { type: 'RUN_FINISHED', runId: 'run-1', + threadId: 'thread-1', finishReason: 'stop', timestamp: Date.now(), }, @@ -62,6 +78,60 @@ function createVideoChunks(jobId: string, url: string): Array { ] as unknown as Array } +const videoResumeSnapshot: GenerationResumeSnapshot = { + resumeState: { + threadId: 'thread-resume', + runId: 'run-resume', + }, + status: 'running', +} + +function createReplayVideoChunks(): Array { + return [ + { + type: 'RUN_STARTED', + runId: 'run-resume', + threadId: 'thread-resume', + timestamp: Date.now(), + }, + { + type: 'CUSTOM', + name: 'generation:result', + value: { + jobId: 'job-replay', + status: 'completed', + url: 'https://example.com/video.mp4', + }, + timestamp: Date.now(), + }, + { + type: 'RUN_FINISHED', + runId: 'run-resume', + threadId: 'thread-resume', + timestamp: Date.now(), + }, + ] as unknown as Array +} + +function createRunContextCaptureAdapter(chunks: Array): { + adapter: ConnectConnectionAdapter + connect: ReturnType + runContexts: Array +} { + const runContexts: Array = [] + const connect = vi.fn() + const adapter: ConnectConnectionAdapter = { + async *connect(_messages, _data, _signal, runContext) { + connect(runContext) + runContexts.push(runContext) + for (const chunk of chunks) { + yield chunk + } + }, + } + return { adapter, connect, runContexts } +} + // Helper to create error stream chunks function createErrorChunks(message: string): Array { return [ @@ -187,6 +257,36 @@ describe('useGeneration', () => { expect(result.status.value).toBe('error') expect(result.error.value?.message).toBe('Generation failed') }) + + it('does not auto-fire a generation on mount from a persisted running snapshot', async () => { + // Regression guard for the removed generation resume surface. + const snapshot: GenerationResumeSnapshot = { + resumeState: { threadId: 'thread-resume', runId: 'run-resume' }, + status: 'running', + } + const { adapter, connect } = createRunContextCaptureAdapter( + createGenerationChunks({ id: '1' }), + ) + const getItem = vi.fn(() => snapshot) + const { result } = renderHook(() => + useGeneration({ + id: 'no-auto-fire', + connection: adapter, + persistence: { getItem, setItem: vi.fn(), removeItem: vi.fn() }, + initialResumeSnapshot: snapshot, + }), + ) + + await flushPromises() + await nextTick() + + expect(connect).not.toHaveBeenCalled() + expect(getItem).not.toHaveBeenCalled() + expect(result.isLoading.value).toBe(false) + expect(result.status.value).toBe('idle') + // The persisted snapshot remains exposed as read-only state. + expect(result.resumeState.value).toEqual(snapshot.resumeState) + }) }) describe('stop and reset', () => { @@ -245,6 +345,117 @@ describe('useGeneration', () => { }) }) + describe('persistence', () => { + it('repaints status from a stored complete snapshot on mount without starting a run', async () => { + const { adapter, connect } = createRunContextCaptureAdapter( + createGenerationChunks({ id: '1' }), + ) + const getItem = vi.fn(() => ({ + resumeState: null, + status: 'complete' as const, + result: { id: 'result-1', model: 'image-model' }, + })) + + const { result } = renderHook(() => + useGeneration({ + id: 'hydrated', + connection: adapter, + persistence: { getItem, setItem: vi.fn(), removeItem: vi.fn() }, + }), + ) + + await flushPromises() + await nextTick() + + // A completed snapshot repaints the normal `status` field to `success`. + expect(result.status.value).toBe('success') + expect(getItem).toHaveBeenCalledWith('generation:hydrated') + expect(connect).not.toHaveBeenCalled() + // The base hook injects no reconstructResult, so `result` stays null. + expect(result.result.value).toBeNull() + expect(result.resumeState.value).toBeNull() + }) + + it('repaints resumeState from a stored running snapshot on mount', async () => { + const { adapter, connect } = createRunContextCaptureAdapter( + createGenerationChunks({ id: '1' }), + ) + const getItem = vi.fn(() => ({ + resumeState: { threadId: 'thread-resume', runId: 'run-resume' }, + status: 'running' as const, + })) + + const { result } = renderHook(() => + useGeneration({ + id: 'running-hydrate', + connection: adapter, + persistence: { getItem, setItem: vi.fn(), removeItem: vi.fn() }, + }), + ) + + await flushPromises() + await nextTick() + + expect(result.resumeState.value).toEqual({ + threadId: 'thread-resume', + runId: 'run-resume', + }) + // A restored running run presents as `generating` but never auto-tails. + expect(result.status.value).toBe('generating') + expect(result.isLoading.value).toBe(false) + expect(connect).not.toHaveBeenCalled() + }) + + it('server-driven (persistence: true) hydrates from the connection by threadId without a local store', async () => { + const { adapter, connect } = createRunContextCaptureAdapter( + createGenerationChunks({ id: '1' }), + ) + const hydrateGeneration = vi.fn(async () => ({ + resumeSnapshot: { + schemaVersion: 1 as const, + resumeState: null, + status: 'complete' as const, + result: { id: 'server-result', model: 'image-model' }, + }, + activeRun: null, + })) + + const { result } = renderHook(() => + useGeneration({ + threadId: 'thread-server', + connection: { ...adapter, hydrateGeneration }, + persistence: true, + }), + ) + + await flushPromises() + await nextTick() + + // The server snapshot repaints `status` to `success`. + expect(result.status.value).toBe('success') + expect(hydrateGeneration).toHaveBeenCalledWith('thread-server') + expect(connect).not.toHaveBeenCalled() + }) + + it('clears resumeState once a streamed run finishes', async () => { + const adapter = createMockConnectionAdapter({ + chunks: createGenerationChunks({ id: '1' }), + }) + + const { result } = renderHook(() => + useGeneration({ connection: adapter }), + ) + + await result.generate({ prompt: 'replay' }) + await flushPromises() + await nextTick() + + expect(result.status.value).toBe('success') + // Once the run ends the in-flight identity is gone. + expect(result.resumeState.value).toBeNull() + }) + }) + describe('error handling', () => { it('should require either connection or fetcher', () => { expect(() => { @@ -354,6 +565,61 @@ describe('useGenerateImage', () => { expect(typeof result.stop).toBe('function') expect(typeof result.reset).toBe('function') }) + + it('restores a completed image result from a durable artifact url', async () => { + // useGenerateImage injects `reconstructImageResult`, so a restored complete + // snapshot repaints `result` with the durable serve url — as if the run had + // just finished. + const artifact: PersistedArtifactRef = { + role: 'output', + artifactId: 'artifact-image-1', + threadId: 'thread-img', + runId: 'run-img', + name: 'image.png', + mimeType: 'image/png', + size: 2048, + createdAt: '2026-07-06T00:00:00.000Z', + url: '/api/artifacts/artifact-image-1', + source: { + activity: 'image', + path: 'runs/run-img/image.png', + provider: 'test', + model: 'test-image', + mediaType: 'image', + }, + } + const getItem = vi.fn(() => ({ + resumeState: null, + status: 'complete' as const, + activity: 'image' as const, + result: { + id: 'img-restored', + model: 'test-image', + artifacts: [artifact], + }, + })) + const adapter = createMockConnectionAdapter() + + const { result } = renderHook(() => + useGenerateImage({ + id: 'img-hydrate', + connection: adapter, + persistence: { getItem, setItem: vi.fn(), removeItem: vi.fn() }, + }), + ) + + await flushPromises() + await nextTick() + + expect(result.status.value).toBe('success') + expect(result.result.value).toEqual({ + id: 'img-restored', + model: 'test-image', + images: [{ url: '/api/artifacts/artifact-image-1' }], + artifacts: [artifact], + }) + expect(result.resumeState.value).toBeNull() + }) }) describe('useGenerateSpeech', () => { @@ -586,7 +852,7 @@ describe('useSummarize', () => { it('should summarize text using fetcher', async () => { const mockResult = { summary: 'A brief summary', - model: 'gpt-4', + model: 'gpt-5.5', } const { result } = renderHook(() => @@ -604,7 +870,7 @@ describe('useSummarize', () => { }) it('should summarize text using connection', async () => { - const mockResult = { summary: 'A brief summary', model: 'gpt-4' } + const mockResult = { summary: 'A brief summary', model: 'gpt-5.5' } const chunks = createGenerationChunks(mockResult) const adapter = createMockConnectionAdapter({ chunks }) @@ -761,6 +1027,32 @@ describe('useGenerateVideo', () => { expect(result.status.value).toBe('idle') }) + it('does not auto-fire a video generation on mount from a persisted running snapshot', async () => { + // Regression guard for the removed generation resume surface (video). + const { adapter, connect } = createRunContextCaptureAdapter( + createReplayVideoChunks(), + ) + const getItem = vi.fn(() => videoResumeSnapshot) + const { result } = renderHook(() => + useGenerateVideo({ + id: 'video-no-auto-fire', + connection: adapter, + persistence: { getItem, setItem: vi.fn(), removeItem: vi.fn() }, + initialResumeSnapshot: videoResumeSnapshot, + }), + ) + + await flushPromises() + await nextTick() + + expect(connect).not.toHaveBeenCalled() + expect(getItem).not.toHaveBeenCalled() + expect(result.isLoading.value).toBe(false) + expect(result.status.value).toBe('idle') + // The seeded in-flight identity is exposed as read-only `resumeState`. + expect(result.resumeState.value).toEqual(videoResumeSnapshot.resumeState) + }) + it('should require either connection or fetcher', () => { expect(() => { renderHook(() => useGenerateVideo({} as any)) diff --git a/packages/ai/skills/ai-core/client-persistence/SKILL.md b/packages/ai/skills/ai-core/client-persistence/SKILL.md index 3e6eb9b3b..343d50194 100644 --- a/packages/ai/skills/ai-core/client-persistence/SKILL.md +++ b/packages/ai/skills/ai-core/client-persistence/SKILL.md @@ -7,6 +7,10 @@ description: > client cache). Reload restore, pending interrupts, mid-stream rejoin with delivery durability. Use for SPA reload durability — NOT server history alone. + Also covers generation hooks (useGenerateImage etc.), same two modes as chat: + client-driven (adapter) persists a lightweight resume snapshot under + generation:; server-driven (persistence: true + threadId) hydrates the last + generation from the server on mount, nothing cached. No extra package: the adapters ship in the framework packages. type: sub-skill library: tanstack-ai @@ -110,6 +114,101 @@ chat's identity _is_ its `threadId`. Without a stable one, each load is a new chat. Generate it server-side or from a route param the user owns; do not randomize per mount. +## Generation hooks: two modes, mirroring chat + +The generation hooks (`useGenerateImage`, `useGenerateVideo`, `useGeneration`, +`useSummarize`, `useTranscription`, …) take the **same `persistence` option** as +`useChat` — `boolean | adapter` — with the same two-mode split. **The hooks are +transparent, mirroring `useChat`:** whichever mode, a reload repaints the hook's +**normal** fields — `status` (`'idle'` / `'generating'` / `'success'` / +`'error'`), `error`, and `result` — as if the run had just finished. There is +**no** `resumeSnapshot`, `pendingArtifacts`, or `resultArtifacts` field. The one +extra field is `resumeState`: the in-flight run identity (`{ threadId, runId, +pendingArtifacts? }`) or `null` — non-null only while a run streams. The +persisted record holds run identity, status, error, and result metadata (ids, +model, video `jobId`), **never the generated media bytes**. + +The hook return is exactly `generate` / `result` / `isLoading` / `error` / +`status` / `stop` / `reset` / `resumeState`. + +### Mode A — client-driven (a storage adapter) + +```tsx +const image = useGenerateImage({ + id: 'hero-image', // stable — the storage key is `generation:` + connection: fetchServerSentEvents('/api/generate/image'), + persistence: localStoragePersistence(), // bare — no type argument +}) +// After a reload: image.status / image.result / image.error are the last run's +// outcome, read exactly like a fresh run. image.resumeState is non-null only +// WHILE a run is streaming. +``` + +- The lightweight snapshot is cached in the browser under `generation:` as a + run streams, and read back on mount. +- `localStoragePersistence()` (and the session / IndexedDB factories) take **no** + type argument here — a bare call defaults to the generation snapshot shape. +- Hydration is automatic on mount and validated + (`parseGenerationResumeSnapshot`); an explicit `initialResumeSnapshot` seed + skips it. +- The `generation:` key segment means a chat and a generation client can share + an id and an adapter without colliding. +- `result` needs byte storage to come back. A client snapshot never holds the + bytes, so on its own a reload restores `status` and `error` while `result` + stays `null` (see byte storage below). + +### Mode B — server-driven (`persistence: true`) + +```tsx +const image = useGenerateImage({ + threadId, // stable — the key the last generation is hydrated under (falls back to id) + connection: fetchServerSentEvents('/api/generate/image'), + persistence: true, +}) +// After a reload: image.status / image.result / image.error are the last +// generation for `threadId`, fetched from the server — nothing was cached. +``` + +- Nothing is cached client-side. On mount the client hydrates the **last + generation** for its `threadId` from the server via the connection's + `hydrateGeneration` handler (the SSE/HTTP adapters issue a `GET` with + `?threadId=` to the same endpoint URL) and repaints it into the normal fields. +- The server `GET` returns `reconstructGeneration(persistence, request)` from + `@tanstack/ai-persistence` — it resolves the job by `?jobId=` (preferred) or + the latest job linked to `?threadId=`, and needs `stores.jobs`. Pair it with + `withGenerationPersistence` on the generation route. See + `ai-core/media-generation` and `ai-persistence`. +- Best for multi-device / compliance (no generation metadata in browser + storage), exactly like chat Mode B. + +### Restoring media: byte storage + `artifactUrl` + +`result` comes back with its media only when the **server** persists the bytes +(`stores.artifacts` + `stores.blobs`) AND `withGenerationPersistence` is given an +`artifactUrl` mapper: + +```ts +withGenerationPersistence(persistence, { + artifactUrl: (ref) => `/api/generate/image/artifact?id=${ref.artifactId}`, +}) +``` + +`artifactUrl` stamps a durable app-origin URL onto each persisted ref and +rewrites the live result's media to it, so live and restored results match. The +durable refs travel on `result.artifacts`; on restore the hook rebuilds `result` +from them, so `result.images[i].url` (or a video's `result.url`) serves from your +own origin. Final refs live on `result.artifacts`; in-flight ones on +`resumeState.pendingArtifacts`. Without byte storage, a reload restores `status` +/ `error` and `result` stays `null`. + +Common to both modes: + +- `stop()` marks the record no longer resumable; `reset()` deletes it (Mode A) or + clears the in-memory snapshot (Mode B). +- Nothing auto-runs from a hydrated record — `generate(...)` is always explicit. +- Use `status` / `result` for a finished run; use `resumeState` to tell that a + run was still generating when the page closed. + ## Common mistakes ### HIGH: No `threadId` diff --git a/packages/ai/skills/ai-core/media-generation/SKILL.md b/packages/ai/skills/ai-core/media-generation/SKILL.md index 268a85f74..0de7843cd 100644 --- a/packages/ai/skills/ai-core/media-generation/SKILL.md +++ b/packages/ai/skills/ai-core/media-generation/SKILL.md @@ -563,21 +563,106 @@ if (result.usage?.unitsBilled != null) { For video, the units arrive with the completed result: `getVideoJobStatus()` returns `usage` and emits a `video:usage` devtools event when fal reports it. +### 7. Durable persistence (job lifecycle + artifact bytes) + +To make generations survive a server restart and be re-served later, add +`withGenerationPersistence` from `@tanstack/ai-persistence` as generation +middleware. It requires `stores.jobs` (a `GenerationJobStore`, keyed on `jobId` — +not a chat `threadId`) and, when you also pass an `stores.artifacts` + +`stores.blobs` **pair** (both or neither), it persists the generated media bytes +at blob key `artifacts//` with an `ArtifactRecord` per file. +`memoryPersistence()` ships all three for dev/tests. + +```typescript +import { generateImage, toServerSentEventsResponse } from '@tanstack/ai' +import { openaiImage } from '@tanstack/ai-openai' +import { + withGenerationPersistence, + memoryPersistence, + retrieveArtifact, + retrieveBlob, + reconstructGeneration, +} from '@tanstack/ai-persistence' + +const persistence = memoryPersistence() // swap for your DB/object-store adapter + +export async function POST(req: Request) { + const { prompt, threadId } = await req.json() + return toServerSentEventsResponse( + generateImage({ + adapter: openaiImage('gpt-image-1'), + prompt, + threadId, // optional link recorded on the job + artifacts + stream: true, + middleware: [ + withGenerationPersistence(persistence, { + // Stamp a durable app-origin serve URL (the GET route below) onto + // each persisted artifact ref, and rewrite the live result's media to + // it. Both live and restored results then render from your origin. + artifactUrl: (ref) => `/api/artifacts?id=${ref.artifactId}`, + }), + ], + }), + ) +} + +// Serve the stored bytes back (GET /api/artifacts?id=…): +export async function GET(req: Request) { + const id = new URL(req.url).searchParams.get('id') ?? '' + const record = await retrieveArtifact(persistence, id) + if (!record) return new Response('Not found', { status: 404 }) + const blob = await retrieveBlob(persistence, record) + if (!blob?.body) return new Response('Not found', { status: 404 }) + return new Response(blob.body, { + headers: { 'content-type': record.mimeType }, + }) +} +``` + +On the client, the generation hooks mirror chat's two persistence modes: +`persistence: ` (client-driven, caches a lightweight snapshot under +`generation:`) or `persistence: true` + a stable `threadId` (server-driven — +hydrates the last generation for the thread on mount via the connection's +`hydrateGeneration` handler, backed by a `reconstructGeneration` GET route). The +hooks are transparent (like `useChat`): a reload repaints `status` / `result` / +`error`, not a separate `resumeSnapshot`. Neither mode stores media bytes on the +client — but because the `artifactUrl` above stamps a durable URL onto each ref +(carried on `result.artifacts`), the restored `result` rebuilds its media from +those refs and serves from your own origin. Without byte storage, a reload +restores `status` / `error` and `result` stays `null`. + +- Building the R2/D1-backed byte stores for a Cloudflare Worker: + **ai-persistence/build-cloudflare-artifact-store**. +- Store contracts, `composePersistence`, and the two client modes end-to-end: + `docs/persistence/generation-persistence.md` and the + `ai-core/client-persistence` sub-skill. + --- ## Common Hook API All generation hooks return the same shape: -| Property | Type | Description | -| ----------- | -------------------------- | ------------------------------------------------ | -| `generate` | `(input) => Promise` | Trigger generation | -| `result` | `T \| null` | Result (optionally transformed via `onResult`) | -| `isLoading` | `boolean` | Whether generation is in progress | -| `error` | `Error \| undefined` | Current error | -| `status` | `GenerationClientState` | `'idle' \| 'generating' \| 'success' \| 'error'` | -| `stop` | `() => void` | Abort current generation | -| `reset` | `() => void` | Clear state, return to idle | +| Property | Type | Description | +| ------------- | ------------------------------- | ------------------------------------------------ | +| `generate` | `(input) => Promise` | Trigger generation | +| `result` | `T \| null` | Result (optionally transformed via `onResult`) | +| `isLoading` | `boolean` | Whether generation is in progress | +| `error` | `Error \| undefined` | Current error | +| `status` | `GenerationClientState` | `'idle' \| 'generating' \| 'success' \| 'error'` | +| `stop` | `() => void` | Abort current generation | +| `reset` | `() => void` | Clear state (and any persisted snapshot) | +| `resumeState` | `GenerationResumeState \| null` | Identity of the run WHILE it streams; null after | + +The hook is **transparent**, mirroring `useChat`: there is no `resumeSnapshot`, +`pendingArtifacts`, or `resultArtifacts` field. Hooks also accept `persistence` +(a browser storage adapter such as `localStoragePersistence()` — a bare call, no +type argument) plus a stable `id`: the client then writes a lightweight snapshot +under `generation:` as the run streams and hydrates it back on mount, +repainting the **normal** `status` / `result` / `error` fields so the last run +survives a reload (metadata only — never media bytes; `result`'s media returns +only with server byte storage + `artifactUrl`). See `ai-core/client-persistence` +for details. Provide either `connection` (streaming SSE transport) or `fetcher` (direct async call / server function returning `Response`). Use `onResult` diff --git a/packages/ai/src/activities/generateAudio/index.ts b/packages/ai/src/activities/generateAudio/index.ts index ec2e72890..cb25e2a0b 100644 --- a/packages/ai/src/activities/generateAudio/index.ts +++ b/packages/ai/src/activities/generateAudio/index.ts @@ -9,6 +9,7 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { streamGenerationResult } from '../stream-generation-result.js' import { resolveDebugOption } from '../../logger/resolve' import { + applyGenerationResultTransforms, createGenerationContext, runGenerationError, runGenerationFinish, @@ -84,6 +85,10 @@ export interface AudioActivityOptions< * `GenerationMiddleware` contract for a custom backend. */ middleware?: Array + /** Stable conversation/thread id for correlating this run when persisted. */ + threadId?: string + /** Stable run id for correlating this run when persisted. */ + runId?: string } // =========================== @@ -134,8 +139,9 @@ export function generateAudio< options: AudioActivityOptions, ): AudioActivityResult { if (options.stream) { - return streamGenerationResult(() => - runGenerateAudio(options), + return streamGenerationResult( + (resolved) => runGenerateAudio({ ...options, ...resolved }), + options, ) as AudioActivityResult } return runGenerateAudio(options) as AudioActivityResult @@ -154,6 +160,8 @@ async function runGenerateAudio< stream: _stream, debug: _debug, middleware, + threadId, + runId, ...rest } = options const model = adapter.model @@ -171,6 +179,9 @@ async function runGenerateAudio< provider: adapter.name, model, modelOptions: rest.modelOptions, + threadId, + runId, + artifactInputs: { prompt: rest.prompt, duration: rest.duration }, createId, }) @@ -192,7 +203,8 @@ async function runGenerateAudio< }) try { - const result = await adapter.generateAudio({ ...rest, model, logger }) + const rawResult = await adapter.generateAudio({ ...rest, model, logger }) + const result = await applyGenerationResultTransforms(mwCtx, rawResult) const elapsedMs = Date.now() - startTime aiEventClient.emit('audio:request:completed', { diff --git a/packages/ai/src/activities/generateImage/index.ts b/packages/ai/src/activities/generateImage/index.ts index 8021e0ce2..5e49baad1 100644 --- a/packages/ai/src/activities/generateImage/index.ts +++ b/packages/ai/src/activities/generateImage/index.ts @@ -9,6 +9,7 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { streamGenerationResult } from '../stream-generation-result.js' import { resolveDebugOption } from '../../logger/resolve' import { + applyGenerationResultTransforms, createGenerationContext, runGenerationError, runGenerationFinish, @@ -137,6 +138,10 @@ export type ImageActivityOptions< * `GenerationMiddleware` contract for a custom backend. */ middleware?: Array + /** Stable conversation/thread id for correlating this run when persisted. */ + threadId?: string + /** Stable run id for correlating this run when persisted. */ + runId?: string } & ({} extends ImageProviderOptionsForModel ? { /** Provider-specific options for image generation */ modelOptions?: ImageProviderOptionsForModel< @@ -225,8 +230,9 @@ export function generateImage< options: ImageActivityOptions, ): ImageActivityResult { if (options.stream) { - return streamGenerationResult(() => - runGenerateImage(options), + return streamGenerationResult( + (resolved) => runGenerateImage({ ...options, ...resolved }), + options, ) as ImageActivityResult } @@ -247,6 +253,8 @@ async function runGenerateImage< stream: _stream, debug: _debug, middleware, + threadId, + runId, ...rest } = options const model = adapter.model @@ -260,6 +268,9 @@ async function runGenerateImage< provider: adapter.name, model, modelOptions: rest.modelOptions, + threadId, + runId, + artifactInputs: { prompt: rest.prompt }, createId, }) @@ -295,7 +306,8 @@ async function runGenerateImage< }) try { - const result = await adapter.generateImages({ ...rest, model, logger }) + const rawResult = await adapter.generateImages({ ...rest, model, logger }) + const result = await applyGenerationResultTransforms(mwCtx, rawResult) const duration = Date.now() - startTime aiEventClient.emit('image:request:completed', { diff --git a/packages/ai/src/activities/generateSpeech/index.ts b/packages/ai/src/activities/generateSpeech/index.ts index 06ac193c8..0a134f4cf 100644 --- a/packages/ai/src/activities/generateSpeech/index.ts +++ b/packages/ai/src/activities/generateSpeech/index.ts @@ -9,6 +9,7 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { streamGenerationResult } from '../stream-generation-result.js' import { resolveDebugOption } from '../../logger/resolve' import { + applyGenerationResultTransforms, createGenerationContext, runGenerationError, runGenerationFinish, @@ -87,6 +88,10 @@ export interface TTSActivityOptions< * `GenerationMiddleware` contract for a custom backend. */ middleware?: Array + /** Stable conversation/thread id for correlating this run when persisted. */ + threadId?: string + /** Stable run id for correlating this run when persisted. */ + runId?: string } // =========================== @@ -144,8 +149,9 @@ export function generateSpeech< TStream extends boolean = false, >(options: TTSActivityOptions): TTSActivityResult { if (options.stream) { - return streamGenerationResult(() => - runGenerateSpeech(options), + return streamGenerationResult( + (resolved) => runGenerateSpeech({ ...options, ...resolved }), + options, ) as TTSActivityResult } return runGenerateSpeech(options) as TTSActivityResult @@ -162,6 +168,8 @@ async function runGenerateSpeech< stream: _stream, debug: _debug, middleware, + threadId, + runId, ...rest } = options const model = adapter.model @@ -179,6 +187,14 @@ async function runGenerateSpeech< provider: adapter.name, model, modelOptions: rest.modelOptions, + artifactInputs: { + text: rest.text, + voice: rest.voice, + format: rest.format, + speed: rest.speed, + }, + threadId, + runId, createId, }) @@ -202,7 +218,8 @@ async function runGenerateSpeech< }) try { - const result = await adapter.generateSpeech({ ...rest, model, logger }) + const rawResult = await adapter.generateSpeech({ ...rest, model, logger }) + const result = await applyGenerationResultTransforms(mwCtx, rawResult) const duration = Date.now() - startTime aiEventClient.emit('speech:request:completed', { diff --git a/packages/ai/src/activities/generateTranscription/index.ts b/packages/ai/src/activities/generateTranscription/index.ts index 03c6a9ed1..ccf85f235 100644 --- a/packages/ai/src/activities/generateTranscription/index.ts +++ b/packages/ai/src/activities/generateTranscription/index.ts @@ -9,6 +9,7 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { streamGenerationResult } from '../stream-generation-result.js' import { resolveDebugOption } from '../../logger/resolve' import { + applyGenerationResultTransforms, createGenerationContext, runGenerationError, runGenerationFinish, @@ -94,6 +95,10 @@ export interface TranscriptionActivityOptions< * `GenerationMiddleware` contract for a custom backend. */ middleware?: Array + /** Stable conversation/thread id for correlating this run when persisted. */ + threadId?: string + /** Stable run id for correlating this run when persisted. */ + runId?: string } // =========================== @@ -171,8 +176,9 @@ export function generateTranscription< options: TranscriptionActivityOptions, ): TranscriptionActivityResult { if (options.stream) { - return streamGenerationResult(() => - runGenerateTranscription(options), + return streamGenerationResult( + (resolved) => runGenerateTranscription({ ...options, ...resolved }), + options, ) as TranscriptionActivityResult } @@ -197,6 +203,8 @@ async function runGenerateTranscription< stream: _stream, debug: _debug, middleware, + threadId, + runId, ...rest } = options const model = adapter.model @@ -214,6 +222,14 @@ async function runGenerateTranscription< provider: adapter.name, model, modelOptions: rest.modelOptions, + artifactInputs: { + audio: rest.audio, + language: rest.language, + prompt: rest.prompt, + responseFormat: rest.responseFormat, + }, + threadId, + runId, createId, }) @@ -236,7 +252,8 @@ async function runGenerateTranscription< }) try { - const result = await adapter.transcribe({ ...rest, model, logger }) + const rawResult = await adapter.transcribe({ ...rest, model, logger }) + const result = await applyGenerationResultTransforms(mwCtx, rawResult) const duration = Date.now() - startTime aiEventClient.emit('transcription:request:completed', { diff --git a/packages/ai/src/activities/middleware/index.ts b/packages/ai/src/activities/middleware/index.ts index 79454ac5c..d5123374d 100644 --- a/packages/ai/src/activities/middleware/index.ts +++ b/packages/ai/src/activities/middleware/index.ts @@ -9,6 +9,8 @@ export type { GenerationAbortInfo, GenerationErrorInfo, AnyGenerationMiddleware, + GenerationResultTransform, + GenerationResultTransformContext, } from './types' export { createGenerationContext, diff --git a/packages/ai/src/activities/middleware/run.ts b/packages/ai/src/activities/middleware/run.ts index 6983b1e55..a1793cda9 100644 --- a/packages/ai/src/activities/middleware/run.ts +++ b/packages/ai/src/activities/middleware/run.ts @@ -4,6 +4,7 @@ import type { GenerationFinishInfo, GenerationMiddleware, GenerationMiddlewareContext, + GenerationResultTransformContext, GenerationUsageInfo, } from './types' @@ -19,6 +20,9 @@ export function createGenerationContext(args: { provider: string model: string modelOptions?: unknown + threadId?: string + runId?: string + artifactInputs?: unknown createId: (prefix: string) => string }): GenerationMiddlewareContext { return { @@ -27,9 +31,13 @@ export function createGenerationContext(args: { provider: args.provider, model: args.model, modelOptions: args.modelOptions, + threadId: args.threadId, + runId: args.runId, source: 'server', createId: args.createId, context: undefined, + resultTransforms: [], + artifactInputs: args.artifactInputs, } } @@ -86,3 +94,26 @@ export function runGenerationError( ): Promise { return run(middleware, (mw) => mw.onError?.(ctx, info)) } + +/** + * Apply the result transforms middleware registered on the context, in order, + * to the raw adapter result. Each transform may return a replacement result or + * `undefined` to leave it unchanged. Runs after the adapter result exists and + * before the final result is returned or streamed. + */ +export async function applyGenerationResultTransforms( + ctx: GenerationMiddlewareContext, + result: TResult, +): Promise { + let current = result + const transformCtx: GenerationResultTransformContext = { middleware: ctx } + + for (const transform of ctx.resultTransforms ?? []) { + const transformed = await transform(current, transformCtx) + if (transformed !== undefined) { + current = transformed as TResult + } + } + + return current +} diff --git a/packages/ai/src/activities/middleware/types.ts b/packages/ai/src/activities/middleware/types.ts index 99ae3e44e..ef1136413 100644 --- a/packages/ai/src/activities/middleware/types.ts +++ b/packages/ai/src/activities/middleware/types.ts @@ -60,6 +60,10 @@ export interface GenerationMiddlewareContext { provider: string /** Model id. Emitted as `gen_ai.request.model`. */ model: string + /** Stable conversation/thread id, when supplied by the caller. */ + threadId?: string + /** Stable run id, when supplied by the caller. */ + runId?: string /** * Provider-specific options passed to the activity, if any. Typed `unknown` * because each activity's options are strongly typed per model; a supertype @@ -72,8 +76,36 @@ export interface GenerationMiddlewareContext { createId: (prefix: string) => string /** Runtime context provided by the activity options, if any. */ context: TContext + /** + * Result transforms registered by middleware during this activity call. + * Transforms run after the raw adapter result exists and before the final + * result is returned or streamed. Push multiple transforms to run them in + * registration order. + */ + resultTransforms?: Array> + /** + * Activity inputs captured for middleware that needs to transform or persist + * the result together with reconstructable request metadata. + */ + artifactInputs?: unknown } +/** Stable context handed to each {@link GenerationResultTransform}. */ +export interface GenerationResultTransformContext { + /** The activity call being transformed. */ + middleware: GenerationMiddlewareContext +} + +/** + * A transform middleware registers on `ctx.resultTransforms` to rewrite the raw + * adapter result before it is returned or streamed. Return a new result to + * replace it, or `undefined` to leave it unchanged. + */ +export type GenerationResultTransform = ( + result: TResult, + ctx: GenerationResultTransformContext, +) => TResult | undefined | Promise + // =========================== // Hook payloads // =========================== diff --git a/packages/ai/src/activities/stream-generation-result.ts b/packages/ai/src/activities/stream-generation-result.ts index 1bb530179..721b7a349 100644 --- a/packages/ai/src/activities/stream-generation-result.ts +++ b/packages/ai/src/activities/stream-generation-result.ts @@ -12,6 +12,19 @@ function createId(prefix: string): string { return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 9)}` } +/** + * Persisted artifact refs a middleware may have attached to the result. Read + * defensively: the result shape is activity-specific and `artifacts` is only + * present when generation persistence is wired with an artifact + blob store. + */ +function artifactsFromResult(result: unknown): Array | undefined { + if (typeof result !== 'object' || result === null) return undefined + const artifacts = (result as { artifacts?: unknown }).artifacts + return Array.isArray(artifacts) && artifacts.length > 0 + ? artifacts + : undefined +} + /** * Wrap a one-shot generation result as a StreamChunk async iterable. * @@ -23,7 +36,10 @@ function createId(prefix: string): string { * @returns An AsyncIterable of StreamChunks with RUN_STARTED, CUSTOM(generation:result), and RUN_FINISHED events on success, or RUN_STARTED and RUN_ERROR on failure */ export async function* streamGenerationResult( - generator: () => Promise, + generator: (resolved: { + runId: string + threadId: string + }) => Promise, options?: { runId?: string; threadId?: string }, ): AsyncIterable { const runId = options?.runId ?? createId('run') @@ -37,7 +53,19 @@ export async function* streamGenerationResult( } try { - const result = await generator() + const result = await generator({ runId, threadId }) + + // Emit persisted artifact refs (if a middleware attached any) before the + // result, so the client records them as the run streams. + const artifacts = artifactsFromResult(result) + if (artifacts) { + yield { + type: EventType.CUSTOM, + name: 'generation:artifacts', + value: artifacts, + timestamp: Date.now(), + } + } yield { type: EventType.CUSTOM, diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index c903e3e94..c263b85d3 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -202,6 +202,8 @@ export type { GenerationAbortInfo, GenerationErrorInfo, AnyGenerationMiddleware, + GenerationResultTransform, + GenerationResultTransformContext, } from './activities/middleware/index' // Capability primitives + middleware builder export { diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 6fea4bac8..935de84fe 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -2175,7 +2175,16 @@ export interface PersistedArtifactRef { mimeType: string size: number createdAt: string + /** The provider's original media URL that was fetched (may be expiring). */ externalUrl?: string + /** + * Durable app-origin URL that serves this artifact's persisted bytes (your + * `GET` route around `retrieveArtifact` / `retrieveBlob`). Stamped by + * `withGenerationPersistence`'s `artifactUrl` option, so clients render and + * restore durable media from your own origin rather than the provider's + * expiring link. + */ + url?: string source: { activity: PersistedArtifactActivity path: string diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index caec892a8..d818b7ebc 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -2075,6 +2075,10 @@ importers: version: 4.3.6 packages/ai-persistence: + dependencies: + '@tanstack/ai-utils': + specifier: workspace:* + version: link:../ai-utils devDependencies: '@tanstack/ai': specifier: workspace:* diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index 4abac986d..cdfaec2b3 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -14,6 +14,8 @@ import { Route as PersistenceDurabilityRouteImport } from './routes/persistence- import { Route as MiddlewareTestRouteImport } from './routes/middleware-test' import { Route as MarkdownCjkRouteImport } from './routes/markdown-cjk' import { Route as InterruptsTestRouteImport } from './routes/interrupts-test' +import { Route as GenerationPersistenceServerRouteImport } from './routes/generation-persistence-server' +import { Route as GenerationPersistenceRouteImport } from './routes/generation-persistence' import { Route as ForeignInterruptRouteImport } from './routes/foreign-interrupt' import { Route as DevtoolsToolsRouteImport } from './routes/devtools-tools' import { Route as DevtoolsStructuredRouteImport } from './routes/devtools-structured' @@ -52,6 +54,8 @@ import { Route as ApiMaxToolCallsWireRouteImport } from './routes/api.max-tool-c import { Route as ApiLazyToolsWireRouteImport } from './routes/api.lazy-tools-wire' import { Route as ApiInterruptsTestRouteImport } from './routes/api.interrupts-test' import { Route as ApiImageRouteImport } from './routes/api.image' +import { Route as ApiGenerationPersistenceServerRouteImport } from './routes/api.generation-persistence-server' +import { Route as ApiGenerationPersistenceRouteImport } from './routes/api.generation-persistence' import { Route as ApiForeignInterruptRouteImport } from './routes/api.foreign-interrupt' import { Route as ApiDurableDeliveryRouteImport } from './routes/api.durable-delivery' import { Route as ApiDevtoolsMemoryRouteImport } from './routes/api.devtools-memory' @@ -93,6 +97,17 @@ const InterruptsTestRoute = InterruptsTestRouteImport.update({ path: '/interrupts-test', getParentRoute: () => rootRouteImport, } as any) +const GenerationPersistenceServerRoute = + GenerationPersistenceServerRouteImport.update({ + id: '/generation-persistence-server', + path: '/generation-persistence-server', + getParentRoute: () => rootRouteImport, + } as any) +const GenerationPersistenceRoute = GenerationPersistenceRouteImport.update({ + id: '/generation-persistence', + path: '/generation-persistence', + getParentRoute: () => rootRouteImport, +} as any) const ForeignInterruptRoute = ForeignInterruptRouteImport.update({ id: '/foreign-interrupt', path: '/foreign-interrupt', @@ -288,6 +303,18 @@ const ApiImageRoute = ApiImageRouteImport.update({ path: '/api/image', getParentRoute: () => rootRouteImport, } as any) +const ApiGenerationPersistenceServerRoute = + ApiGenerationPersistenceServerRouteImport.update({ + id: '/api/generation-persistence-server', + path: '/api/generation-persistence-server', + getParentRoute: () => rootRouteImport, + } as any) +const ApiGenerationPersistenceRoute = + ApiGenerationPersistenceRouteImport.update({ + id: '/api/generation-persistence', + path: '/api/generation-persistence', + getParentRoute: () => rootRouteImport, + } as any) const ApiForeignInterruptRoute = ApiForeignInterruptRouteImport.update({ id: '/api/foreign-interrupt', path: '/api/foreign-interrupt', @@ -376,6 +403,8 @@ export interface FileRoutesByFullPath { '/devtools-structured': typeof DevtoolsStructuredRoute '/devtools-tools': typeof DevtoolsToolsRoute '/foreign-interrupt': typeof ForeignInterruptRoute + '/generation-persistence': typeof GenerationPersistenceRoute + '/generation-persistence-server': typeof GenerationPersistenceServerRoute '/interrupts-test': typeof InterruptsTestRoute '/markdown-cjk': typeof MarkdownCjkRoute '/middleware-test': typeof MiddlewareTestRoute @@ -391,6 +420,8 @@ export interface FileRoutesByFullPath { '/api/devtools-memory': typeof ApiDevtoolsMemoryRoute '/api/durable-delivery': typeof ApiDurableDeliveryRoute '/api/foreign-interrupt': typeof ApiForeignInterruptRoute + '/api/generation-persistence': typeof ApiGenerationPersistenceRoute + '/api/generation-persistence-server': typeof ApiGenerationPersistenceServerRoute '/api/image': typeof ApiImageRouteWithChildren '/api/interrupts-test': typeof ApiInterruptsTestRoute '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute @@ -436,6 +467,8 @@ export interface FileRoutesByTo { '/devtools-structured': typeof DevtoolsStructuredRoute '/devtools-tools': typeof DevtoolsToolsRoute '/foreign-interrupt': typeof ForeignInterruptRoute + '/generation-persistence': typeof GenerationPersistenceRoute + '/generation-persistence-server': typeof GenerationPersistenceServerRoute '/interrupts-test': typeof InterruptsTestRoute '/markdown-cjk': typeof MarkdownCjkRoute '/middleware-test': typeof MiddlewareTestRoute @@ -451,6 +484,8 @@ export interface FileRoutesByTo { '/api/devtools-memory': typeof ApiDevtoolsMemoryRoute '/api/durable-delivery': typeof ApiDurableDeliveryRoute '/api/foreign-interrupt': typeof ApiForeignInterruptRoute + '/api/generation-persistence': typeof ApiGenerationPersistenceRoute + '/api/generation-persistence-server': typeof ApiGenerationPersistenceServerRoute '/api/image': typeof ApiImageRouteWithChildren '/api/interrupts-test': typeof ApiInterruptsTestRoute '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute @@ -497,6 +532,8 @@ export interface FileRoutesById { '/devtools-structured': typeof DevtoolsStructuredRoute '/devtools-tools': typeof DevtoolsToolsRoute '/foreign-interrupt': typeof ForeignInterruptRoute + '/generation-persistence': typeof GenerationPersistenceRoute + '/generation-persistence-server': typeof GenerationPersistenceServerRoute '/interrupts-test': typeof InterruptsTestRoute '/markdown-cjk': typeof MarkdownCjkRoute '/middleware-test': typeof MiddlewareTestRoute @@ -512,6 +549,8 @@ export interface FileRoutesById { '/api/devtools-memory': typeof ApiDevtoolsMemoryRoute '/api/durable-delivery': typeof ApiDurableDeliveryRoute '/api/foreign-interrupt': typeof ApiForeignInterruptRoute + '/api/generation-persistence': typeof ApiGenerationPersistenceRoute + '/api/generation-persistence-server': typeof ApiGenerationPersistenceServerRoute '/api/image': typeof ApiImageRouteWithChildren '/api/interrupts-test': typeof ApiInterruptsTestRoute '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute @@ -559,6 +598,8 @@ export interface FileRouteTypes { | '/devtools-structured' | '/devtools-tools' | '/foreign-interrupt' + | '/generation-persistence' + | '/generation-persistence-server' | '/interrupts-test' | '/markdown-cjk' | '/middleware-test' @@ -574,6 +615,8 @@ export interface FileRouteTypes { | '/api/devtools-memory' | '/api/durable-delivery' | '/api/foreign-interrupt' + | '/api/generation-persistence' + | '/api/generation-persistence-server' | '/api/image' | '/api/interrupts-test' | '/api/lazy-tools-wire' @@ -619,6 +662,8 @@ export interface FileRouteTypes { | '/devtools-structured' | '/devtools-tools' | '/foreign-interrupt' + | '/generation-persistence' + | '/generation-persistence-server' | '/interrupts-test' | '/markdown-cjk' | '/middleware-test' @@ -634,6 +679,8 @@ export interface FileRouteTypes { | '/api/devtools-memory' | '/api/durable-delivery' | '/api/foreign-interrupt' + | '/api/generation-persistence' + | '/api/generation-persistence-server' | '/api/image' | '/api/interrupts-test' | '/api/lazy-tools-wire' @@ -679,6 +726,8 @@ export interface FileRouteTypes { | '/devtools-structured' | '/devtools-tools' | '/foreign-interrupt' + | '/generation-persistence' + | '/generation-persistence-server' | '/interrupts-test' | '/markdown-cjk' | '/middleware-test' @@ -694,6 +743,8 @@ export interface FileRouteTypes { | '/api/devtools-memory' | '/api/durable-delivery' | '/api/foreign-interrupt' + | '/api/generation-persistence' + | '/api/generation-persistence-server' | '/api/image' | '/api/interrupts-test' | '/api/lazy-tools-wire' @@ -740,6 +791,8 @@ export interface RootRouteChildren { DevtoolsStructuredRoute: typeof DevtoolsStructuredRoute DevtoolsToolsRoute: typeof DevtoolsToolsRoute ForeignInterruptRoute: typeof ForeignInterruptRoute + GenerationPersistenceRoute: typeof GenerationPersistenceRoute + GenerationPersistenceServerRoute: typeof GenerationPersistenceServerRoute InterruptsTestRoute: typeof InterruptsTestRoute MarkdownCjkRoute: typeof MarkdownCjkRoute MiddlewareTestRoute: typeof MiddlewareTestRoute @@ -755,6 +808,8 @@ export interface RootRouteChildren { ApiDevtoolsMemoryRoute: typeof ApiDevtoolsMemoryRoute ApiDurableDeliveryRoute: typeof ApiDurableDeliveryRoute ApiForeignInterruptRoute: typeof ApiForeignInterruptRoute + ApiGenerationPersistenceRoute: typeof ApiGenerationPersistenceRoute + ApiGenerationPersistenceServerRoute: typeof ApiGenerationPersistenceServerRoute ApiImageRoute: typeof ApiImageRouteWithChildren ApiInterruptsTestRoute: typeof ApiInterruptsTestRoute ApiLazyToolsWireRoute: typeof ApiLazyToolsWireRoute @@ -822,6 +877,20 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof InterruptsTestRouteImport parentRoute: typeof rootRouteImport } + '/generation-persistence-server': { + id: '/generation-persistence-server' + path: '/generation-persistence-server' + fullPath: '/generation-persistence-server' + preLoaderRoute: typeof GenerationPersistenceServerRouteImport + parentRoute: typeof rootRouteImport + } + '/generation-persistence': { + id: '/generation-persistence' + path: '/generation-persistence' + fullPath: '/generation-persistence' + preLoaderRoute: typeof GenerationPersistenceRouteImport + parentRoute: typeof rootRouteImport + } '/foreign-interrupt': { id: '/foreign-interrupt' path: '/foreign-interrupt' @@ -1088,6 +1157,20 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiImageRouteImport parentRoute: typeof rootRouteImport } + '/api/generation-persistence-server': { + id: '/api/generation-persistence-server' + path: '/api/generation-persistence-server' + fullPath: '/api/generation-persistence-server' + preLoaderRoute: typeof ApiGenerationPersistenceServerRouteImport + parentRoute: typeof rootRouteImport + } + '/api/generation-persistence': { + id: '/api/generation-persistence' + path: '/api/generation-persistence' + fullPath: '/api/generation-persistence' + preLoaderRoute: typeof ApiGenerationPersistenceRouteImport + parentRoute: typeof rootRouteImport + } '/api/foreign-interrupt': { id: '/api/foreign-interrupt' path: '/api/foreign-interrupt' @@ -1265,6 +1348,8 @@ const rootRouteChildren: RootRouteChildren = { DevtoolsStructuredRoute: DevtoolsStructuredRoute, DevtoolsToolsRoute: DevtoolsToolsRoute, ForeignInterruptRoute: ForeignInterruptRoute, + GenerationPersistenceRoute: GenerationPersistenceRoute, + GenerationPersistenceServerRoute: GenerationPersistenceServerRoute, InterruptsTestRoute: InterruptsTestRoute, MarkdownCjkRoute: MarkdownCjkRoute, MiddlewareTestRoute: MiddlewareTestRoute, @@ -1280,6 +1365,8 @@ const rootRouteChildren: RootRouteChildren = { ApiDevtoolsMemoryRoute: ApiDevtoolsMemoryRoute, ApiDurableDeliveryRoute: ApiDurableDeliveryRoute, ApiForeignInterruptRoute: ApiForeignInterruptRoute, + ApiGenerationPersistenceRoute: ApiGenerationPersistenceRoute, + ApiGenerationPersistenceServerRoute: ApiGenerationPersistenceServerRoute, ApiImageRoute: ApiImageRouteWithChildren, ApiInterruptsTestRoute: ApiInterruptsTestRoute, ApiLazyToolsWireRoute: ApiLazyToolsWireRoute, diff --git a/testing/e2e/src/routes/api.generation-persistence-server.ts b/testing/e2e/src/routes/api.generation-persistence-server.ts new file mode 100644 index 000000000..857d93f2d --- /dev/null +++ b/testing/e2e/src/routes/api.generation-persistence-server.ts @@ -0,0 +1,157 @@ +import { createFileRoute } from '@tanstack/react-router' +import { toServerSentEventsResponse } from '@tanstack/ai' +import type { StreamChunk } from '@tanstack/ai' + +/** + * Provider-free harness route for the SERVER-DRIVEN generation-persistence + * story (`persistence: true` + `threadId`). It is the server-authoritative + * counterpart to `api.generation-persistence.ts` (the client-driven variant). + * + * The client keeps NO local store; on mount it probes the GET below with a + * `?threadId=` query and restores transparently from the returned + * `reconstructGeneration`-shaped JSON (`{ resumeSnapshot, activeRun }`) into the + * normal `result` / `status` fields. To make the round-trip real, POST records + * the finished job in a module-level in-memory map keyed by `threadId`, and GET + * reads it back — so a full `page.reload()` (empty client storage) still + * restores the last run's status + a `result` whose image comes from the durable + * artifact URL, exactly the path `reconstructGeneration` serves in production. + * + * We hand-build the JSON here rather than pull `@tanstack/ai-persistence` into + * the e2e app (it is not a dependency), mirroring the `server-interrupt` + * scenario in `api.persistence-durability.ts`. `reconstructGeneration` itself is + * unit-tested in `@tanstack/ai-persistence`; this proves the CLIENT restore. + * + * Exempt from the aimock policy: this route streams a fixed AG-UI sequence and + * never reaches an LLM provider's HTTP layer, so there is nothing to mock. + */ + +// 1x1 transparent PNG — the live result's inline bytes, never persisted. +const TINY_PNG_B64 = + 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==' + +// Durable app-origin serve URL for the generated image (as +// `withGenerationPersistence`'s `artifactUrl` would stamp it). +const DURABLE_IMAGE_URL = '/durable/generation-server/image-1.png' + +// Server-authoritative record of the last completed generation per thread. In +// production this is a `GenerationJobStore` row; here a process-lifetime map is +// enough for the reload round-trip (the e2e server stays up across reloads). +const completedByThread = new Map>() + +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 +} + +function imageArtifact(threadId: string, runId: string) { + return { + role: 'output', + artifactId: 'artifact-image-1', + threadId, + runId, + name: 'image-1.png', + mimeType: 'image/png', + size: 68, + createdAt: new Date(0).toISOString(), + url: DURABLE_IMAGE_URL, + source: { + activity: 'image', + path: 'images.0', + provider: 'mock', + model: 'mock-image-model', + mediaType: 'image', + }, + } +} + +/** The metadata + durable artifact ref the server persists (never the bytes). */ +function persistedResult(threadId: string, runId: string) { + return { + id: 'image-1', + model: 'mock-image-model', + artifacts: [imageArtifact(threadId, runId)], + } +} + +function imageRun(threadId: string, runId: string): AsyncIterable { + return (async function* () { + yield { + type: 'RUN_STARTED', + threadId, + runId, + timestamp: Date.now(), + } as StreamChunk + yield { + type: 'CUSTOM', + name: 'generation:result', + value: { + id: 'image-1', + model: 'mock-image-model', + images: [{ b64Json: TINY_PNG_B64 }], + artifacts: [imageArtifact(threadId, runId)], + }, + threadId, + runId, + timestamp: Date.now(), + } as StreamChunk + yield { + type: 'RUN_FINISHED', + threadId, + runId, + timestamp: Date.now(), + } as StreamChunk + // The run finished: record the job the way `withGenerationPersistence` + // would, so the GET restore below can rebuild it after a reload. + completedByThread.set(threadId, persistedResult(threadId, runId)) + })() +} + +export const Route = createFileRoute('/api/generation-persistence-server')({ + server: { + handlers: { + POST: async ({ request }) => { + // The client sends an AG-UI RunAgentInput body carrying the hook's + // stable `threadId` (via runContext) — the same id the GET restore + // probe queries — plus the run id in the X-Run-Id header. + const body: unknown = await request.json() + const threadId = stringField(body, 'threadId') ?? 'generation-server' + const runId = + request.headers.get('X-Run-Id') ?? + stringField(body, 'runId') ?? + `run-${Date.now()}` + return toServerSentEventsResponse(imageRun(threadId, runId)) + }, + + // Server-authoritative restore: the `persistence: true` client's mount + // probe (`?threadId=`). Returns the same `{ resumeSnapshot, activeRun }` + // shape `reconstructGeneration` produces — a `complete` snapshot whose + // result carries the durable artifact ref once the thread has a recorded + // job, else the empty first-load answer. + GET: ({ request }) => { + const threadId = new URL(request.url).searchParams.get('threadId') ?? '' + const result = threadId ? completedByThread.get(threadId) : undefined + const body = result + ? { + resumeSnapshot: { + schemaVersion: 1, + resumeState: null, + status: 'complete', + activity: 'image', + result, + }, + activeRun: null, + } + : { resumeSnapshot: null, activeRun: null } + return new Response(JSON.stringify(body), { + headers: { + 'content-type': 'application/json', + 'cache-control': 'no-store', + }, + }) + }, + }, + }, +}) diff --git a/testing/e2e/src/routes/api.generation-persistence.ts b/testing/e2e/src/routes/api.generation-persistence.ts new file mode 100644 index 000000000..a1accaf07 --- /dev/null +++ b/testing/e2e/src/routes/api.generation-persistence.ts @@ -0,0 +1,96 @@ +import { createFileRoute } from '@tanstack/react-router' +import { toServerSentEventsResponse } from '@tanstack/ai' +import type { StreamChunk } from '@tanstack/ai' + +/** + * Provider-free harness route for the generation-persistence reload story. + * Streams a FIXED generation AG-UI sequence (started → progress → result → + * finished) instead of calling an image model, so the e2e is deterministic + * with nothing to mock. + * + * Exempt from the aimock policy: this route streams a fixed AG-UI sequence and + * never reaches an LLM provider's HTTP layer, so there is nothing to mock. + */ + +// 1x1 transparent PNG — the live result's inline bytes. Never persisted; the +// snapshot keeps only metadata + the durable artifact ref below. +const TINY_PNG_B64 = + 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==' + +// Durable app-origin serve URL for the generated image, the way +// `withGenerationPersistence`'s `artifactUrl` would stamp it. The reload path +// rebuilds `result.images` from this, so the restored image renders from our +// own origin instead of the (never-persisted) inline bytes. +const DURABLE_IMAGE_URL = '/durable/generation-persistence/image-1.png' + +function imageArtifact(threadId: string, runId: string) { + return { + role: 'output', + artifactId: 'artifact-image-1', + threadId, + runId, + name: 'image-1.png', + mimeType: 'image/png', + size: 68, + createdAt: new Date(0).toISOString(), + url: DURABLE_IMAGE_URL, + source: { + activity: 'image', + path: 'images.0', + provider: 'mock', + model: 'mock-image-model', + mediaType: 'image', + }, + } +} + +function imageRun(threadId: string, runId: string): AsyncIterable { + return (async function* () { + yield { + type: 'RUN_STARTED', + threadId, + runId, + timestamp: Date.now(), + } as StreamChunk + yield { + type: 'CUSTOM', + name: 'generation:progress', + value: { progress: 50, message: 'Painting pixels' }, + threadId, + runId, + timestamp: Date.now(), + } as StreamChunk + yield { + type: 'CUSTOM', + name: 'generation:result', + value: { + id: 'image-1', + model: 'mock-image-model', + images: [{ b64Json: TINY_PNG_B64 }], + artifacts: [imageArtifact(threadId, runId)], + }, + threadId, + runId, + timestamp: Date.now(), + } as StreamChunk + yield { + type: 'RUN_FINISHED', + threadId, + runId, + timestamp: Date.now(), + } as StreamChunk + })() +} + +export const Route = createFileRoute('/api/generation-persistence')({ + server: { + handlers: { + POST: async ({ request }) => { + const runId = request.headers.get('X-Run-Id') ?? `run-${Date.now()}` + const threadId = + request.headers.get('X-Thread-Id') ?? 'generation-persistence' + return toServerSentEventsResponse(imageRun(threadId, runId)) + }, + }, + }, +}) diff --git a/testing/e2e/src/routes/generation-persistence-server.tsx b/testing/e2e/src/routes/generation-persistence-server.tsx new file mode 100644 index 000000000..48bb96789 --- /dev/null +++ b/testing/e2e/src/routes/generation-persistence-server.tsx @@ -0,0 +1,65 @@ +import { useEffect, useState } from 'react' +import { createFileRoute } from '@tanstack/react-router' +import { fetchServerSentEvents, useGenerateImage } from '@tanstack/ai-react' + +/** + * Server-driven generation-persistence harness (server half). + * + * `persistence: true` + a stable `threadId` — the hook keeps NO local store. + * On mount it hydrates from the endpoint's GET (a `?threadId=` probe answered by + * `/api/generation-persistence-server`) and restores transparently into the + * normal fields, so a full `page.reload()` brings back `status: 'success'` and a + * `result` whose image renders from the durable serve URL — FROM THE SERVER, + * with `localStorage` empty. There is no `resumeSnapshot` field. + */ + +const THREAD_ID = 'generation-server-thread' +const connection = fetchServerSentEvents('/api/generation-persistence-server') + +export const Route = createFileRoute('/generation-persistence-server')({ + component: GenerationPersistenceServerPage, +}) + +function GenerationPersistenceServerPage() { + const image = useGenerateImage({ + id: 'generation-server', + threadId: THREAD_ID, + connection, + persistence: true, + }) + + // The page is SSR'd; the spec must not click the server-rendered button + // before React attaches handlers. This flag flips only after hydration. + const [hydrated, setHydrated] = useState(false) + useEffect(() => setHydrated(true), []) + + return ( +
+

Generation persistence (server-driven)

+ {hydrated ?
: null} + + +
{image.status}
+
{image.result?.id ?? 'none'}
+
{image.error?.message ?? 'none'}
+ + {image.result?.images.map((img, i) => ( + generated + ))} +
+ ) +} diff --git a/testing/e2e/src/routes/generation-persistence.tsx b/testing/e2e/src/routes/generation-persistence.tsx new file mode 100644 index 000000000..f5fd4b511 --- /dev/null +++ b/testing/e2e/src/routes/generation-persistence.tsx @@ -0,0 +1,75 @@ +import { useEffect, useState } from 'react' +import { createFileRoute } from '@tanstack/react-router' +import { + fetchServerSentEvents, + localStoragePersistence, + useGenerateImage, +} from '@tanstack/ai-react' + +/** + * Browser-refresh persistence harness for generation hooks (client-driven half). + * + * A bare `localStoragePersistence()` adapter stores the lightweight resume + * snapshot under `tanstack-ai:generation:`. A full `page.reload()` restores + * it transparently into the NORMAL hook fields — `status` becomes `success`, + * `error` returns, and `result` is rebuilt from the durable artifact URL (the + * generated bytes themselves are never written). There is no `resumeSnapshot` + * field; persistence is invisible, exactly like `useChat` restoring `messages`. + * The provider-free endpoint is `/api/generation-persistence`. + */ + +const snapshots = localStoragePersistence() +const connection = fetchServerSentEvents('/api/generation-persistence') + +export const Route = createFileRoute('/generation-persistence')({ + component: GenerationPersistencePage, +}) + +function GenerationPersistencePage() { + const image = useGenerateImage({ + id: 'generation-persistence', + connection, + persistence: snapshots, + }) + + // The page is SSR'd; the spec must not click the server-rendered button + // before React attaches handlers. This flag flips only after hydration. + const [hydrated, setHydrated] = useState(false) + useEffect(() => setHydrated(true), []) + + return ( +
+

Generation persistence

+ {hydrated ?
: null} + + + +
{image.status}
+
{image.result?.id ?? 'none'}
+
{image.error?.message ?? 'none'}
+
+ {image.resumeState ? image.resumeState.runId : 'none'} +
+ + {image.result?.images.map((img, i) => ( + generated + ))} +
+ ) +} diff --git a/testing/e2e/tests/generation-persistence-server.spec.ts b/testing/e2e/tests/generation-persistence-server.spec.ts new file mode 100644 index 000000000..958b40682 --- /dev/null +++ b/testing/e2e/tests/generation-persistence-server.spec.ts @@ -0,0 +1,52 @@ +import { expect, test } from '@playwright/test' + +/** + * Server-driven generation persistence (`persistence: true`). + * + * The counterpart to `generation-persistence.spec.ts` (client-driven). Here the + * hook keeps NO local store: as a run streams the server records the job, and on + * mount / after a full `page.reload()` the client restores the last run FROM THE + * SERVER via a `?threadId=` GET (the `reconstructGeneration` shape) straight into + * the normal `result` / `status` fields. The restored image renders from the + * durable serve URL, `localStorage` stays empty, and no run is auto-started. + * + * Provider-free: `/api/generation-persistence-server` streams a fixed AG-UI + * sequence and answers the restore GET from an in-memory job record (exempt from + * the aimock policy). + */ + +test.describe('generation persistence (server-driven)', () => { + test('records the run server-side and restores it after reload with no local store', async ({ + page, + }) => { + await page.goto('/generation-persistence-server') + await expect(page.getByTestId('hydration-marker')).toBeAttached() + // Empty first load: the server has no job for this thread yet. + await expect(page.getByTestId('status')).toHaveText('idle') + await expect(page.getByTestId('result-id')).toHaveText('none') + + await page.getByTestId('generate-button').click() + await expect(page.getByTestId('status')).toHaveText('success') + await expect(page.getByTestId('result-id')).toHaveText('image-1') + await expect(page.getByTestId('generated-image')).toHaveCount(1) + + // Server-driven mode writes NOTHING to browser storage — the record lives + // on the server. No localStorage key mentions the generation id. + const localKeys = await page.evaluate(() => + Object.keys(window.localStorage), + ) + expect(localKeys.some((k) => k.includes('generation-server'))).toBe(false) + + // Reload with empty client storage: the run is restored purely from the + // server, into the normal fields, and the image comes from the durable URL. + await page.reload() + await expect(page.getByTestId('hydration-marker')).toBeAttached() + await expect(page.getByTestId('status')).toHaveText('success') + await expect(page.getByTestId('result-id')).toHaveText('image-1') + await expect(page.getByTestId('generated-image')).toHaveCount(1) + await expect(page.getByTestId('generated-image')).toHaveAttribute( + 'src', + '/durable/generation-server/image-1.png', + ) + }) +}) diff --git a/testing/e2e/tests/generation-persistence.spec.ts b/testing/e2e/tests/generation-persistence.spec.ts new file mode 100644 index 000000000..ac94224cc --- /dev/null +++ b/testing/e2e/tests/generation-persistence.spec.ts @@ -0,0 +1,73 @@ +import { expect, test } from '@playwright/test' + +/** + * Generation resume persistence (browser refresh, client-driven). + * + * Proves the transparent restore wired by `localStoragePersistence` + + * `useGenerateImage({ persistence, id })`: as a run streams, the client writes a + * lightweight snapshot under `tanstack-ai:generation:` (metadata + the + * durable artifact URL, never the image bytes). A full `page.reload()` restores + * it straight into the NORMAL hook fields — `status` is `success`, and `result` + * is rebuilt so the image renders from the durable serve URL. There is no + * `resumeSnapshot` field; `reset()` clears everything. + * + * Provider-free: `/api/generation-persistence` streams a fixed AG-UI sequence + * (exempt from the aimock policy). + */ + +const STORAGE_KEY = 'tanstack-ai:generation:generation-persistence' + +test.describe('generation persistence (browser refresh)', () => { + test('restores status + result into the normal fields after reload, clears on reset', async ({ + page, + }) => { + await page.goto('/generation-persistence') + await expect(page.getByTestId('hydration-marker')).toBeAttached() + await expect(page.getByTestId('status')).toHaveText('idle') + await expect(page.getByTestId('result-id')).toHaveText('none') + + await page.getByTestId('generate-button').click() + await expect(page.getByTestId('status')).toHaveText('success') + await expect(page.getByTestId('result-id')).toHaveText('image-1') + await expect(page.getByTestId('generated-image')).toHaveCount(1) + + // The persisted record holds metadata + the durable artifact URL, never the + // inline image bytes. + const stored = await page.evaluate( + (key) => window.localStorage.getItem(key), + STORAGE_KEY, + ) + expect(stored).not.toBeNull() + expect(stored).toContain('"status":"complete"') + expect(stored).toContain('/durable/generation-persistence/image-1.png') + expect(stored).not.toContain('iVBOR') + + // Reload: the run restores transparently into status + result; nothing + // auto-runs (resumeState stays null) and the image renders from the durable + // serve URL, not the inline bytes. + await page.reload() + await expect(page.getByTestId('hydration-marker')).toBeAttached() + await expect(page.getByTestId('status')).toHaveText('success') + await expect(page.getByTestId('result-id')).toHaveText('image-1') + await expect(page.getByTestId('resume-state')).toHaveText('none') + await expect(page.getByTestId('generated-image')).toHaveCount(1) + await expect(page.getByTestId('generated-image')).toHaveAttribute( + 'src', + '/durable/generation-persistence/image-1.png', + ) + + // Reset clears the in-memory state and deletes the persisted record. + await page.getByTestId('reset-button').click() + await expect(page.getByTestId('status')).toHaveText('idle') + await expect(page.getByTestId('result-id')).toHaveText('none') + await expect + .poll(() => + page.evaluate((key) => window.localStorage.getItem(key), STORAGE_KEY), + ) + .toBeNull() + + await page.reload() + await expect(page.getByTestId('hydration-marker')).toBeAttached() + await expect(page.getByTestId('status')).toHaveText('idle') + }) +})