diff --git a/.changeset/generation-persistence.md b/.changeset/generation-persistence.md new file mode 100644 index 000000000..0ae75ceb5 --- /dev/null +++ b/.changeset/generation-persistence.md @@ -0,0 +1,22 @@ +--- +'@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: a lightweight client resume snapshot plus optional durable storage of the generated media bytes. + +**Media byte storage.** `withGenerationPersistence` now persists the generated bytes when the persistence backend provides both an `artifacts` (`ArtifactStore`) and a `blobs` (`BlobStore`) store. As an image/audio/speech/transcription run finishes, the middleware writes each artifact's bytes to the blob store (key `artifacts//`), records an `ArtifactRecord`, attaches `PersistedArtifactRef`s to the result, and emits a `generation:artifacts` event so the client records them. Extraction is customizable via `extractArtifacts` / `nameArtifact`. `memoryPersistence()` ships in-memory `artifacts`/`blobs` stores; the generation activities gained `threadId` / `runId` options and run their result through middleware result transforms. `@tanstack/ai-utils` adds `base64ToUint8Array`. To serve a stored artifact, `@tanstack/ai-persistence` exports `retrieveArtifact(persistence, id)` and `retrieveBlob(persistence, idOrRecord)` (plus `artifactBlobKey`). + +Add client-side generation persistence: a lightweight, read-only resume snapshot for media generation activities. + +Generation hooks (`useGenerateImage`, `useGenerateVideo`, `useGenerateAudio`, `useGenerateSpeech`, `useGeneration`, `useSummarize`, `useTranscription`, and their Solid/Vue/Svelte/Angular equivalents) now accept a `persistence` storage adapter and an `initialResumeSnapshot`, and expose `resumeSnapshot` / `resumeState` (plus observed `pendingArtifacts` / `resultArtifacts`). As a run streams, the client builds a `GenerationResumeSnapshot` — run identity, status, errors, and result metadata, but **never** the generated media bytes — and writes it to the adapter. The `persistence` option reuses the same `ChatStorageAdapter` contract as chat, so the shared `localStoragePersistence` / `sessionStoragePersistence` / `indexedDBPersistence` factories work for generations too. On reload the snapshot is surfaced for observability; it exposes no `resume()` action and never restarts provider work — generation still only begins when `generate(...)` is called. + +This pairs with the existing `withGenerationPersistence` server middleware, which records run status in the shared `RunStore`. diff --git a/docs/config.json b/docs/config.json index 9106ef679..c09231918 100644 --- a/docs/config.json +++ b/docs/config.json @@ -253,6 +253,11 @@ "addedAt": "2026-07-22", "updatedAt": "2026-07-23" }, + { + "label": "Generation Persistence", + "to": "persistence/generation-persistence", + "addedAt": "2026-07-23" + }, { "label": "Controls", "to": "persistence/controls", diff --git a/docs/persistence/generation-persistence.md b/docs/persistence/generation-persistence.md new file mode 100644 index 000000000..9cefed13a --- /dev/null +++ b/docs/persistence/generation-persistence.md @@ -0,0 +1,249 @@ +--- +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 track of. +Generation persistence keeps a small record of each run so your app can pick +things back up. + +It helps with three things: + +- **After a reload**, show what the last run was: its id, whether it finished, + and any error. This is a small read-only snapshot kept in the browser. +- **Keep the generated files.** On the server, save the generated bytes to your + own storage so they outlive the provider's expiring URLs. +- **While a run is still streaming**, let a dropped connection re-attach to it + instead of starting over. This reuses the same resumable streams the chat + client uses. + +## When to use it + +Use it when a run is long enough that a reload or a dropped connection actually +matters, or when you need to keep the output: video, batch images, long audio, +transcription of a big file. For a quick one-shot image you show and forget, you +can skip it. + +The browser snapshot never holds the generated bytes, only run identity, status, +errors, and references to the output. Storing the bytes themselves is a +server-side opt-in, shown in [Store the generated bytes](#store-the-generated-bytes). + +## Create the server endpoint + +Record each run in a store, and wrap the stream so a reload can re-attach to it: + +```ts group=generation-persistence +import { + generateImage, + generationParamsFromRequest, + memoryStream, + resumeServerSentEventsResponse, + toServerSentEventsResponse, +} from '@tanstack/ai' +import { openaiImage } from '@tanstack/ai-openai' +import { withGenerationPersistence } from '@tanstack/ai-persistence' +import { sqlitePersistence } from '@tanstack/ai-persistence-drizzle/sqlite' + +const persistence = sqlitePersistence({ + url: 'file:.tanstack-ai/generation.sqlite', + migrate: true, +}) + +export async function POST(request: Request) { + const durability = memoryStream(request) + const { input, threadId, runId } = + await generationParamsFromRequest('image', request) + + if (typeof input.prompt !== 'string') { + throw new Error('This endpoint accepts text image prompts only.') + } + + const stream = generateImage({ + ...(threadId ? { threadId } : {}), + ...(runId ? { runId } : {}), + adapter: openaiImage('gpt-image-2'), + prompt: input.prompt, + stream: true, + middleware: [withGenerationPersistence(persistence)], + }) + + // withGenerationPersistence records the run's status and result. + // durability records the stream so a reload can re-attach to it. + return toServerSentEventsResponse(stream, { + durability: { adapter: durability }, + }) +} + +export async function GET(request: Request) { + // Replays an in-flight run from the durability log. No provider call here. + return resumeServerSentEventsResponse({ adapter: memoryStream(request) }) +} +``` + +Use the matching request kind for audio, TTS, video, or transcription. +`withGenerationPersistence` records runs whenever a `runs` store is present. In +production, swap `memoryStream` for `durableStream` from +`@tanstack/ai-durable-stream`, where requests span processes. + +Keep run ids unique across chat and generation when they share a backend, +because `RunStore` is keyed by `runId`. + +## Store the generated bytes + +Provider URLs for generated media expire. To keep the output, give your +persistence backend an `artifacts` store (metadata) and a `blobs` store (the +bytes). When both are present, `withGenerationPersistence` writes each generated +file's bytes to the blob store, records an `ArtifactRecord`, and attaches +durable references to the result. `memoryPersistence()` ships both stores, so +it works out of the box; any backend that implements `ArtifactStore` and +`BlobStore` works the same way. + +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 by its id and streams the bytes +from your own origin. This is a plain file endpoint: it serves one stored file +and nothing more. It does not resume a run or rebuild a conversation. + +```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, threadId, runId } = + await generationParamsFromRequest('image', request) + + if (typeof input.prompt !== 'string') { + throw new Error('This endpoint accepts text image prompts only.') + } + + const stream = generateImage({ + ...(threadId ? { threadId } : {}), + ...(runId ? { runId } : {}), + adapter: openaiImage('gpt-image-2'), + prompt: input.prompt, + stream: true, + middleware: [withGenerationPersistence(persistence)], + }) + + return toServerSentEventsResponse(stream) +} + +// Serve a stored artifact's bytes by id. +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 `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. + +### Fetching artifacts later + +Two pieces cooperate, and neither one rebuilds a chat: + +- The run **remembers which files it produced**. Each reference carries an + `artifactId` and shows up on the generation hook as `resultArtifacts` (and + `pendingArtifacts` while streaming). The client snapshot keeps these across a + reload, and you can also store them in your own database. +- The **`GET` route turns an `artifactId` into bytes**. Point an `` `src`, + a download link, or a later request at `/api/generate/image?id=` + and it streams the stored file. + +So a page that generated an image yesterday can show it today: take the +`artifactId` you kept and hit the serve route. That is the whole loop, one id in, +one file out. Rebuilding a conversation's stored messages is a separate concern +handled by the chat `reconstructChat` helper, not by anything here. + +## Show the last run after a reload + +Pass a storage adapter as `persistence`. The client writes a snapshot as the run +streams and reads it back on load: + +```tsx +import { localStoragePersistence } from '@tanstack/ai-client' +import { fetchServerSentEvents, useGenerateImage } from '@tanstack/ai-react' + +const snapshots = localStoragePersistence({ keyPrefix: 'my-app:generation:' }) + +export function HeroImageGenerator() { + const image = useGenerateImage({ + id: 'hero-image', + connection: fetchServerSentEvents('/api/generate/image'), + persistence: snapshots, + }) + + return ( +
+ + + {image.resumeState ?

Last run: {image.resumeState.runId}

: null} +
+ ) +} +``` + +`image.resumeState` holds the last run's id once a run has streamed. The +snapshot is read-only, so it never re-runs the provider. A generation starts +only when you call `generate(...)`. + +## Reconnect to a run that is still streaming + +The server endpoint above already wires this: a `durability` adapter on +`toServerSentEventsResponse`, plus a `GET` handler that replays the run from the +log. 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. + +A full page reload is different. The hooks do not start a run on mount, so they +will not reconnect on their own. What survives the reload is the snapshot, which +holds the `runId`, so you can trigger a reconnect from it yourself. See +[Resumable Streams](../resumable-streams/overview) for the durability contract, +production adapters, and the one-time-side-effects note. + +## What the browser snapshot holds + +The browser snapshot never holds the generated bytes, only references to them. +Without an `artifacts` + `blobs` backend those references point at the provider's +own URL, which usually expires, so a snapshot opened much later can point at +media that is gone. Add the two stores (see +[Store the generated bytes](#store-the-generated-bytes)) to keep the files and +serve them from your own origin instead. diff --git a/examples/ts-react-chat/src/routes/generations.image.tsx b/examples/ts-react-chat/src/routes/generations.image.tsx index 044570498..fa67ef1bd 100644 --- a/examples/ts-react-chat/src/routes/generations.image.tsx +++ b/examples/ts-react-chat/src/routes/generations.image.tsx @@ -2,10 +2,20 @@ 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, read-only +// generation resume snapshot. Only run identity, status, errors, and result +// metadata are stored — never the generated image bytes. +const imageSnapshots = localStoragePersistence({ + keyPrefix: 'example:generation:', +}) + function StreamingImageGeneration() { const [prompt, setPrompt] = useState('') const [numberOfImages, setNumberOfImages] = useState(1) @@ -69,6 +79,42 @@ 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 ( +
+
+

+ Resume status: {hookReturn.status} +

+ {hookReturn.resumeState ? ( +

+ Last run:{' '} + {hookReturn.resumeState.runId} +

+ ) : ( +

No persisted run yet.

+ )} +
+ +
+ ) +} + function ImageGenerationUI({ prompt, setPrompt, @@ -170,9 +216,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 +261,16 @@ function ImageGenerationPage() { > Server Fn +
@@ -225,8 +281,10 @@ function ImageGenerationPage() { ) : mode === 'direct' ? ( - ) : ( + ) : mode === 'server-fn' ? ( + ) : ( + )} diff --git a/packages/ai-angular/src/inject-generate-audio.ts b/packages/ai-angular/src/inject-generate-audio.ts index 188a032e0..fbb57607c 100644 --- a/packages/ai-angular/src/inject-generate-audio.ts +++ b/packages/ai-angular/src/inject-generate-audio.ts @@ -10,13 +10,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' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for audio generation */ @@ -48,7 +57,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 +70,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 +116,19 @@ 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, + }) return { - generate: generate as (input: AudioGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, + ...generation, + generate: generation.generate as ( + input: AudioGenerateInput, + ) => Promise, } } diff --git a/packages/ai-angular/src/inject-generate-image.ts b/packages/ai-angular/src/inject-generate-image.ts index 3c3c2e4fe..16b08e8e6 100644 --- a/packages/ai-angular/src/inject-generate-image.ts +++ b/packages/ai-angular/src/inject-generate-image.ts @@ -6,7 +6,10 @@ 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, @@ -15,14 +18,14 @@ export type InjectGenerateImageOptions = Omit< 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 +41,18 @@ export function injectGenerateImage( hookName: 'injectGenerateImage', outputKind: 'image' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - injectGeneration({ - ...options, - devtools, - }) + const generation = injectGeneration< + ImageGenerateInput, + ImageGenerationResult, + TTransformed + >({ + ...options, + devtools, + }) return { - generate: generate as (input: ImageGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, + ...generation, + generate: generation.generate as ( + input: ImageGenerateInput, + ) => Promise, } } diff --git a/packages/ai-angular/src/inject-generate-speech.ts b/packages/ai-angular/src/inject-generate-speech.ts index ce3e08549..d78b804ec 100644 --- a/packages/ai-angular/src/inject-generate-speech.ts +++ b/packages/ai-angular/src/inject-generate-speech.ts @@ -6,7 +6,10 @@ 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, @@ -15,14 +18,15 @@ export type InjectGenerateSpeechOptions = Omit< 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,18 @@ export function injectGenerateSpeech( hookName: 'injectGenerateSpeech', outputKind: 'audio' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - injectGeneration({ - ...options, - devtools, - }) + const generation = injectGeneration< + SpeechGenerateInput, + TTSResult, + TTransformed + >({ + ...options, + devtools, + }) return { - generate: generate as (input: SpeechGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, + ...generation, + generate: generation.generate as ( + input: SpeechGenerateInput, + ) => Promise, } } diff --git a/packages/ai-angular/src/inject-generate-video.ts b/packages/ai-angular/src/inject-generate-video.ts index c493e39ee..ce1a62a6f 100644 --- a/packages/ai-angular/src/inject-generate-video.ts +++ b/packages/ai-angular/src/inject-generate-video.ts @@ -17,12 +17,17 @@ import type { ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPendingArtifact, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, VideoGenerateInput, VideoGenerateResult, VideoStatusInfo, } from '@tanstack/ai-client' import type { StreamChunk } from '@tanstack/ai' +import type { PersistedArtifactRef } from '@tanstack/ai/client' let nextId = 0 @@ -32,6 +37,8 @@ export interface InjectGenerateVideoOptions { id?: string body?: ReactiveOption> devtools?: AIDevtoolsDisplayOptions + persistence?: GenerationPersistence + initialResumeSnapshot?: GenerationResumeSnapshot onResult?: (result: VideoGenerateResult) => TOutput | null | void onError?: (error: Error) => void onProgress?: (progress: number, message?: string) => void @@ -50,6 +57,10 @@ export interface InjectGenerateVideoResult { status: Signal stop: () => void reset: () => void + resumeSnapshot: Signal + resumeState: Signal + pendingArtifacts: Signal> + resultArtifacts: Signal> } // `TTransformed` infers from the `onResult` return position so the callback @@ -79,6 +90,29 @@ export function injectGenerateVideo( const isLoading = signal(false) const error = signal(undefined) const status = signal('idle') + const resumeSnapshot = signal( + options.initialResumeSnapshot, + ) + const resumeState = signal( + options.initialResumeSnapshot?.resumeState ?? null, + ) + const pendingArtifacts = signal>( + options.initialResumeSnapshot?.pendingArtifacts ?? [], + ) + const resultArtifacts = signal>( + options.initialResumeSnapshot?.result?.artifacts ?? [], + ) + let disposed = false + + const setResumeSnapshotState = ( + snapshot: GenerationResumeSnapshot | undefined, + ) => { + if (disposed) return + resumeSnapshot.set(snapshot) + resumeState.set(snapshot?.resumeState ?? null) + pendingArtifacts.set(snapshot?.pendingArtifacts ?? []) + resultArtifacts.set(snapshot?.result?.artifacts ?? []) + } const bodySource = options.body !== undefined ? toReactive(options.body) : undefined @@ -86,6 +120,12 @@ export function injectGenerateVideo( const baseOptions = { id: clientId, ...(bodySource !== undefined && { body: bodySource() }), + ...(options.persistence !== undefined && { + persistence: options.persistence, + }), + ...(options.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: options.initialResumeSnapshot, + }), devtoolsBridgeFactory: createVideoDevtoolsBridge, devtools: { ...options.devtools, @@ -99,17 +139,40 @@ 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) + }, + onResumeSnapshotChange: setResumeSnapshotState, } let client: VideoGenerationClient @@ -132,14 +195,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 +226,9 @@ export function injectGenerateVideo( status: status.asReadonly(), stop: () => client.stop(), reset: () => client.reset(), + resumeSnapshot: resumeSnapshot.asReadonly(), + resumeState: resumeState.asReadonly(), + pendingArtifacts: pendingArtifacts.asReadonly(), + resultArtifacts: resultArtifacts.asReadonly(), } } diff --git a/packages/ai-angular/src/inject-generation.ts b/packages/ai-angular/src/inject-generation.ts index 177842e01..7aef1daa9 100644 --- a/packages/ai-angular/src/inject-generation.ts +++ b/packages/ai-angular/src/inject-generation.ts @@ -18,8 +18,13 @@ import type { GenerationClientOptions, GenerationClientState, GenerationFetcher, + GenerationPendingArtifact, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, } from '@tanstack/ai-client' +import type { PersistedArtifactRef } from '@tanstack/ai/client' import type { ReactiveOption } from './internal/to-reactive' let nextId = 0 @@ -35,6 +40,10 @@ export interface InjectGenerationOptions { body?: ReactiveOption> /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** Server-side lightweight generation state persistence. */ + persistence?: GenerationPersistence + /** Initial lightweight resume snapshot restored by the app (read-only state). */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when a result is received. Can optionally return a transformed value. * @@ -66,13 +75,21 @@ export interface InjectGenerationResult { stop: () => void /** Clear result, error, and return to idle */ reset: () => void + /** Lightweight generation resume snapshot, if one is available */ + resumeSnapshot: Signal + /** Observed run/cursor metadata from the snapshot (read-only state) */ + resumeState: Signal + /** Pending persisted artifact references observed during generation/replay */ + pendingArtifacts: Signal> + /** Final persisted artifact references observed from a replayed result */ + resultArtifacts: Signal> } // `TTransformed` infers from the `onResult` return position (a covariant // inference site that works even for an optional nested property), which types // the callback parameter as `TResult` and narrows `result`. Inferring the // whole callback as a defaulted type parameter instead collapses to the -// default, leaving the parameter `any` — a hard error under `strict`. See +// default, leaving the parameter `any` — a hard error under `strict`. See // issue #848. export function injectGeneration< TInput extends Record, @@ -97,6 +114,29 @@ export function injectGeneration< const isLoading = signal(false) const error = signal(undefined) const status = signal('idle') + const resumeSnapshot = signal( + options.initialResumeSnapshot, + ) + const resumeState = signal( + options.initialResumeSnapshot?.resumeState ?? null, + ) + const pendingArtifacts = signal>( + options.initialResumeSnapshot?.pendingArtifacts ?? [], + ) + const resultArtifacts = signal>( + options.initialResumeSnapshot?.result?.artifacts ?? [], + ) + let disposed = false + + const setResumeSnapshotState = ( + snapshot: GenerationResumeSnapshot | undefined, + ) => { + if (disposed) return + resumeSnapshot.set(snapshot) + resumeState.set(snapshot?.resumeState ?? null) + pendingArtifacts.set(snapshot?.pendingArtifacts ?? []) + resultArtifacts.set(snapshot?.result?.artifacts ?? []) + } const bodySource = options.body !== undefined ? toReactive(options.body) : undefined @@ -104,6 +144,12 @@ export function injectGeneration< const clientOptions: GenerationClientOptions = { id: clientId, ...(bodySource !== undefined && { body: bodySource() }), + ...(options.persistence !== undefined && { + persistence: options.persistence, + }), + ...(options.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: options.initialResumeSnapshot, + }), devtoolsBridgeFactory: createGenerationDevtoolsBridge, devtools: { ...options.devtools, @@ -116,13 +162,28 @@ 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) + }, + onResumeSnapshotChange: setResumeSnapshotState, } let client: GenerationClient @@ -145,14 +206,26 @@ 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 ( @@ -164,5 +237,9 @@ export function injectGeneration< status: status.asReadonly(), stop: () => client.stop(), reset: () => client.reset(), + resumeSnapshot: resumeSnapshot.asReadonly(), + resumeState: resumeState.asReadonly(), + pendingArtifacts: pendingArtifacts.asReadonly(), + resultArtifacts: resultArtifacts.asReadonly(), } } diff --git a/packages/ai-angular/src/inject-summarize.ts b/packages/ai-angular/src/inject-summarize.ts index fa41a76fa..db6b1c6e8 100644 --- a/packages/ai-angular/src/inject-summarize.ts +++ b/packages/ai-angular/src/inject-summarize.ts @@ -6,7 +6,10 @@ 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, @@ -15,14 +18,14 @@ export type InjectSummarizeOptions = Omit< 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 +41,18 @@ export function injectSummarize( hookName: 'injectSummarize', outputKind: 'text' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - injectGeneration( - { - ...options, - devtools, - }, - ) + const generation = injectGeneration< + SummarizeGenerateInput, + SummarizationResult, + TTransformed + >({ + ...options, + devtools, + }) return { - generate: generate as (input: SummarizeGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, + ...generation, + generate: generation.generate as ( + input: SummarizeGenerateInput, + ) => Promise, } } diff --git a/packages/ai-angular/src/inject-transcription.ts b/packages/ai-angular/src/inject-transcription.ts index ef2066ad2..b7bc832af 100644 --- a/packages/ai-angular/src/inject-transcription.ts +++ b/packages/ai-angular/src/inject-transcription.ts @@ -6,7 +6,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< @@ -19,14 +22,14 @@ export type InjectTranscriptionOptions = Omit< 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 +45,18 @@ export function injectTranscription( hookName: 'injectTranscription', outputKind: 'text' as const, } - const { generate, result, isLoading, error, status, stop, reset } = - injectGeneration< - TranscriptionGenerateInput, - TranscriptionResult, - TTransformed - >({ - ...options, - devtools, - }) + const generation = injectGeneration< + TranscriptionGenerateInput, + TranscriptionResult, + TTransformed + >({ + ...options, + devtools, + }) return { - generate: generate as (input: TranscriptionGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, + ...generation, + generate: generation.generate as ( + input: TranscriptionGenerateInput, + ) => Promise, } } diff --git a/packages/ai-angular/tests/inject-generation.test.ts b/packages/ai-angular/tests/inject-generation.test.ts index 1ae52833e..350dcdbb6 100644 --- a/packages/ai-angular/tests/inject-generation.test.ts +++ b/packages/ai-angular/tests/inject-generation.test.ts @@ -6,6 +6,13 @@ 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 type { StreamChunk } from '@tanstack/ai' +import type { + ConnectConnectionAdapter, + 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 +41,53 @@ 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(), + } +} + +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 } +} + describe('injectGeneration', () => { it('initializes idle with a fetcher and generates a result', async () => { const fetcher = vi.fn(async () => ({ value: 42 })) @@ -68,4 +119,53 @@ 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) + }) +}) + +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 persisted snapshot remains exposed as read-only state. + expect(result.resumeSnapshot()).toEqual(videoResumeSnapshot) + expect(result.resumeState()).toEqual(videoResumeSnapshot.resumeState) + }) }) diff --git a/packages/ai-client/src/generation-client.ts b/packages/ai-client/src/generation-client.ts index 38bc56a0b..abf77cdb1 100644 --- a/packages/ai-client/src/generation-client.ts +++ b/packages/ai-client/src/generation-client.ts @@ -1,4 +1,7 @@ -import { GENERATION_EVENTS } from './generation-types' +import { + GENERATION_EVENTS, + updateGenerationResumeSnapshot, +} from './generation-types' import { createNoOpGenerationDevtoolsBridge } from './devtools-noop' import { parseSSEResponse } from './sse-parser' import type { StreamChunk } from '@tanstack/ai/client' @@ -16,6 +19,8 @@ import type { GenerationClientOptions, GenerationClientState, GenerationFetcher, + GenerationResumeSnapshot, + GenerationPersistence, } from './generation-types' /** @@ -33,6 +38,9 @@ interface GenerationCallbacks { onLoadingChange?: ((isLoading: boolean) => void) | undefined onErrorChange?: ((error: Error | undefined) => void) | undefined onStatusChange?: ((status: GenerationClientState) => void) | undefined + onResumeSnapshotChange?: + | ((snapshot: GenerationResumeSnapshot) => void) + | undefined } /** @@ -81,6 +89,7 @@ export class GenerationClient< private readonly devtoolsMetadata: AIDevtoolsClientMetadata private readonly devtoolsBridge: GenerationDevtoolsBridge private readonly threadId: string + private readonly serverPersistence: GenerationPersistence | undefined private body: Record private result: TOutput | null = null private input: TInput | null = null @@ -88,9 +97,13 @@ 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 resumePersistenceError: Error | undefined = undefined private abortController: AbortController | null = null private readonly callbacksRef: GenerationCallbacks private devtoolsMounted = false + private disposed = false constructor( options: GenerationClientOptions & @@ -107,6 +120,8 @@ export class GenerationClient< this.connection = options.connection this.fetcher = options.fetcher this.body = options.body ?? {} + this.serverPersistence = options.persistence + this.resumeSnapshot = options.initialResumeSnapshot this.callbacksRef = { onResult: options.onResult, @@ -117,6 +132,7 @@ export class GenerationClient< onLoadingChange: options.onLoadingChange, onErrorChange: options.onErrorChange, onStatusChange: options.onStatusChange, + onResumeSnapshotChange: options.onResumeSnapshotChange, } this.devtoolsMetadata = this.createDevtoolsMetadata(options.devtools) @@ -159,6 +175,7 @@ export class GenerationClient< */ async generate(input: TInput): Promise { this.mountDevtools() + if (this.disposed) return if (this.isLoading) return this.input = input @@ -179,11 +196,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() } } else if (this.connection) { // Streaming adapter path @@ -194,7 +216,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', @@ -225,8 +247,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 +260,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 @@ -352,6 +378,7 @@ export class GenerationClient< } dispose(): void { + this.disposed = true this.stop() this.devtoolsBridge.dispose() this.devtoolsMounted = false @@ -373,10 +400,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 +524,58 @@ export class GenerationClient< runId, } } + + private observeResumeSnapshot(chunk: StreamChunk): void { + this.resumeSnapshot = updateGenerationResumeSnapshot( + this.resumeSnapshot, + chunk, + ) + this.callbacksRef.onResumeSnapshotChange?.(this.resumeSnapshot) + void this.persistResumeSnapshot(this.resumeSnapshot) + } + + private completePlainFetcherResumeSnapshot(): void { + if (!this.resumeSnapshot) { + return + } + this.resumeSnapshot = { + ...this.resumeSnapshot, + resumeState: null, + status: 'complete', + } + this.callbacksRef.onResumeSnapshotChange?.(this.resumeSnapshot) + void this.persistResumeSnapshot(this.resumeSnapshot) + } + + private async persistResumeSnapshot( + snapshot: GenerationResumeSnapshot, + ): Promise { + if (!this.serverPersistence) { + return + } + + this.resumeSnapshotPersistenceQueue = + this.resumeSnapshotPersistenceQueue.then( + () => this.writeResumeSnapshot(snapshot), + () => this.writeResumeSnapshot(snapshot), + ) + await this.resumeSnapshotPersistenceQueue + } + + private async writeResumeSnapshot( + snapshot: GenerationResumeSnapshot, + ): Promise { + try { + await this.serverPersistence?.setItem(this.threadId, snapshot) + } catch (error) { + this.resumePersistenceError = + error instanceof Error ? error : new Error(String(error)) + console.warn( + '[TanStack AI] Failed to persist generation resume snapshot', + error, + ) + } + } } function completeProgressValue( diff --git a/packages/ai-client/src/generation-types.ts b/packages/ai-client/src/generation-types.ts index b11e8ca2a..95084a9a9 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 { + MediaPrompt, + PersistedArtifactRef, + StreamChunk, +} from '@tanstack/ai/client' import type { 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,54 @@ export type InferGenerationOutput = TFn extends ( */ export type GenerationClientState = 'idle' | 'generating' | 'success' | 'error' +export type GenerationResumeStatus = 'idle' | 'running' | 'complete' | 'error' + +export interface GenerationResumeState { + threadId: string + runId: string +} + +export type GenerationPendingArtifact = PersistedArtifactRef + +export interface GenerationResultSnapshot { + id?: string + model?: string + status?: string + jobId?: string + expiresAt?: string + artifacts?: Array +} + +export interface GenerationErrorSnapshot { + message: string + code?: string +} + +export interface GenerationEventSnapshot { + type: StreamChunk['type'] + name?: string + timestamp?: number +} + +export interface GenerationResumeSnapshot { + 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 + // =========================== // Event Constants // =========================== @@ -70,6 +123,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 */ @@ -135,6 +190,25 @@ export interface GenerationClientOptions<_TInput, TResult, TOutput = TResult> { /** Metadata used to register this generation hook with TanStack AI Devtools */ devtools?: Partial + /** + * Initial lightweight resume snapshot restored by framework hooks. Contains + * only observed run metadata, errors, and persisted artifact refs. 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 + + /** + * Optional storage adapter for the lightweight generation resume snapshot. + * Accepts any {@link ChatStorageAdapter} — including the shared + * `localStoragePersistence` / `sessionStoragePersistence` / + * `indexedDBPersistence` factories. Generated media bytes are never written. + */ + persistence?: GenerationPersistence + /** * Factory that constructs the devtools bridge. Default is a no-op * factory; the real implementation lives in `@tanstack/ai-client/devtools`. @@ -166,6 +240,63 @@ 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 */ + onResumeSnapshotChange?: (snapshot: GenerationResumeSnapshot) => void +} + +export function updateGenerationResumeSnapshot( + previous: GenerationResumeSnapshot | null | undefined, + chunk: StreamChunk, +): GenerationResumeSnapshot { + const threadId = stringField(chunk, 'threadId') + const runId = stringField(chunk, 'runId') + const previousArtifacts = previous?.pendingArtifacts ?? [] + const next: GenerationResumeSnapshot = { + resumeState: previous?.resumeState ?? null, + status: previous?.status ?? 'idle', + ...(previous?.activity ? { activity: previous.activity } : {}), + ...(previousArtifacts.length > 0 + ? { pendingArtifacts: [...previousArtifacts] } + : {}), + ...(previous?.result ? { result: { ...previous.result } } : {}), + ...(previous?.error ? { error: { ...previous.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.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 } // =========================== @@ -200,6 +331,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 +461,214 @@ 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 } : {}), + } +} + +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') + if (id) snapshot.id = id + if (model) snapshot.model = model + if (status) snapshot.status = status + if (jobId) snapshot.jobId = jobId + 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 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 } : {}), + 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 + } +} + +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 cca4e4be6..5df5dbc6f 100644 --- a/packages/ai-client/src/index.ts +++ b/packages/ai-client/src/index.ts @@ -71,6 +71,14 @@ export type { InferGenerationOutput, InferGenerationOutputFromReturn, GenerationClientState, + GenerationResumeState, + GenerationResumeStatus, + GenerationResumeSnapshot, + GenerationPendingArtifact, + GenerationResultSnapshot, + GenerationErrorSnapshot, + GenerationEventSnapshot, + GenerationPersistence, GenerationClientOptions, GenerationFetcher, GenerationFetcherOptions, @@ -85,7 +93,10 @@ export type { SummarizeGenerateInput, VideoGenerateInput, } from './generation-types' -export { GENERATION_EVENTS } from './generation-types' +export { + GENERATION_EVENTS, + updateGenerationResumeSnapshot, +} from './generation-types' export { UnsupportedResponseStreamError } from './response-stream' export { clientTools, createChatClientOptions } from './types' // Web storage adapters for durable chat persistence (messages + resume snapshot) diff --git a/packages/ai-client/src/storage-adapters.ts b/packages/ai-client/src/storage-adapters.ts index 23130495e..a5c456940 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,15 @@ 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. `TValue` is value-agnostic + * by default, so `localStoragePersistence()` drops straight into any + * `persistence` option — chat or generation — with no type argument; the option + * you pass it to constrains the stored value. Pass a codec only for values JSON + * can't round-trip losslessly, and a type argument to lock the store's value + * type at the call site. */ -export function localStoragePersistence( +// oxlint-disable-next-line typescript/no-explicit-any -- value-agnostic default; the consuming `persistence` option constrains the value type +export function localStoragePersistence( options: WebStoragePersistenceOptions = {}, ): ChatStorageAdapter { return createWebStoragePersistence('localStorage', options) @@ -102,11 +105,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:` + * every other respect: value-agnostic default `TValue`, `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( +// oxlint-disable-next-line typescript/no-explicit-any -- value-agnostic default; the consuming `persistence` option constrains the value type +export function sessionStoragePersistence( options: WebStoragePersistenceOptions = {}, ): ChatStorageAdapter { return createWebStoragePersistence('sessionStorage', options) @@ -121,9 +125,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. `TValue` is value-agnostic by default. */ -export function indexedDBPersistence( +// oxlint-disable-next-line typescript/no-explicit-any -- value-agnostic default; the consuming `persistence` option constrains the value type +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..1d3433b71 100644 --- a/packages/ai-client/src/video-generation-client.ts +++ b/packages/ai-client/src/video-generation-client.ts @@ -1,4 +1,7 @@ -import { GENERATION_EVENTS } from './generation-types' +import { + GENERATION_EVENTS, + updateGenerationResumeSnapshot, +} from './generation-types' import { createNoOpVideoDevtoolsBridge } from './devtools-noop' import { parseSSEResponse } from './sse-parser' import type { StreamChunk } from '@tanstack/ai/client' @@ -15,6 +18,8 @@ import type { import type { GenerationClientState, GenerationFetcher, + GenerationResumeSnapshot, + GenerationPersistence, VideoGenerateInput, VideoGenerateResult, VideoGenerationClientOptions, @@ -42,6 +47,9 @@ interface VideoCallbacks { onStatusChange?: ((status: GenerationClientState) => void) | undefined onJobIdChange?: ((jobId: string | null) => void) | undefined onVideoStatusChange?: ((status: VideoStatusInfo | null) => void) | undefined + onResumeSnapshotChange?: + | ((snapshot: GenerationResumeSnapshot) => void) + | undefined } /** @@ -88,6 +96,7 @@ export class VideoGenerationClient { private readonly devtoolsMetadata: AIDevtoolsClientMetadata private readonly devtoolsBridge: VideoDevtoolsBridge private readonly threadId: string + private readonly serverPersistence: GenerationPersistence | undefined private body: Record private result: TOutput | null = null @@ -98,9 +107,13 @@ 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 resumePersistenceError: Error | undefined = undefined private abortController: AbortController | null = null private readonly callbacksRef: VideoCallbacks private devtoolsMounted = false + private disposed = false constructor( options: VideoGenerationClientOptions & @@ -117,6 +130,8 @@ export class VideoGenerationClient { this.connection = options.connection this.fetcher = options.fetcher this.body = options.body ?? {} + this.serverPersistence = options.persistence + this.resumeSnapshot = options.initialResumeSnapshot this.callbacksRef = { onResult: options.onResult, @@ -131,6 +146,7 @@ export class VideoGenerationClient { onStatusChange: options.onStatusChange, onJobIdChange: options.onJobIdChange, onVideoStatusChange: options.onVideoStatusChange, + onResumeSnapshotChange: options.onResumeSnapshotChange, } this.devtoolsMetadata = this.createDevtoolsMetadata(options.devtools) @@ -174,6 +190,7 @@ export class VideoGenerationClient { */ async generate(input: VideoGenerateInput): Promise { this.mountDevtools() + if (this.disposed) return if (this.isLoading) return this.input = input @@ -200,7 +217,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', @@ -226,8 +243,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 +266,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() } } @@ -262,13 +282,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 @@ -403,6 +425,7 @@ export class VideoGenerationClient { } dispose(): void { + this.disposed = true this.stop() this.devtoolsBridge.dispose() this.devtoolsMounted = false @@ -432,10 +455,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 +610,56 @@ export class VideoGenerationClient { runId, } } + + private observeResumeSnapshot(chunk: StreamChunk): void { + this.resumeSnapshot = updateGenerationResumeSnapshot( + this.resumeSnapshot, + chunk, + ) + this.callbacksRef.onResumeSnapshotChange?.(this.resumeSnapshot) + void this.persistResumeSnapshot(this.resumeSnapshot) + } + + private completePlainFetcherResumeSnapshot(): void { + if (!this.resumeSnapshot) { + return + } + this.resumeSnapshot = { + ...this.resumeSnapshot, + resumeState: null, + status: 'complete', + } + this.callbacksRef.onResumeSnapshotChange?.(this.resumeSnapshot) + void this.persistResumeSnapshot(this.resumeSnapshot) + } + + private async persistResumeSnapshot( + snapshot: GenerationResumeSnapshot, + ): Promise { + if (!this.serverPersistence) { + return + } + + this.resumeSnapshotPersistenceQueue = + this.resumeSnapshotPersistenceQueue.then( + () => this.writeResumeSnapshot(snapshot), + () => this.writeResumeSnapshot(snapshot), + ) + await this.resumeSnapshotPersistenceQueue + } + + private async writeResumeSnapshot( + snapshot: GenerationResumeSnapshot, + ): Promise { + try { + await this.serverPersistence?.setItem(this.threadId, snapshot) + } catch (error) { + this.resumePersistenceError = + error instanceof Error ? error : new Error(String(error)) + console.warn( + '[TanStack AI] Failed to persist generation resume snapshot', + error, + ) + } + } } diff --git a/packages/ai-client/tests/generation-client.test.ts b/packages/ai-client/tests/generation-client.test.ts index fb921a89a..99f1fcd41 100644 --- a/packages/ai-client/tests/generation-client.test.ts +++ b/packages/ai-client/tests/generation-client.test.ts @@ -1,8 +1,13 @@ import { describe, expect, it, vi } from 'vitest' import { EventType } from '@tanstack/ai/client' -import { GenerationClient, UnsupportedResponseStreamError } from '../src' +import { + GenerationClient, + UnsupportedResponseStreamError, + VideoGenerationClient, +} from '../src' import type { StreamChunk } from '@tanstack/ai/client' import type { ConnectConnectionAdapter } from '../src/connection-adapters' +import type { GenerationResumeSnapshot, GenerationPersistence } from '../src' // Helper to create a mock connect-based adapter from StreamChunks function createMockConnection( @@ -17,6 +22,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 +466,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 +1324,131 @@ 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' }, + }) + }) + }) + }) }) 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..74f8fad78 --- /dev/null +++ b/packages/ai-client/tests/generation-resume-state.test.ts @@ -0,0 +1,259 @@ +import { describe, expect, it } from 'vitest' +import { EventType } from '@tanstack/ai/client' +import { + GENERATION_EVENTS, + 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('omits non-durable and oversized result URLs from resume snapshots', () => { + const oversizedUrl = `https://example.com/${'x'.repeat(4096)}` + const unsafeUrls = [ + 'data:image/png;base64,raw-image-bytes', + 'blob:https://example.com/raw-image-bytes', + oversizedUrl, + ] + + for (const url of unsafeUrls) { + const snapshot = reduceChunks([ + { + type: EventType.CUSTOM, + name: GENERATION_EVENTS.RESULT, + value: { + id: 'result-1', + model: 'image-model', + url, + artifacts: [artifactRef], + }, + threadId: 'thread-1', + runId: 'run-1', + timestamp: 2, + }, + ]) + + expect(snapshot.result).toMatchObject({ + id: 'result-1', + model: 'image-model', + artifacts: [artifactRef], + }) + expect(snapshot.result).not.toHaveProperty('url') + expect(JSON.stringify(snapshot)).not.toContain('raw-image-bytes') + expect(JSON.stringify(snapshot)).not.toContain(oversizedUrl) + } + }) + + 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('does not model explicit stop as durable cancel state', () => { + const snapshot = reduceChunks([ + { + type: EventType.RUN_STARTED, + threadId: 'thread-1', + runId: 'run-1', + timestamp: 1, + }, + ]) + + expect(snapshot.status).toBe('running') + expect(snapshot).not.toHaveProperty('cancelled') + expect(snapshot).not.toHaveProperty('cancelEndpoint') + }) +}) diff --git a/packages/ai-event-client/src/index.ts b/packages/ai-event-client/src/index.ts index 0bd1fdc12..815c8a565 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 05f20cd0f..abd902dbd 100644 --- a/packages/ai-persistence/package.json +++ b/packages/ai-persistence/package.json @@ -46,6 +46,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/src/index.ts b/packages/ai-persistence/src/index.ts index 8b2f1226a..32398c30b 100644 --- a/packages/ai-persistence/src/index.ts +++ b/packages/ai-persistence/src/index.ts @@ -9,19 +9,44 @@ export type { InterruptStatus, InterruptStore, MetadataStore, + ArtifactRecord, + ArtifactStore, + BlobBody, + BlobRecord, + BlobObject, + BlobListPage, + BlobPutOptions, + BlobListOptions, + BlobStore, AIPersistence, AIPersistenceStores, AIPersistenceOverrides, ComposedAIPersistenceStores, } from './types' +// Core artifact wire types (re-exported for convenience) +export type { + PersistedArtifactActivity, + PersistedArtifactRef, + PersistedArtifactRole, +} from '@tanstack/ai' + // Middleware 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' +// Server helpers: retrieve a persisted generation artifact + its bytes +export { retrieveArtifact, retrieveBlob, artifactBlobKey } from './retrieve' + // Reference in-memory implementation export { memoryPersistence } from './memory' diff --git a/packages/ai-persistence/src/memory.ts b/packages/ai-persistence/src/memory.ts index da9d0d80c..1246b12ab 100644 --- a/packages/ai-persistence/src/memory.ts +++ b/packages/ai-persistence/src/memory.ts @@ -3,6 +3,13 @@ import { defineAIPersistence } from './types' import type { LockStore } from './locks' import type { ModelMessage } from '@tanstack/ai' import type { + ArtifactRecord, + ArtifactStore, + BlobBody, + BlobListOptions, + BlobObject, + BlobRecord, + BlobStore, InterruptRecord, InterruptStore, MessageStore, @@ -145,11 +152,207 @@ 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 interrupts: InterruptStore metadata: MetadataStore + artifacts: ArtifactStore + blobs: BlobStore locks: LockStore } @@ -159,6 +362,8 @@ export function memoryPersistence() { runs: new MemoryRunStore(), interrupts: new MemoryInterruptStore(), metadata: new MemoryMetadataStore(), + artifacts: new MemoryArtifactStore(), + blobs: new MemoryBlobStore(), locks: new InMemoryLockStore(), } return defineAIPersistence({ stores }) diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index 9c4f5de80..5ab8fb6ff 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, LocksCapability, @@ -9,7 +10,7 @@ import { } from './capabilities' import { validateChatPersistenceStores, - validatePersistenceStoreKeys, + validateGenerationPersistenceStores, } from './types' import type { AbortInfo, @@ -25,6 +26,9 @@ import type { GenerationMiddleware, GenerationMiddlewareContext, ModelMessage, + PersistedArtifactActivity, + PersistedArtifactRef, + PersistedArtifactRole, RunAgentResumeItem, StreamChunk, ToolApprovalResolution, @@ -33,9 +37,54 @@ import type { import type { AIPersistence, AIPersistenceStores, + ArtifactRecord, + BlobBody, InterruptRecord, RunStore, } from './types' +import { artifactBlobKey } from './retrieve' + +export interface WithPersistenceOptions { + extractArtifacts?: ( + input: GenerationArtifactExtractionInput, + ) => + | Array + | Promise> + nameArtifact?: (input: GenerationArtifactNameInput) => string +} + +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 @@ -239,6 +288,467 @@ 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, + }, + }) + } + + return refs +} + // --------------------------------------------------------------------------- // Shared store / feature plan // --------------------------------------------------------------------------- @@ -247,6 +757,7 @@ interface PersistencePlan { wantsMessages: boolean wantsInterrupts: boolean wantsLocks: boolean + wantsArtifactPersistence: boolean runs: AIPersistence['stores']['runs'] } @@ -255,6 +766,9 @@ function resolvePersistencePlan(persistence: AIPersistence): PersistencePlan { wantsMessages: persistence.stores.messages !== undefined, wantsInterrupts: persistence.stores.interrupts !== undefined, wantsLocks: persistence.stores.locks !== undefined, + wantsArtifactPersistence: + persistence.stores.artifacts !== undefined && + persistence.stores.blobs !== undefined, runs: persistence.stores.runs, } } @@ -284,9 +798,19 @@ type InvalidChatPersistence = ? StoreIsDefinitelyAbsent : false +type InvalidGenerationPersistence = + StoreIsDefinitelyPresent extends true + ? StoreIsDefinitelyAbsent + : StoreIsDefinitelyPresent extends true + ? StoreIsDefinitelyAbsent + : false + type ValidChatPersistence = InvalidChatPersistence extends true ? never : unknown +type ValidGenerationPersistence = + InvalidGenerationPersistence extends true ? never : unknown + async function createOrResumeRun( runs: RunStore | undefined, runId: string, @@ -503,17 +1027,21 @@ export function withPersistence(persistence: AIPersistence): ChatMiddleware { // --------------------------------------------------------------------------- /** - * Generation-only persistence middleware. Tracks run status (run records) for - * image, audio, TTS, video, and transcription activities. + * Generation-only persistence middleware. Tracks run status and optionally + * persists media artifacts/blobs for image, audio, TTS, video, and + * transcription activities. */ export function withGenerationPersistence( - persistence: AIPersistence, + persistence: AIPersistence & ValidGenerationPersistence, + opts?: WithPersistenceOptions, ): GenerationMiddleware export function withGenerationPersistence( persistence: AIPersistence, + opts?: WithPersistenceOptions, ): GenerationMiddleware { - validatePersistenceStoreKeys(persistence) - const { runs } = resolvePersistencePlan(persistence) + validateGenerationPersistenceStores(persistence) + const plan = resolvePersistencePlan(persistence) + const { wantsArtifactPersistence, runs } = plan // A generation activity has no thread or agent run: its only stable identity // is `requestId`, so the run record is keyed by it on both axes. @@ -522,6 +1050,21 @@ export function withGenerationPersistence( async onStart(ctx: GenerationMiddlewareContext) { await createOrResumeRun(runs, ctx.requestId, ctx.requestId) + if (!wantsArtifactPersistence) return + ctx.resultTransforms?.push(async (result) => { + const refs = await persistGenerationArtifacts( + persistence, + opts, + ctx, + result, + ) + if (refs.length === 0) return undefined + const existing = objectValue(result)?.artifacts + return { + ...(objectValue(result) ?? {}), + artifacts: [...(Array.isArray(existing) ? existing : []), ...refs], + } + }) }, async onFinish( 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 6213c4e36..a2f2b337b 100644 --- a/packages/ai-persistence/src/types.ts +++ b/packages/ai-persistence/src/types.ts @@ -25,9 +25,10 @@ import type { LockStore } from './locks' // // 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. @@ -196,12 +197,137 @@ export interface MetadataStore { delete: (scope: string, key: string) => Promise } +/** + * 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 +} + export interface AIPersistenceStores { messages?: MessageStore runs?: RunStore interrupts?: InterruptStore metadata?: MetadataStore locks?: LockStore + artifacts?: ArtifactStore + blobs?: BlobStore } export interface AIPersistence< @@ -316,6 +442,8 @@ const storeKeys = [ 'interrupts', 'metadata', 'locks', + 'artifacts', + 'blobs', ] satisfies Array const storeKeySet = new Set(storeKeys) @@ -341,6 +469,19 @@ export function validateChatPersistenceStores( } } +export function validateGenerationPersistenceStores( + persistence: AIPersistence, +): void { + validatePersistenceStoreKeys(persistence) + 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.', + ) + } +} + export function defineAIPersistence( persistence: AIPersistence>, ): AIPersistence { 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..f2642cecb --- /dev/null +++ b/packages/ai-persistence/tests/generation-artifacts.test.ts @@ -0,0 +1,496 @@ +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('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 run tracking without artifact stores', () => { + const full = memoryPersistence() + const persistence = defineAIPersistence({ + stores: { + runs: full.stores.runs, + }, + }) + + expect(() => withGenerationPersistence(persistence)).not.toThrow() + }) + + 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 560069ecf..245716673 100644 --- a/packages/ai-persistence/tests/memory.test.ts +++ b/packages/ai-persistence/tests/memory.test.ts @@ -13,6 +13,8 @@ describe('memoryPersistence', () => { it('exposes the complete state store set', () => { expect(Object.keys(memoryPersistence().stores).sort()).toEqual([ + 'artifacts', + 'blobs', 'interrupts', 'locks', 'messages', diff --git a/packages/ai-react/src/use-generate-audio.ts b/packages/ai-react/src/use-generate-audio.ts index a5f8223c0..98122c4ec 100644 --- a/packages/ai-react/src/use-generate-audio.ts +++ b/packages/ai-react/src/use-generate-audio.ts @@ -6,8 +6,13 @@ import type { ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPendingArtifact, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, } from '@tanstack/ai-client' +import type { PersistedArtifactRef } from '@tanstack/ai/client' /** * Options for the useGenerateAudio hook. @@ -25,6 +30,10 @@ export interface UseGenerateAudioOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** Server-side lightweight generation state persistence. */ + persistence?: GenerationPersistence + /** Initial lightweight resume snapshot restored by the app. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when audio is generated. Can optionally return a transformed value. * @@ -61,14 +70,22 @@ export interface UseGenerateAudioReturn { stop: () => void /** Clear result, error, and return to idle */ reset: () => void + /** Lightweight generation resume snapshot, if one is available */ + resumeSnapshot: GenerationResumeSnapshot | undefined + /** Current resumable run/cursor state, if one is available */ + resumeState: GenerationResumeState | null + /** Pending persisted artifact references observed during generation/replay */ + pendingArtifacts: Array + /** Final persisted artifact references observed from a replayed result */ + resultArtifacts: Array } /** * React hook for generating audio (music, sound effects) using AI models. * * Supports two transport modes: - * - **ConnectConnectionAdapter** — Streaming transport (SSE, HTTP stream, custom) - * - **Fetcher** — Direct async function call + * - **ConnectConnectionAdapter** — Streaming transport (SSE, HTTP stream, custom) + * - **Fetcher** — Direct async function call * * @example * ```tsx @@ -106,19 +123,19 @@ 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, + }) return { - generate: generate as (input: AudioGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, + ...generation, + generate: generation.generate as ( + input: AudioGenerateInput, + ) => Promise, } } diff --git a/packages/ai-react/src/use-generate-image.ts b/packages/ai-react/src/use-generate-image.ts index 078a792d1..452b1a838 100644 --- a/packages/ai-react/src/use-generate-image.ts +++ b/packages/ai-react/src/use-generate-image.ts @@ -5,9 +5,14 @@ import type { ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPendingArtifact, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, ImageGenerateInput, InferGenerationOutputFromReturn, } from '@tanstack/ai-client' +import type { PersistedArtifactRef } from '@tanstack/ai/client' /** * Options for the useGenerateImage hook. @@ -25,6 +30,10 @@ export interface UseGenerateImageOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** Server-side lightweight generation state persistence. */ + persistence?: GenerationPersistence + /** Initial lightweight resume snapshot restored by the app. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when images are generated. Can optionally return a transformed value. * @@ -61,14 +70,22 @@ export interface UseGenerateImageReturn { stop: () => void /** Clear result, error, and return to idle */ reset: () => void + /** Lightweight generation resume snapshot, if one is available */ + resumeSnapshot: GenerationResumeSnapshot | undefined + /** Current resumable run/cursor state, if one is available */ + resumeState: GenerationResumeState | null + /** Pending persisted artifact references observed during generation/replay */ + pendingArtifacts: Array + /** Final persisted artifact references observed from a replayed result */ + resultArtifacts: Array } /** * React hook for generating images using AI models. * * Supports two transport modes: - * - **ConnectConnectionAdapter** — Streaming transport (SSE, HTTP stream, custom) - * - **Fetcher** — Direct async function call + * - **ConnectConnectionAdapter** — Streaming transport (SSE, HTTP stream, custom) + * - **Fetcher** — Direct async function call * * @example * ```tsx @@ -108,19 +125,19 @@ 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, + }) return { - generate: generate as (input: ImageGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, + ...generation, + generate: generation.generate as ( + input: ImageGenerateInput, + ) => Promise, } } diff --git a/packages/ai-react/src/use-generate-speech.ts b/packages/ai-react/src/use-generate-speech.ts index b22199ff5..793e5b3a8 100644 --- a/packages/ai-react/src/use-generate-speech.ts +++ b/packages/ai-react/src/use-generate-speech.ts @@ -5,9 +5,14 @@ import type { ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPendingArtifact, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, SpeechGenerateInput, } from '@tanstack/ai-client' +import type { PersistedArtifactRef } from '@tanstack/ai/client' /** * Options for the useGenerateSpeech hook. @@ -25,6 +30,10 @@ export interface UseGenerateSpeechOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** Server-side lightweight generation state persistence. */ + persistence?: GenerationPersistence + /** Initial lightweight resume snapshot restored by the app. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when speech is generated. Can optionally return a transformed value. * @@ -61,6 +70,14 @@ export interface UseGenerateSpeechReturn { stop: () => void /** Clear result, error, and return to idle */ reset: () => void + /** Lightweight generation resume snapshot, if one is available */ + resumeSnapshot: GenerationResumeSnapshot | undefined + /** Current resumable run/cursor state, if one is available */ + resumeState: GenerationResumeState | null + /** Pending persisted artifact references observed during generation/replay */ + pendingArtifacts: Array + /** Final persisted artifact references observed from a replayed result */ + resultArtifacts: Array } /** @@ -102,19 +119,19 @@ 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, + ...generation, + generate: generation.generate as ( + input: SpeechGenerateInput, + ) => Promise, } } diff --git a/packages/ai-react/src/use-generate-video.ts b/packages/ai-react/src/use-generate-video.ts index 3aece73c5..a84af20a0 100644 --- a/packages/ai-react/src/use-generate-video.ts +++ b/packages/ai-react/src/use-generate-video.ts @@ -7,11 +7,16 @@ import type { ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPendingArtifact, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, VideoGenerateInput, VideoGenerateResult, VideoStatusInfo, } from '@tanstack/ai-client' +import type { PersistedArtifactRef } from '@tanstack/ai/client' /** * Options for the useGenerateVideo hook. @@ -27,6 +32,10 @@ export interface UseGenerateVideoOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** Server-side lightweight generation state persistence. */ + persistence?: GenerationPersistence + /** Initial lightweight resume snapshot restored by the app (read-only state). */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when video generation completes. Can optionally return a transformed value. * @@ -71,6 +80,14 @@ export interface UseGenerateVideoReturn { stop: () => void /** Clear all state and return to idle */ reset: () => void + /** Lightweight generation resume snapshot, if one is available */ + resumeSnapshot: GenerationResumeSnapshot | undefined + /** Current resumable run/cursor state, if one is available */ + resumeState: GenerationResumeState | null + /** Pending persisted artifact references observed during generation/replay */ + pendingArtifacts: Array + /** Final persisted artifact references observed from a replayed result */ + resultArtifacts: Array } /** @@ -127,6 +144,9 @@ export function useGenerateVideo( const [isLoading, setIsLoading] = useState(false) const [error, setError] = useState(undefined) const [status, setStatus] = useState('idle') + const [resumeSnapshot, setResumeSnapshot] = useState< + GenerationResumeSnapshot | undefined + >(options.initialResumeSnapshot) const optionsRef = useRef(options) optionsRef.current = options @@ -142,6 +162,10 @@ export function useGenerateVideo( const baseOptions = { id: clientId, body: opts.body, + ...(opts.persistence !== undefined && { persistence: opts.persistence }), + ...(opts.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: opts.initialResumeSnapshot, + }), devtoolsBridgeFactory: createVideoDevtoolsBridge, devtools: { ...opts.devtools, @@ -177,6 +201,7 @@ export function useGenerateVideo( onStatusChange: setStatus, onJobIdChange: setJobId, onVideoStatusChange: setVideoStatus, + onResumeSnapshotChange: setResumeSnapshot, } if (opts.connection) { @@ -206,7 +231,8 @@ 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 read-only for display. useEffect(() => { client.mountDevtools() @@ -240,5 +266,9 @@ export function useGenerateVideo( status, stop, reset, + resumeSnapshot, + resumeState: resumeSnapshot?.resumeState ?? null, + pendingArtifacts: resumeSnapshot?.pendingArtifacts ?? [], + resultArtifacts: resumeSnapshot?.result?.artifacts ?? [], } } diff --git a/packages/ai-react/src/use-generation.ts b/packages/ai-react/src/use-generation.ts index 6e5557d9b..ea54e21ff 100644 --- a/packages/ai-react/src/use-generation.ts +++ b/packages/ai-react/src/use-generation.ts @@ -8,8 +8,13 @@ import type { GenerationClientOptions, GenerationClientState, GenerationFetcher, + GenerationPendingArtifact, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, } from '@tanstack/ai-client' +import type { PersistedArtifactRef } from '@tanstack/ai/client' /** * Options for the useGeneration hook. @@ -31,6 +36,10 @@ export interface UseGenerationOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** Server-side lightweight generation state persistence. */ + persistence?: GenerationPersistence + /** Initial lightweight resume snapshot restored by the app (read-only state). */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when a result is received. Can optionally return a transformed value. * @@ -67,6 +76,14 @@ export interface UseGenerationReturn { stop: () => void /** Clear result, error, and return to idle */ reset: () => void + /** Lightweight generation state snapshot, if one is available */ + resumeSnapshot: GenerationResumeSnapshot | undefined + /** Observed run/cursor metadata from the snapshot (read-only state) */ + resumeState: GenerationResumeState | null + /** Pending persisted artifact references observed during generation/replay */ + pendingArtifacts: Array + /** Final persisted artifact references observed from a replayed result */ + resultArtifacts: Array } /** @@ -92,7 +109,7 @@ export interface UseGenerationReturn { // inference site that works even for an optional nested property), which types // the callback parameter as `TResult` and narrows `result`. Inferring the // whole callback as a defaulted type parameter instead collapses to the -// default, leaving the parameter `any` — a hard error under `strict`. See +// default, leaving the parameter `any` — a hard error under `strict`. See // issue #848. export function useGeneration< TInput extends Record, @@ -111,6 +128,9 @@ export function useGeneration< const [isLoading, setIsLoading] = useState(false) const [error, setError] = useState(undefined) const [status, setStatus] = useState('idle') + const [resumeSnapshot, setResumeSnapshot] = useState< + GenerationResumeSnapshot | undefined + >(options.initialResumeSnapshot) const optionsRef = useRef(options) optionsRef.current = options @@ -125,6 +145,10 @@ export function useGeneration< const clientOptions: GenerationClientOptions = { id: clientId, body: opts.body, + ...(opts.persistence !== undefined && { persistence: opts.persistence }), + ...(opts.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: opts.initialResumeSnapshot, + }), devtoolsBridgeFactory: createGenerationDevtoolsBridge, devtools: { hookName: 'useGeneration', @@ -150,6 +174,7 @@ export function useGeneration< onLoadingChange: setIsLoading, onErrorChange: setError, onStatusChange: setStatus, + onResumeSnapshotChange: setResumeSnapshot, } if (opts.connection) { @@ -179,7 +204,8 @@ 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 read-only for display. useEffect(() => { client.mountDevtools() @@ -211,5 +237,9 @@ export function useGeneration< status, stop, reset, + resumeSnapshot, + resumeState: resumeSnapshot?.resumeState ?? null, + pendingArtifacts: resumeSnapshot?.pendingArtifacts ?? [], + resultArtifacts: resumeSnapshot?.result?.artifacts ?? [], } } diff --git a/packages/ai-react/src/use-summarize.ts b/packages/ai-react/src/use-summarize.ts index 9de8eeb71..0deac64bf 100644 --- a/packages/ai-react/src/use-summarize.ts +++ b/packages/ai-react/src/use-summarize.ts @@ -5,9 +5,14 @@ import type { ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPendingArtifact, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, SummarizeGenerateInput, } from '@tanstack/ai-client' +import type { PersistedArtifactRef } from '@tanstack/ai/client' /** * Options for the useSummarize hook. @@ -25,6 +30,10 @@ export interface UseSummarizeOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** Server-side lightweight generation state persistence. */ + persistence?: GenerationPersistence + /** Initial lightweight resume snapshot restored by the app. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when summarization is complete. Can optionally return a transformed value. * @@ -61,6 +70,14 @@ export interface UseSummarizeReturn { stop: () => void /** Clear result, error, and return to idle */ reset: () => void + /** Lightweight generation resume snapshot, if one is available */ + resumeSnapshot: GenerationResumeSnapshot | undefined + /** Current resumable run/cursor state, if one is available */ + resumeState: GenerationResumeState | null + /** Pending persisted artifact references observed during generation/replay */ + pendingArtifacts: Array + /** Final persisted artifact references observed from a replayed result */ + resultArtifacts: Array } /** @@ -105,19 +122,19 @@ 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, + }) return { - generate: generate as (input: SummarizeGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, + ...generation, + generate: generation.generate as ( + input: SummarizeGenerateInput, + ) => Promise, } } diff --git a/packages/ai-react/src/use-transcription.ts b/packages/ai-react/src/use-transcription.ts index d0e01df8c..ed22ad3a5 100644 --- a/packages/ai-react/src/use-transcription.ts +++ b/packages/ai-react/src/use-transcription.ts @@ -5,9 +5,14 @@ import type { ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPendingArtifact, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, TranscriptionGenerateInput, } from '@tanstack/ai-client' +import type { PersistedArtifactRef } from '@tanstack/ai/client' /** * Options for the useTranscription hook. @@ -25,6 +30,10 @@ export interface UseTranscriptionOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** Server-side lightweight generation state persistence. */ + persistence?: GenerationPersistence + /** Initial lightweight resume snapshot restored by the app. */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when transcription is complete. Can optionally return a transformed value. * @@ -61,6 +70,14 @@ export interface UseTranscriptionReturn { stop: () => void /** Clear result, error, and return to idle */ reset: () => void + /** Lightweight generation resume snapshot, if one is available */ + resumeSnapshot: GenerationResumeSnapshot | undefined + /** Current resumable run/cursor state, if one is available */ + resumeState: GenerationResumeState | null + /** Pending persisted artifact references observed during generation/replay */ + pendingArtifacts: Array + /** Final persisted artifact references observed from a replayed result */ + resultArtifacts: Array } /** @@ -110,20 +127,16 @@ 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 }) return { - generate: generate as (input: TranscriptionGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, + ...generation, + generate: generation.generate as ( + input: TranscriptionGenerateInput, + ) => Promise, } } diff --git a/packages/ai-react/tests/use-generation.test.ts b/packages/ai-react/tests/use-generation.test.ts index f3540b520..67a8af348 100644 --- a/packages/ai-react/tests/use-generation.test.ts +++ b/packages/ai-react/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,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 +368,52 @@ 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() + // Persisted state is read-only for display; the client never reads it + // back to drive a resume, so getItem is not consulted on mount. + expect(persistence.getItem).not.toHaveBeenCalled() + expect(result.current.isLoading).toBe(false) + expect(result.current.status).toBe('idle') + // The persisted snapshot is still exposed as read-only state. + expect(result.current.resumeState).toEqual({ + threadId: 'thread-resume', + runId: 'run-resume', + }) + }) }) }) @@ -535,7 +673,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 +692,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 +819,39 @@ 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 persisted snapshot remains exposed as read-only state. + expect(result.current.resumeSnapshot).toEqual(videoResumeSnapshot) + expect(result.current.resumeState).toEqual(videoResumeSnapshot.resumeState) + }) }) describe('onResult transform', () => { diff --git a/packages/ai-solid/src/use-generate-audio.ts b/packages/ai-solid/src/use-generate-audio.ts index 94bfac297..fdca21fac 100644 --- a/packages/ai-solid/src/use-generate-audio.ts +++ b/packages/ai-solid/src/use-generate-audio.ts @@ -1,4 +1,8 @@ import { useGeneration } from './use-generation' +import type { + UseGenerationOptions, + UseGenerationReturn, +} from './use-generation' import type { AudioGenerationResult, StreamChunk } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -15,7 +19,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' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for audio generation */ @@ -47,7 +56,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 +69,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 +108,19 @@ 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, + }) return { - generate: generate as (input: AudioGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, + ...generation, + generate: generation.generate as ( + input: AudioGenerateInput, + ) => Promise, } } diff --git a/packages/ai-solid/src/use-generate-image.ts b/packages/ai-solid/src/use-generate-image.ts index b88902163..5a85ba6f3 100644 --- a/packages/ai-solid/src/use-generate-image.ts +++ b/packages/ai-solid/src/use-generate-image.ts @@ -1,4 +1,8 @@ import { useGeneration } from './use-generation' +import type { + UseGenerationOptions, + UseGenerationReturn, +} from './use-generation' import type { ImageGenerationResult, StreamChunk } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -15,7 +19,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' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for image generation */ @@ -47,7 +56,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 +69,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 +116,19 @@ 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, + }) return { - generate: generate as (input: ImageGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, + ...generation, + generate: generation.generate as ( + input: ImageGenerateInput, + ) => Promise, } } diff --git a/packages/ai-solid/src/use-generate-speech.ts b/packages/ai-solid/src/use-generate-speech.ts index b682e71a4..96fd5a5cc 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' | '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,19 @@ 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, + ...generation, + generate: generation.generate as ( + input: SpeechGenerateInput, + ) => Promise, } } diff --git a/packages/ai-solid/src/use-generate-video.ts b/packages/ai-solid/src/use-generate-video.ts index 4fd68b5de..d99323e5d 100644 --- a/packages/ai-solid/src/use-generate-video.ts +++ b/packages/ai-solid/src/use-generate-video.ts @@ -14,11 +14,16 @@ import type { ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPendingArtifact, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, VideoGenerateInput, VideoGenerateResult, VideoStatusInfo, } from '@tanstack/ai-client' +import type { PersistedArtifactRef } from '@tanstack/ai/client' import type { Accessor } from 'solid-js' /** @@ -37,6 +42,10 @@ export interface UseGenerateVideoOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** Server-side lightweight generation state persistence. */ + persistence?: GenerationPersistence + /** Initial lightweight resume snapshot restored by the app (read-only state). */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when video generation completes. Can optionally return a transformed value. * @@ -81,6 +90,14 @@ export interface UseGenerateVideoReturn { stop: () => void /** Clear all state and return to idle */ reset: () => void + /** Lightweight generation resume snapshot, if one is available */ + resumeSnapshot: Accessor + /** Observed run/cursor metadata from the snapshot (read-only state) */ + resumeState: Accessor + /** Pending persisted artifact references observed during generation/replay */ + pendingArtifacts: Accessor> + /** Final persisted artifact references observed from a replayed result */ + resultArtifacts: Accessor> } /** @@ -139,6 +156,10 @@ export function useGenerateVideo( const [isLoading, setIsLoading] = createSignal(false) const [error, setError] = createSignal(undefined) const [status, setStatus] = createSignal('idle') + const [resumeSnapshot, setResumeSnapshot] = createSignal< + GenerationResumeSnapshot | undefined + >(options.initialResumeSnapshot) + let disposed = false const client = createMemo(() => { // Conditional spread on `body`: VideoGenerationClientOptions.body @@ -146,6 +167,12 @@ export function useGenerateVideo( const baseOptions = { id: clientId, body: options.body, + ...(options.persistence !== undefined && { + persistence: options.persistence, + }), + ...(options.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: options.initialResumeSnapshot, + }), devtoolsBridgeFactory: createVideoDevtoolsBridge, devtools: { ...options.devtools, @@ -159,17 +186,44 @@ 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) + }, + onResumeSnapshotChange: ( + snapshot: GenerationResumeSnapshot | undefined, + ) => { + if (!disposed) setResumeSnapshot(snapshot) + }, } if (options.connection) { @@ -199,12 +253,15 @@ export function useGenerateVideo( }) }) + // Mount devtools only. Generation runs are never auto-started on mount — + // persisted state is read-only for display. onMount(() => { client().mountDevtools() }) // Cleanup on unmount: stop any in-flight requests and unregister devtools onCleanup(() => { + disposed = true client().dispose() }) @@ -230,5 +287,9 @@ export function useGenerateVideo( status, stop, reset, + resumeSnapshot, + resumeState: () => resumeSnapshot()?.resumeState ?? null, + pendingArtifacts: () => resumeSnapshot()?.pendingArtifacts ?? [], + resultArtifacts: () => resumeSnapshot()?.result?.artifacts ?? [], } } diff --git a/packages/ai-solid/src/use-generation.ts b/packages/ai-solid/src/use-generation.ts index 2db94a579..5098b91ec 100644 --- a/packages/ai-solid/src/use-generation.ts +++ b/packages/ai-solid/src/use-generation.ts @@ -15,8 +15,13 @@ import type { GenerationClientOptions, GenerationClientState, GenerationFetcher, + GenerationPendingArtifact, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, } from '@tanstack/ai-client' +import type { PersistedArtifactRef } from '@tanstack/ai/client' import type { Accessor } from 'solid-js' /** @@ -39,6 +44,10 @@ export interface UseGenerationOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** Server-side lightweight generation state persistence. */ + persistence?: GenerationPersistence + /** Initial lightweight resume snapshot restored by the app (read-only state). */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when a result is received. Can optionally return a transformed value. * @@ -75,6 +84,14 @@ export interface UseGenerationReturn { stop: () => void /** Clear result, error, and return to idle */ reset: () => void + /** Lightweight generation resume snapshot, if one is available */ + resumeSnapshot: Accessor + /** Observed run/cursor metadata from the snapshot (read-only state) */ + resumeState: Accessor + /** Pending persisted artifact references observed during generation/replay */ + pendingArtifacts: Accessor> + /** Final persisted artifact references observed from a replayed result */ + resultArtifacts: Accessor> } /** @@ -101,7 +118,7 @@ export interface UseGenerationReturn { // inference site that works even for an optional nested property), which types // the callback parameter as `TResult` and narrows `result`. Inferring the // whole callback as a defaulted type parameter instead collapses to the -// default, leaving the parameter `any` — a hard error under `strict`. See +// default, leaving the parameter `any` — a hard error under `strict`. See // issue #848. export function useGeneration< TInput extends Record, @@ -120,6 +137,10 @@ export function useGeneration< const [isLoading, setIsLoading] = createSignal(false) const [error, setError] = createSignal(undefined) const [status, setStatus] = createSignal('idle') + const [resumeSnapshot, setResumeSnapshot] = createSignal< + GenerationResumeSnapshot | undefined + >(options.initialResumeSnapshot) + let disposed = false const client = createMemo(() => { // Conditional spread on `body`: `GenerationClientOptions.body` is a @@ -128,6 +149,12 @@ export function useGeneration< const clientOptions: GenerationClientOptions = { id: clientId, body: options.body, + ...(options.persistence !== undefined && { + persistence: options.persistence, + }), + ...(options.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: options.initialResumeSnapshot, + }), devtoolsBridgeFactory: createGenerationDevtoolsBridge, devtools: { ...options.devtools, @@ -140,13 +167,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) + }, + onResumeSnapshotChange: (snapshot) => { + if (!disposed) setResumeSnapshot(snapshot) + }, } if (options.connection) { @@ -176,12 +220,15 @@ export function useGeneration< }) }) + // Mount devtools only. Generation runs are never auto-started on mount — + // persisted state is read-only for display. onMount(() => { client().mountDevtools() }) // Cleanup on unmount: stop any in-flight requests and unregister devtools onCleanup(() => { + disposed = true client().dispose() }) @@ -205,5 +252,9 @@ export function useGeneration< status, stop, reset, + resumeSnapshot, + resumeState: () => resumeSnapshot()?.resumeState ?? null, + pendingArtifacts: () => resumeSnapshot()?.pendingArtifacts ?? [], + resultArtifacts: () => resumeSnapshot()?.result?.artifacts ?? [], } } diff --git a/packages/ai-solid/src/use-summarize.ts b/packages/ai-solid/src/use-summarize.ts index b6a9eefc6..804d7cd6d 100644 --- a/packages/ai-solid/src/use-summarize.ts +++ b/packages/ai-solid/src/use-summarize.ts @@ -1,4 +1,8 @@ import { useGeneration } from './use-generation' +import type { + UseGenerationOptions, + UseGenerationReturn, +} from './use-generation' import type { StreamChunk, SummarizationResult } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -15,7 +19,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' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for summarization */ @@ -47,7 +56,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 +70,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 +114,19 @@ 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, + }) return { - generate: generate as (input: SummarizeGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, + ...generation, + generate: generation.generate as ( + input: SummarizeGenerateInput, + ) => Promise, } } diff --git a/packages/ai-solid/src/use-transcription.ts b/packages/ai-solid/src/use-transcription.ts index e3e19ee7d..fadb5aabf 100644 --- a/packages/ai-solid/src/use-transcription.ts +++ b/packages/ai-solid/src/use-transcription.ts @@ -1,4 +1,8 @@ import { useGeneration } from './use-generation' +import type { + UseGenerationOptions, + UseGenerationReturn, +} from './use-generation' import type { StreamChunk, TranscriptionResult } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -15,7 +19,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' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for transcription */ @@ -47,7 +60,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 +73,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 +123,16 @@ 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 }) return { - generate: generate as (input: TranscriptionGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, + ...generation, + generate: generation.generate as ( + input: TranscriptionGenerateInput, + ) => Promise, } } diff --git a/packages/ai-solid/tests/use-generation.test.ts b/packages/ai-solid/tests/use-generation.test.ts index 562268aa0..605bfe45c 100644 --- a/packages/ai-solid/tests/use-generation.test.ts +++ b/packages/ai-solid/tests/use-generation.test.ts @@ -10,6 +10,12 @@ import { useGenerateVideo } from '../src/use-generate-video' import { createMockConnectionAdapter } from './test-utils' import type { 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 +77,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 +869,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 +911,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 +929,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 +1134,38 @@ 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() + expect(persistence.getItem).not.toHaveBeenCalled() + expect(result.isLoading()).toBe(false) + expect(result.status()).toBe('idle') + // The persisted snapshot remains exposed as read-only state. + expect(result.resumeSnapshot()).toEqual(videoResumeSnapshot) + expect(result.resumeState()).toEqual(videoResumeSnapshot.resumeState) + }) }) describe('error handling', () => { diff --git a/packages/ai-svelte/src/create-generate-audio.svelte.ts b/packages/ai-svelte/src/create-generate-audio.svelte.ts index 5838ee6f2..87ded3f2a 100644 --- a/packages/ai-svelte/src/create-generate-audio.svelte.ts +++ b/packages/ai-svelte/src/create-generate-audio.svelte.ts @@ -1,4 +1,8 @@ import { createGeneration } from './create-generation.svelte' +import type { + CreateGenerationOptions, + CreateGenerationReturn, +} from './create-generation.svelte' import type { AudioGenerationResult, StreamChunk } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -14,7 +18,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' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for audio generation */ @@ -46,7 +55,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 +68,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 } /** @@ -130,5 +135,18 @@ export function createGenerateAudio( stop: gen.stop, reset: gen.reset, updateBody: gen.updateBody, + dispose: gen.dispose, + get resumeSnapshot() { + return gen.resumeSnapshot + }, + get resumeState() { + return gen.resumeState + }, + get pendingArtifacts() { + return gen.pendingArtifacts + }, + get resultArtifacts() { + return gen.resultArtifacts + }, } } diff --git a/packages/ai-svelte/src/create-generate-image.svelte.ts b/packages/ai-svelte/src/create-generate-image.svelte.ts index 0909ddca2..e42b70d55 100644 --- a/packages/ai-svelte/src/create-generate-image.svelte.ts +++ b/packages/ai-svelte/src/create-generate-image.svelte.ts @@ -1,4 +1,8 @@ import { createGeneration } from './create-generation.svelte' +import type { + CreateGenerationOptions, + CreateGenerationReturn, +} from './create-generation.svelte' import type { ImageGenerationResult, StreamChunk } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -14,7 +18,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' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for image generation */ @@ -46,7 +55,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 +68,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 } /** @@ -139,5 +144,18 @@ export function createGenerateImage( stop: gen.stop, reset: gen.reset, updateBody: gen.updateBody, + dispose: gen.dispose, + get resumeSnapshot() { + return gen.resumeSnapshot + }, + get resumeState() { + return gen.resumeState + }, + get pendingArtifacts() { + return gen.pendingArtifacts + }, + get resultArtifacts() { + return gen.resultArtifacts + }, } } diff --git a/packages/ai-svelte/src/create-generate-speech.svelte.ts b/packages/ai-svelte/src/create-generate-speech.svelte.ts index 20ccafbf9..90daf58e4 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' | '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 } /** @@ -126,5 +130,18 @@ export function createGenerateSpeech( stop: gen.stop, reset: gen.reset, updateBody: gen.updateBody, + dispose: gen.dispose, + get resumeSnapshot() { + return gen.resumeSnapshot + }, + get resumeState() { + return gen.resumeState + }, + get pendingArtifacts() { + return gen.pendingArtifacts + }, + get resultArtifacts() { + return gen.resultArtifacts + }, } } diff --git a/packages/ai-svelte/src/create-generate-video.svelte.ts b/packages/ai-svelte/src/create-generate-video.svelte.ts index 05f2d787b..1dd954c75 100644 --- a/packages/ai-svelte/src/create-generate-video.svelte.ts +++ b/packages/ai-svelte/src/create-generate-video.svelte.ts @@ -6,11 +6,16 @@ import type { ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPendingArtifact, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, VideoGenerateInput, VideoGenerateResult, VideoStatusInfo, } from '@tanstack/ai-client' +import type { PersistedArtifactRef } from '@tanstack/ai/client' /** * Options for the createGenerateVideo function. @@ -28,6 +33,10 @@ export interface CreateGenerateVideoOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** Server-side lightweight generation state persistence. */ + persistence?: GenerationPersistence + /** Initial lightweight resume snapshot restored by the app (read-only state). */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when video generation completes. Can optionally return a transformed value. * @@ -76,6 +85,14 @@ export interface CreateGenerateVideoReturn { dispose: () => void /** Update additional body parameters */ updateBody: (body: Record) => void + /** Lightweight generation resume snapshot, if one is available */ + readonly resumeSnapshot: GenerationResumeSnapshot | undefined + /** Observed run/cursor metadata from the snapshot (read-only state) */ + readonly resumeState: GenerationResumeState | null + /** Pending persisted artifact references observed during generation/replay */ + readonly pendingArtifacts: Array + /** Final persisted artifact references observed from a replayed result */ + readonly resultArtifacts: Array } /** @@ -133,6 +150,17 @@ export function createGenerateVideo( let isLoading = $state(false) let error = $state(undefined) let status = $state('idle') + let resumeSnapshot = $state( + options.initialResumeSnapshot, + ) + let disposed = false + + const setResumeSnapshotState = ( + snapshot: GenerationResumeSnapshot | undefined, + ) => { + if (disposed) return + resumeSnapshot = snapshot + } // `body` uses a conditional spread because `VideoGenerationClientOptions.body` // is declared `body?: Record` (absent vs. present) under @@ -141,6 +169,12 @@ export function createGenerateVideo( const baseOptions = { id: clientId, body: options.body, + ...(options.persistence !== undefined && { + persistence: options.persistence, + }), + ...(options.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: options.initialResumeSnapshot, + }), devtoolsBridgeFactory: createVideoDevtoolsBridge, devtools: { ...options.devtools, @@ -154,29 +188,46 @@ 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 }, + onResumeSnapshotChange: setResumeSnapshotState, } 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. @@ -217,6 +270,7 @@ export function createGenerateVideo( } const dispose = () => { + disposed = true client.dispose() } @@ -248,5 +302,17 @@ export function createGenerateVideo( reset, dispose, updateBody, + get resumeSnapshot() { + return resumeSnapshot + }, + get resumeState() { + return resumeSnapshot?.resumeState ?? null + }, + get pendingArtifacts() { + return resumeSnapshot?.pendingArtifacts ?? [] + }, + get resultArtifacts() { + return resumeSnapshot?.result?.artifacts ?? [] + }, } } diff --git a/packages/ai-svelte/src/create-generation.svelte.ts b/packages/ai-svelte/src/create-generation.svelte.ts index a731ac398..24e36e7c4 100644 --- a/packages/ai-svelte/src/create-generation.svelte.ts +++ b/packages/ai-svelte/src/create-generation.svelte.ts @@ -7,8 +7,13 @@ import type { GenerationClientOptions, GenerationClientState, GenerationFetcher, + GenerationPendingArtifact, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, } from '@tanstack/ai-client' +import type { PersistedArtifactRef } from '@tanstack/ai/client' /** * Options for the createGeneration function. @@ -30,6 +35,10 @@ export interface CreateGenerationOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** Server-side lightweight generation state persistence. */ + persistence?: GenerationPersistence + /** Initial lightweight resume snapshot restored by the app (read-only state). */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when a result is received. Can optionally return a transformed value. * @@ -70,6 +79,14 @@ export interface CreateGenerationReturn { dispose: () => void /** Update additional body parameters */ updateBody: (body: Record) => void + /** Lightweight generation resume snapshot, if one is available */ + readonly resumeSnapshot: GenerationResumeSnapshot | undefined + /** Observed run/cursor metadata from the snapshot (read-only state) */ + readonly resumeState: GenerationResumeState | null + /** Pending persisted artifact references observed during generation/replay */ + readonly pendingArtifacts: Array + /** Final persisted artifact references observed from a replayed result */ + readonly resultArtifacts: Array } /** @@ -107,7 +124,7 @@ export interface CreateGenerationReturn { // inference site that works even for an optional nested property), which types // the callback parameter as `TResult` and narrows `result`. Inferring the // whole callback as a defaulted type parameter instead collapses to the -// default, leaving the parameter `any` — a hard error under `strict`. See +// default, leaving the parameter `any` — a hard error under `strict`. See // issue #848. export function createGeneration< TInput extends Record, @@ -129,15 +146,32 @@ export function createGeneration< let isLoading = $state(false) let error = $state(undefined) let status = $state('idle') + let resumeSnapshot = $state( + options.initialResumeSnapshot, + ) + let disposed = false + + const setResumeSnapshotState = ( + snapshot: GenerationResumeSnapshot | undefined, + ) => { + if (disposed) return + resumeSnapshot = snapshot + } // `body` uses a conditional spread because `GenerationClientOptions.body` // is declared `body?: Record` (absent vs. present) under // `exactOptionalPropertyTypes`. Assigning `undefined` directly would be - // rejected — the optional caller `options.body` may be undefined, in which + // rejected — the optional caller `options.body` may be undefined, in which // case we want the key to be absent. const clientOptions: GenerationClientOptions = { id: clientId, body: options.body, + ...(options.persistence !== undefined && { + persistence: options.persistence, + }), + ...(options.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: options.initialResumeSnapshot, + }), devtoolsBridgeFactory: createGenerationDevtoolsBridge, devtools: { ...options.devtools, @@ -150,21 +184,32 @@ 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 }, + onResumeSnapshotChange: setResumeSnapshotState, } let client: GenerationClient @@ -185,6 +230,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. @@ -205,6 +252,7 @@ export function createGeneration< } const dispose = () => { + disposed = true client.dispose() } @@ -230,5 +278,17 @@ export function createGeneration< reset, dispose, updateBody, + get resumeSnapshot() { + return resumeSnapshot + }, + get resumeState() { + return resumeSnapshot?.resumeState ?? null + }, + get pendingArtifacts() { + return resumeSnapshot?.pendingArtifacts ?? [] + }, + get resultArtifacts() { + return resumeSnapshot?.result?.artifacts ?? [] + }, } } diff --git a/packages/ai-svelte/src/create-summarize.svelte.ts b/packages/ai-svelte/src/create-summarize.svelte.ts index 0009eee15..5d582de21 100644 --- a/packages/ai-svelte/src/create-summarize.svelte.ts +++ b/packages/ai-svelte/src/create-summarize.svelte.ts @@ -1,4 +1,8 @@ import { createGeneration } from './create-generation.svelte' +import type { + CreateGenerationOptions, + CreateGenerationReturn, +} from './create-generation.svelte' import type { StreamChunk, SummarizationResult } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -14,7 +18,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' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for summarization */ @@ -46,7 +55,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 +68,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 } /** @@ -134,5 +139,18 @@ export function createSummarize( stop: gen.stop, reset: gen.reset, updateBody: gen.updateBody, + dispose: gen.dispose, + get resumeSnapshot() { + return gen.resumeSnapshot + }, + get resumeState() { + return gen.resumeState + }, + get pendingArtifacts() { + return gen.pendingArtifacts + }, + get resultArtifacts() { + return gen.resultArtifacts + }, } } diff --git a/packages/ai-svelte/src/create-transcription.svelte.ts b/packages/ai-svelte/src/create-transcription.svelte.ts index b6583cb98..f6a7fdfae 100644 --- a/packages/ai-svelte/src/create-transcription.svelte.ts +++ b/packages/ai-svelte/src/create-transcription.svelte.ts @@ -1,4 +1,8 @@ import { createGeneration } from './create-generation.svelte' +import type { + CreateGenerationOptions, + CreateGenerationReturn, +} from './create-generation.svelte' import type { StreamChunk, TranscriptionResult } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -14,7 +18,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' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for transcription */ @@ -46,7 +59,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 +72,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 } /** @@ -141,5 +150,18 @@ export function createTranscription( stop: gen.stop, reset: gen.reset, updateBody: gen.updateBody, + dispose: gen.dispose, + get resumeSnapshot() { + return gen.resumeSnapshot + }, + get resumeState() { + return gen.resumeState + }, + get pendingArtifacts() { + return gen.pendingArtifacts + }, + get resultArtifacts() { + return gen.resultArtifacts + }, } } diff --git a/packages/ai-svelte/tests/create-generation.test.ts b/packages/ai-svelte/tests/create-generation.test.ts index ca2761c9f..c9e4e1185 100644 --- a/packages/ai-svelte/tests/create-generation.test.ts +++ b/packages/ai-svelte/tests/create-generation.test.ts @@ -8,6 +8,11 @@ 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 { + ConnectConnectionAdapter, + GenerationResumeSnapshot, + RunAgentInputContext, +} from '@tanstack/ai-client' // Helper to create generation stream chunks function createGenerationChunks(result: unknown): Array { @@ -71,6 +76,33 @@ 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 } +} + // Helper to create error stream chunks function createErrorChunks(message: string): Array { return [ @@ -186,6 +218,31 @@ 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() + 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('stop and reset', () => { @@ -489,7 +546,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 +561,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 +595,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 +715,28 @@ 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 persisted snapshot remains exposed as read-only state. + expect(gen.resumeSnapshot).toEqual(videoResumeSnapshot) + 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/use-generate-audio.ts b/packages/ai-vue/src/use-generate-audio.ts index 91c8eb6a0..48fb347d6 100644 --- a/packages/ai-vue/src/use-generate-audio.ts +++ b/packages/ai-vue/src/use-generate-audio.ts @@ -1,4 +1,8 @@ import { useGeneration } from './use-generation' +import type { + UseGenerationOptions, + UseGenerationReturn, +} from './use-generation' import type { AudioGenerationResult, StreamChunk } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -15,7 +19,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' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for audio generation */ @@ -47,7 +56,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 +69,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 +108,19 @@ 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, + }) return { - generate: generate as (input: AudioGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, + ...generation, + generate: generation.generate as ( + input: AudioGenerateInput, + ) => Promise, } } diff --git a/packages/ai-vue/src/use-generate-image.ts b/packages/ai-vue/src/use-generate-image.ts index f80cc2f5c..f40c88767 100644 --- a/packages/ai-vue/src/use-generate-image.ts +++ b/packages/ai-vue/src/use-generate-image.ts @@ -1,4 +1,8 @@ import { useGeneration } from './use-generation' +import type { + UseGenerationOptions, + UseGenerationReturn, +} from './use-generation' import type { ImageGenerationResult, StreamChunk } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -15,7 +19,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' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for image generation */ @@ -47,7 +56,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 +69,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 +118,19 @@ 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, + }) return { - generate: generate as (input: ImageGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, + ...generation, + generate: generation.generate as ( + input: ImageGenerateInput, + ) => Promise, } } diff --git a/packages/ai-vue/src/use-generate-speech.ts b/packages/ai-vue/src/use-generate-speech.ts index 2302fc766..574d58517 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' | '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,19 @@ 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, + ...generation, + generate: generation.generate as ( + input: SpeechGenerateInput, + ) => Promise, } } diff --git a/packages/ai-vue/src/use-generate-video.ts b/packages/ai-vue/src/use-generate-video.ts index 644558508..06c9f3970 100644 --- a/packages/ai-vue/src/use-generate-video.ts +++ b/packages/ai-vue/src/use-generate-video.ts @@ -14,11 +14,16 @@ import type { ConnectConnectionAdapter, GenerationClientState, GenerationFetcher, + GenerationPendingArtifact, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, VideoGenerateInput, VideoGenerateResult, VideoStatusInfo, } from '@tanstack/ai-client' +import type { PersistedArtifactRef } from '@tanstack/ai/client' import type { DeepReadonly, ShallowRef } from 'vue' /** @@ -37,6 +42,10 @@ export interface UseGenerateVideoOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** Server-side lightweight generation state persistence. */ + persistence?: GenerationPersistence + /** Initial lightweight resume snapshot restored by the app (read-only state). */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when video generation completes. Can optionally return a transformed value. * @@ -81,6 +90,14 @@ export interface UseGenerateVideoReturn { stop: () => void /** Clear all state and return to idle */ reset: () => void + /** Lightweight generation resume snapshot, if one is available */ + resumeSnapshot: DeepReadonly> + /** Observed run/cursor metadata from the snapshot (read-only state) */ + resumeState: DeepReadonly> + /** Pending persisted artifact references observed during generation/replay */ + pendingArtifacts: DeepReadonly>> + /** Final persisted artifact references observed from a replayed result */ + resultArtifacts: DeepReadonly>> } /** @@ -137,6 +154,29 @@ export function useGenerateVideo( const isLoading = shallowRef(false) const error = shallowRef(undefined) const status = shallowRef('idle') + const resumeSnapshot = shallowRef( + options.initialResumeSnapshot, + ) + const resumeState = shallowRef( + options.initialResumeSnapshot?.resumeState ?? null, + ) + const pendingArtifacts = shallowRef>( + options.initialResumeSnapshot?.pendingArtifacts ?? [], + ) + const resultArtifacts = shallowRef>( + options.initialResumeSnapshot?.result?.artifacts ?? [], + ) + let disposed = false + + const setResumeSnapshotState = ( + snapshot: GenerationResumeSnapshot | undefined, + ) => { + if (disposed) return + resumeSnapshot.value = snapshot + resumeState.value = snapshot?.resumeState ?? null + pendingArtifacts.value = snapshot?.pendingArtifacts ?? [] + resultArtifacts.value = snapshot?.result?.artifacts ?? [] + } // Conditional spread on `body`: `VideoGenerationClientOptions.body` is a // strict optional and under EOPT we must omit the key when absent rather @@ -144,6 +184,12 @@ export function useGenerateVideo( const baseOptions = { id: clientId, body: options.body, + ...(options.persistence !== undefined && { + persistence: options.persistence, + }), + ...(options.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: options.initialResumeSnapshot, + }), devtoolsBridgeFactory: createVideoDevtoolsBridge, devtools: { ...options.devtools, @@ -157,29 +203,46 @@ 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 }, + onResumeSnapshotChange: setResumeSnapshotState, } let client: VideoGenerationClient @@ -212,12 +275,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 +313,9 @@ export function useGenerateVideo( status: readonly(status), stop, reset, + resumeSnapshot: readonly(resumeSnapshot), + resumeState: readonly(resumeState), + pendingArtifacts: readonly(pendingArtifacts), + resultArtifacts: readonly(resultArtifacts), } } diff --git a/packages/ai-vue/src/use-generation.ts b/packages/ai-vue/src/use-generation.ts index 596f14e97..f49d51957 100644 --- a/packages/ai-vue/src/use-generation.ts +++ b/packages/ai-vue/src/use-generation.ts @@ -15,8 +15,13 @@ import type { GenerationClientOptions, GenerationClientState, GenerationFetcher, + GenerationPendingArtifact, + GenerationPersistence, + GenerationResumeSnapshot, + GenerationResumeState, InferGenerationOutputFromReturn, } from '@tanstack/ai-client' +import type { PersistedArtifactRef } from '@tanstack/ai/client' import type { DeepReadonly, ShallowRef } from 'vue' /** @@ -39,6 +44,10 @@ export interface UseGenerationOptions { body?: Record /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions + /** Server-side lightweight generation state persistence. */ + persistence?: GenerationPersistence + /** Initial lightweight resume snapshot restored by the app (read-only state). */ + initialResumeSnapshot?: GenerationResumeSnapshot /** * Callback when a result is received. Can optionally return a transformed value. * @@ -75,6 +84,14 @@ export interface UseGenerationReturn { stop: () => void /** Clear result, error, and return to idle */ reset: () => void + /** Lightweight generation resume snapshot, if one is available */ + resumeSnapshot: DeepReadonly> + /** Observed run/cursor metadata from the snapshot (read-only state) */ + resumeState: DeepReadonly> + /** Pending persisted artifact references observed during generation/replay */ + pendingArtifacts: DeepReadonly>> + /** Final persisted artifact references observed from a replayed result */ + resultArtifacts: DeepReadonly>> } /** @@ -103,7 +120,7 @@ export interface UseGenerationReturn { // inference site that works even for an optional nested property), which types // the callback parameter as `TResult` and narrows `result`. Inferring the // whole callback as a defaulted type parameter instead collapses to the -// default, leaving the parameter `any` — a hard error under `strict`. See +// default, leaving the parameter `any` — a hard error under `strict`. See // issue #848. export function useGeneration< TInput extends Record, @@ -122,6 +139,29 @@ export function useGeneration< const isLoading = shallowRef(false) const error = shallowRef(undefined) const status = shallowRef('idle') + const resumeSnapshot = shallowRef( + options.initialResumeSnapshot, + ) + const resumeState = shallowRef( + options.initialResumeSnapshot?.resumeState ?? null, + ) + const pendingArtifacts = shallowRef>( + options.initialResumeSnapshot?.pendingArtifacts ?? [], + ) + const resultArtifacts = shallowRef>( + options.initialResumeSnapshot?.result?.artifacts ?? [], + ) + let disposed = false + + const setResumeSnapshotState = ( + snapshot: GenerationResumeSnapshot | undefined, + ) => { + if (disposed) return + resumeSnapshot.value = snapshot + resumeState.value = snapshot?.resumeState ?? null + pendingArtifacts.value = snapshot?.pendingArtifacts ?? [] + resultArtifacts.value = snapshot?.result?.artifacts ?? [] + } // Conditional spread on `body`: `GenerationClientOptions.body` is a strict // optional (`body?: Record`), and under EOPT we must omit the @@ -129,6 +169,12 @@ export function useGeneration< const clientOptions: GenerationClientOptions = { id: clientId, body: options.body, + ...(options.persistence !== undefined && { + persistence: options.persistence, + }), + ...(options.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: options.initialResumeSnapshot, + }), devtoolsBridgeFactory: createGenerationDevtoolsBridge, devtools: { ...options.devtools, @@ -141,21 +187,32 @@ 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 }, + onResumeSnapshotChange: setResumeSnapshotState, } let client: GenerationClient @@ -188,12 +245,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() }) @@ -221,5 +281,9 @@ export function useGeneration< status: readonly(status), stop, reset, + resumeSnapshot: readonly(resumeSnapshot), + resumeState: readonly(resumeState), + pendingArtifacts: readonly(pendingArtifacts), + resultArtifacts: readonly(resultArtifacts), } } diff --git a/packages/ai-vue/src/use-summarize.ts b/packages/ai-vue/src/use-summarize.ts index 7297c5277..e49a5a1f2 100644 --- a/packages/ai-vue/src/use-summarize.ts +++ b/packages/ai-vue/src/use-summarize.ts @@ -1,4 +1,8 @@ import { useGeneration } from './use-generation' +import type { + UseGenerationOptions, + UseGenerationReturn, +} from './use-generation' import type { StreamChunk, SummarizationResult } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -15,7 +19,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' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for summarization */ @@ -47,7 +56,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 +70,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 +114,19 @@ 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, + }) return { - generate: generate as (input: SummarizeGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, + ...generation, + generate: generation.generate as ( + input: SummarizeGenerateInput, + ) => Promise, } } diff --git a/packages/ai-vue/src/use-transcription.ts b/packages/ai-vue/src/use-transcription.ts index 26e42156c..69410e039 100644 --- a/packages/ai-vue/src/use-transcription.ts +++ b/packages/ai-vue/src/use-transcription.ts @@ -1,4 +1,8 @@ import { useGeneration } from './use-generation' +import type { + UseGenerationOptions, + UseGenerationReturn, +} from './use-generation' import type { StreamChunk, TranscriptionResult } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -15,7 +19,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' | 'initialResumeSnapshot' +> { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter /** Direct async function for transcription */ @@ -47,7 +60,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 +73,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 +122,16 @@ 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 }) return { - generate: generate as (input: TranscriptionGenerateInput) => Promise, - result, - isLoading, - error, - status, - stop, - reset, + ...generation, + generate: generation.generate as ( + input: TranscriptionGenerateInput, + ) => Promise, } } diff --git a/packages/ai-vue/tests/use-generation.test.ts b/packages/ai-vue/tests/use-generation.test.ts index 21b40eafe..7004ba995 100644 --- a/packages/ai-vue/tests/use-generation.test.ts +++ b/packages/ai-vue/tests/use-generation.test.ts @@ -10,6 +10,11 @@ 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 { + ConnectConnectionAdapter, + GenerationResumeSnapshot, + RunAgentInputContext, +} from '@tanstack/ai-client' import type { DeepReadonly } from 'vue' // Helper to create generation stream chunks @@ -62,6 +67,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 +246,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', () => { @@ -586,7 +675,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 +693,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 +850,33 @@ 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 persisted snapshot remains exposed as read-only state. + expect(result.resumeSnapshot.value).toEqual(videoResumeSnapshot) + expect(result.resumeState.value).toEqual(videoResumeSnapshot.resumeState) + }) + it('should require either connection or fetcher', () => { expect(() => { renderHook(() => useGenerateVideo({} as any)) 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 4448c4e44..1fc59f27e 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/pnpm-lock.yaml b/pnpm-lock.yaml index 63d35dbed..3647c95c6 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -2050,6 +2050,10 @@ importers: version: 4.3.6 packages/ai-persistence: + dependencies: + '@tanstack/ai-utils': + specifier: workspace:* + version: link:../ai-utils devDependencies: '@tanstack/ai': specifier: workspace:* @@ -31664,8 +31668,8 @@ snapshots: tinyglobby@0.2.17: dependencies: - fdir: 6.5.0(picomatch@4.0.4) - picomatch: 4.0.4 + fdir: 6.5.0(picomatch@4.0.5) + picomatch: 4.0.5 tinypool@2.1.0: {} @@ -31815,7 +31819,7 @@ snapshots: tsx@4.21.0: dependencies: esbuild: 0.27.7 - get-tsconfig: 4.13.0 + get-tsconfig: 4.14.0 optionalDependencies: fsevents: 2.3.3