diff --git a/.changeset/generation-persistence.md b/.changeset/generation-persistence.md index 0b1014a4d..3602d9f22 100644 --- a/.changeset/generation-persistence.md +++ b/.changeset/generation-persistence.md @@ -1,4 +1,7 @@ --- +'@tanstack/ai': minor +'@tanstack/ai-utils': minor +'@tanstack/ai-persistence': minor '@tanstack/ai-client': minor '@tanstack/ai-react': minor '@tanstack/ai-solid': minor @@ -7,7 +10,11 @@ '@tanstack/ai-angular': minor --- -Add client-side generation persistence: a lightweight resume snapshot for media generation activities. +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 `pendingArtifacts` / `resultArtifacts`, which stay empty until the server-side artifact pipeline ships in a follow-up). diff --git a/docs/persistence/generation-persistence.md b/docs/persistence/generation-persistence.md index f41c0cfc4..e3e059980 100644 --- a/docs/persistence/generation-persistence.md +++ b/docs/persistence/generation-persistence.md @@ -10,11 +10,13 @@ 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 two things: +It helps with three things: - **After a reload**, show what the last run was: whether it finished, what failed, and metadata like the result id or video job id. This is a small snapshot kept in browser storage and read back automatically on mount. +- **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. @@ -22,12 +24,13 @@ It helps with two things: ## When to use it Use it when a run is long enough that a reload or a dropped connection actually -matters: video, batch images, long audio, transcription of a big file. For a -quick one-shot image you show and forget, you can skip it. +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. -It never stores the generated bytes. Only the run's identity, status, errors, -and result metadata (like the result id, model, or video job id) are saved. See -[What it does not store](#what-it-does-not-store). +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 @@ -86,6 +89,104 @@ each run in it, keyed by the request id it generates for the run. 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`, and give the hook a stable `id`. The @@ -174,10 +275,11 @@ reload as informational ("a run was still going when the page closed"). See [Resumable Streams](../resumable-streams/overview) for the durability contract, production adapters, and the one-time-side-effects note. -## What it does not store +## What the browser snapshot holds -The snapshot stores run identity and result metadata — ids, model, status, a -video `jobId`, an expiry timestamp — never the generated bytes, and it does not -carry the provider's media URL. If you need the media itself to survive a -reload, save it to your own storage from `onResult`; durable artifact storage -is coming as a follow-up. +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/packages/ai-persistence/package.json b/packages/ai-persistence/package.json index f86e0563d..670a52c35 100644 --- a/packages/ai-persistence/package.json +++ b/packages/ai-persistence/package.json @@ -48,6 +48,9 @@ "test:types": "tsc", "test:oxlint": "oxlint src --type-aware" }, + "dependencies": { + "@tanstack/ai-utils": "workspace:*" + }, "peerDependencies": { "@tanstack/ai": "workspace:*", "vitest": "^4.1.10" diff --git a/packages/ai-persistence/src/index.ts b/packages/ai-persistence/src/index.ts index 9007aa83b..25e556643 100644 --- a/packages/ai-persistence/src/index.ts +++ b/packages/ai-persistence/src/index.ts @@ -23,6 +23,16 @@ export type { ChatTranscriptPersistence, ChatPersistence, ChatWithInterruptsPersistence, + // Generation media bytes (opt-in; pair `artifacts` with `blobs`) + ArtifactRecord, + ArtifactStore, + BlobBody, + BlobRecord, + BlobObject, + BlobListPage, + BlobPutOptions, + BlobListOptions, + BlobStore, AIPersistence, AIPersistenceOverrides, ComposedAIPersistenceStores, @@ -33,14 +43,30 @@ export type { // AIPersistenceStores is intentionally NOT re-exported — use a named chat // shape or AIPersistence<{ messages: MessageStore, … }>. +// Core artifact wire types (re-exported for convenience) +export type { + PersistedArtifactActivity, + PersistedArtifactRef, + PersistedArtifactRole, +} from '@tanstack/ai' + // Middleware (state only — locks live in @tanstack/ai as withLocks) export { withPersistence, withGenerationPersistence } from './middleware' +export type { + WithPersistenceOptions, + GenerationArtifactDescriptor, + GenerationArtifactExtractionInput, + GenerationArtifactNameInput, +} from './middleware' // Server helper: rehydrate a thread's messages for a client load export { reconstructChat } from './reconstruct' export type { ReconstructChatOptions } from './reconstruct' -// Reference in-memory implementation (state stores only) +// Server helpers: retrieve a persisted generation artifact + its bytes +export { retrieveArtifact, retrieveBlob, artifactBlobKey } from './retrieve' + +// Reference in-memory implementation (state stores + generation media) export { memoryPersistence } from './memory' // Interrupt controller diff --git a/packages/ai-persistence/src/memory.ts b/packages/ai-persistence/src/memory.ts index 70f8368f1..c15c2535a 100644 --- a/packages/ai-persistence/src/memory.ts +++ b/packages/ai-persistence/src/memory.ts @@ -1,7 +1,13 @@ import { defineAIPersistence } from './types' import type { ModelMessage } from '@tanstack/ai' import type { - ChatPersistence, + ArtifactRecord, + ArtifactStore, + BlobBody, + BlobListOptions, + BlobObject, + BlobRecord, + BlobStore, InterruptRecord, InterruptStore, MessageStore, @@ -169,20 +175,227 @@ 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 +} + /** - * In-process reference backend for **chat** state stores. + * In-process reference backend for the **chat** state stores plus the optional + * generation media stores. * - * Returns {@link ChatPersistence} (messages + runs + interrupts + metadata). - * Locks are not included — use `InMemoryLockStore` + `withLocks` from - * `@tanstack/ai` when a test or single-process app needs coordination. + * Returns messages + runs + interrupts + metadata (the chat persistence shape) + * widened with `artifacts` + `blobs`, so it satisfies `withPersistence` and the + * byte-persisting path of `withGenerationPersistence` alike. Locks are not + * included — use `InMemoryLockStore` + `withLocks` from `@tanstack/ai` when a + * test or single-process app needs coordination. */ -export function memoryPersistence(): ChatPersistence { - return defineAIPersistence({ - stores: { - messages: new MemoryMessageStore(), - runs: new MemoryRunStore(), - interrupts: new MemoryInterruptStore(), - metadata: new MemoryMetadataStore(), - }, - }) +export function memoryPersistence() { + const stores: MemoryPersistenceStores = { + messages: new MemoryMessageStore(), + runs: new MemoryRunStore(), + interrupts: new MemoryInterruptStore(), + metadata: new MemoryMetadataStore(), + artifacts: new MemoryArtifactStore(), + blobs: new MemoryBlobStore(), + } + return defineAIPersistence({ stores }) } diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index cd86cfab9..43090aba7 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -1,4 +1,5 @@ import { defineChatMiddleware } from '@tanstack/ai' +import { base64ToUint8Array } from '@tanstack/ai-utils' import { InterruptsCapability, PersistenceCapability, @@ -23,6 +24,9 @@ import type { GenerationMiddleware, GenerationMiddlewareContext, ModelMessage, + PersistedArtifactActivity, + PersistedArtifactRef, + PersistedArtifactRole, RunAgentResumeItem, StreamChunk, ToolApprovalResolution, @@ -31,10 +35,55 @@ import type { import type { AIPersistence, AIPersistenceStores, + ArtifactRecord, + BlobBody, ChatTranscriptStores, InterruptRecord, RunStore, } from './types' +import { artifactBlobKey } from './retrieve' + +export interface WithPersistenceOptions { + extractArtifacts?: ( + input: GenerationArtifactExtractionInput, + ) => + | Array + | Promise> + nameArtifact?: (input: GenerationArtifactNameInput) => string +} + +export interface GenerationArtifactDescriptor { + role: PersistedArtifactRole + path: string + mediaType?: PersistedArtifactRef['source']['mediaType'] + mimeType?: string + bytes?: BlobBody + url?: string + json?: unknown + name?: string + jobId?: string + expiresAt?: string | Date +} + +export interface GenerationArtifactExtractionInput { + activity: PersistedArtifactActivity + provider: string + model: string + threadId: string + runId: string + inputs: unknown + result: unknown +} + +export interface GenerationArtifactNameInput { + descriptor: GenerationArtifactDescriptor + activity: PersistedArtifactActivity + provider: string + model: string + threadId: string + runId: string + index: number +} interface RunStateEntry { merged: boolean @@ -259,18 +308,483 @@ 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 // --------------------------------------------------------------------------- interface PersistencePlan { wantsInterrupts: boolean + wantsArtifactPersistence: boolean runs: AIPersistence['stores']['runs'] } function resolvePersistencePlan(persistence: AIPersistence): PersistencePlan { return { wantsInterrupts: persistence.stores.interrupts !== undefined, + wantsArtifactPersistence: + persistence.stores.artifacts !== undefined && + persistence.stores.blobs !== undefined, runs: persistence.stores.runs, } } @@ -313,9 +827,20 @@ type InvalidChatPersistence = type ValidChatPersistence = InvalidChatPersistence extends true ? never : unknown -/** Generation entrypoint invalid when `runs` is known-absent. */ +/** + * Generation entrypoint invalid when: + * - `runs` is known-absent, or + * - exactly one of `artifacts` / `blobs` is known-present (the byte path writes + * to both, so one without the other is always a misconfiguration). + */ type InvalidGenerationPersistence = - StoreIsDefinitelyAbsent extends true ? true : false + StoreIsDefinitelyAbsent extends true + ? true + : StoreIsDefinitelyPresent extends true + ? StoreIsDefinitelyAbsent + : StoreIsDefinitelyPresent extends true + ? StoreIsDefinitelyAbsent + : false type ValidGenerationPersistence = InvalidGenerationPersistence extends true ? never : unknown @@ -622,10 +1147,12 @@ export function withPersistence( // --------------------------------------------------------------------------- /** - * 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. * - * Requires `stores.runs`. + * Requires `stores.runs`. `stores.artifacts` + `stores.blobs` (as a pair) opt + * into durable media bytes. * * ⚠️ TEMPORARY / WRONG SHAPE — do not extend this design. * @@ -640,15 +1167,24 @@ export function withPersistence( * later artifact storage — not reuse chat `RunStore` / `MessageStore`. An * optional `threadId` on a job is only a *link* to a chat when the product * needs it, never the job's primary identity. + * + * The artifact path below does **not** resolve this: it keys `ArtifactRecord` + * on `ctx.runId ?? ctx.requestId` while the run record above is still keyed on + * `ctx.requestId`, so a caller-supplied `runId` leaves `runs.get(runId)` empty + * while `artifacts.list(runId)` returns rows. The two stores cannot be joined + * until the job store lands. */ -export function withGenerationPersistence< - TStores extends AIPersistenceStores & { runs: RunStore }, ->( +export function withGenerationPersistence( persistence: AIPersistence & ValidGenerationPersistence, + opts?: WithPersistenceOptions, +): GenerationMiddleware +export function withGenerationPersistence( + persistence: AIPersistence, + opts?: WithPersistenceOptions, ): GenerationMiddleware { validateGenerationPersistenceStores(persistence) - const runStore = persistence.stores.runs - if (!runStore) { + const { wantsArtifactPersistence, runs } = resolvePersistencePlan(persistence) + if (!runs) { // validateGenerationPersistenceStores already throws; this narrows for TypeScript. throw new Error('Generation persistence requires stores.runs.') } @@ -658,25 +1194,40 @@ export function withGenerationPersistence< async onStart(ctx: GenerationMiddlewareContext) { // STOPGAP ONLY — see function JSDoc. Do not copy this pattern. - await createOrResumeRun(runStore, ctx.requestId, ctx.requestId) + 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( ctx: GenerationMiddlewareContext, info: GenerationFinishInfo, ) { - await completeRun(runStore, ctx.requestId, info.usage) + await completeRun(runs, ctx.requestId, info.usage) }, async onError(ctx: GenerationMiddlewareContext, info: GenerationErrorInfo) { - await failRun(runStore, ctx.requestId, info.error) + await failRun(runs, ctx.requestId, info.error) }, async onAbort( ctx: GenerationMiddlewareContext, _info: GenerationAbortInfo, ) { - await interruptRun(runStore, ctx.requestId) + await interruptRun(runs, ctx.requestId) }, } } diff --git a/packages/ai-persistence/src/retrieve.ts b/packages/ai-persistence/src/retrieve.ts new file mode 100644 index 000000000..0984e3ab8 --- /dev/null +++ b/packages/ai-persistence/src/retrieve.ts @@ -0,0 +1,45 @@ +import type { AIPersistence, ArtifactRecord, BlobObject } from './types' + +/** + * The blob-store key a generation artifact's bytes are stored under. + * `withGenerationPersistence` writes bytes to this key; `retrieveBlob` reads + * from it. Keep the two in lockstep by using this helper on both sides. + */ +export function artifactBlobKey( + ref: Pick, +): string { + return `artifacts/${ref.runId}/${ref.artifactId}` +} + +/** + * Look up a persisted generation artifact's metadata by id. Returns `null` when + * the persistence has no `artifacts` store or no record matches — so a serve + * handler can map that straight to a 404. + */ +export async function retrieveArtifact( + persistence: AIPersistence, + artifactId: string, +): Promise { + const record = await persistence.stores.artifacts?.get(artifactId) + return record ?? null +} + +/** + * Look up a persisted generation artifact's stored bytes. Pass an `artifactId` + * (resolved to its record first) or an already-loaded {@link ArtifactRecord} + * (no second metadata lookup). Returns `null` when the artifact, its record, or + * its blob is missing, or the stores are not configured. + */ +export async function retrieveBlob( + persistence: AIPersistence, + artifact: string | ArtifactRecord, +): Promise { + const record = + typeof artifact === 'string' + ? await retrieveArtifact(persistence, artifact) + : artifact + if (!record) return null + + const blob = await persistence.stores.blobs?.get(artifactBlobKey(record)) + return blob ?? null +} diff --git a/packages/ai-persistence/src/types.ts b/packages/ai-persistence/src/types.ts index a1fd7f488..240c569c4 100644 --- a/packages/ai-persistence/src/types.ts +++ b/packages/ai-persistence/src/types.ts @@ -31,9 +31,10 @@ export type { Scope } // // TIMESTAMP CONVENTION // -------------------- -// Store *records* (`RunRecord`, `InterruptRecord`) speak **epoch -// milliseconds** (`number`), the native unit for SQL/`BIGINT` columns and -// `Date.now()`. Wire/result references that leave the persistence layer speak +// Store *records* (`RunRecord`, `InterruptRecord`, `ArtifactRecord`, +// `BlobRecord`) speak **epoch milliseconds** (`number`), the native unit for +// SQL/`BIGINT` columns and `Date.now()`. Wire/result references that leave the +// persistence layer (e.g. core's `PersistedArtifactRef.createdAt`) speak // **ISO-8601 strings**. The middleware performs the number→ISO conversion at // the boundary; do not mix the two on a single field. @@ -277,6 +278,129 @@ export function defineMetadataStore(store: MetadataStore): MetadataStore { return store } +/** + * Metadata row describing a persisted artifact (generated media, tool output). + * + * The bytes themselves live in a {@link BlobStore}; this record holds the + * descriptive metadata and an optional `externalUrl` for reference-only + * backends. + * + * @property createdAt - Epoch ms. (Core's wire-facing `PersistedArtifactRef` + * exposes the same instant as an ISO string; see the timestamp convention.) + */ +export interface ArtifactRecord { + artifactId: string + runId: string + threadId: string + name: string + mimeType: string + size: number + externalUrl?: string + createdAt: number +} + +/** Durable store for artifact metadata records. */ +export interface ArtifactStore { + /** Insert or overwrite the artifact metadata record. */ + save: (record: ArtifactRecord) => Promise + /** Return the artifact for `artifactId`, or `null` if none exists. */ + get: (artifactId: string) => Promise + /** All artifacts for a run. Returns `[]` when the run has none. */ + list: (runId: string) => Promise> + /** OPTIONAL: delete a single artifact by id. */ + delete?: (artifactId: string) => Promise + /** OPTIONAL: delete every artifact belonging to `runId`. */ + deleteForRun?: (runId: string) => Promise +} + +/** + * Accepted body shapes for {@link BlobStore.put}. `ArrayBufferView` already + * covers `Uint8Array` and every other typed-array/`DataView`, so no separate + * `Uint8Array` member is needed. + */ +export type BlobBody = + | ReadableStream + | ArrayBuffer + | ArrayBufferView + | string + | Blob + +/** + * Metadata for a stored blob. + * + * @property size - Byte length, when known. + * @property createdAt - Epoch ms first written. + * @property updatedAt - Epoch ms last overwritten. + */ +export interface BlobRecord { + key: string + size?: number + etag?: string + contentType?: string + customMetadata?: Record + createdAt?: number + updatedAt?: number +} + +/** A stored blob's metadata plus lazy accessors for its bytes. */ +export interface BlobObject extends BlobRecord { + arrayBuffer: () => Promise + text: () => Promise + body?: ReadableStream +} + +/** + * One page of a {@link BlobStore.list} scan. + * + * @property cursor - Opaque continuation token; present only when `truncated`. + * @property truncated - `true` when more objects match beyond this page. + */ +export interface BlobListPage { + objects: Array + cursor?: string + truncated?: boolean +} + +export interface BlobPutOptions { + contentType?: string + customMetadata?: Record +} + +export interface BlobListOptions { + prefix?: string + cursor?: string + limit?: number +} + +/** Durable object/blob store (byte-storing or reference-only backends). */ +export interface BlobStore { + /** Insert or overwrite the object at `key`, returning its metadata. */ + put: ( + key: string, + body: BlobBody, + options?: BlobPutOptions, + ) => Promise + /** Return the object at `key` (metadata + byte accessors), or `null`. */ + get: (key: string) => Promise + /** Return only the metadata for `key`, or `null`. */ + head: (key: string) => Promise + /** Remove the object at `key`. A no-op if absent. */ + delete: (key: string) => Promise + /** + * List objects, optionally filtered by `prefix`, in ascending key order. + * + * CURSOR SEMANTICS: `prefix` matches literally and case-sensitively (SQL + * backends must escape LIKE metacharacters, so `run_` matches only the exact + * bytes `run_`, not `_` as a wildcard). When `limit` is given and more keys + * match, the page is `truncated: true` with a `cursor`; passing that `cursor` + * back returns the strictly-following keys (keys `> cursor`). Cursor ordering + * is the same byte ordering as the sort, so paging visits every key exactly + * once with no gaps or repeats. `limit: 0` yields an empty, untruncated page + * with no cursor. + */ + list: (options?: BlobListOptions) => Promise +} + /** * Sparse bag of **state** store keys — composition / validation only. * @@ -293,6 +417,8 @@ export interface AIPersistenceStores { runs?: RunStore interrupts?: InterruptStore metadata?: MetadataStore + artifacts?: ArtifactStore + blobs?: BlobStore } /** @@ -465,6 +591,8 @@ const storeKeys = [ 'runs', 'interrupts', 'metadata', + 'artifacts', + 'blobs', ] satisfies Array const storeKeySet = new Set(storeKeys) @@ -499,8 +627,12 @@ export function validateChatPersistenceStores( } /** - * Generation middleware entrypoint rule: `runs` is required (run lifecycle is - * the only generation state this middleware tracks). + * Generation middleware entrypoint rules: + * - `runs` is required (run lifecycle is the generation state this middleware + * always tracks) + * - `artifacts` and `blobs` are optional, but come as a pair — the byte path + * writes the bytes to `blobs` and the describing row to `artifacts`, so one + * without the other is always a misconfiguration */ export function validateGenerationPersistenceStores( persistence: AIPersistence, @@ -509,6 +641,13 @@ export function validateGenerationPersistenceStores( if (!persistence.stores.runs) { throw new Error('Generation persistence requires stores.runs.') } + const hasArtifacts = persistence.stores.artifacts !== undefined + const hasBlobs = persistence.stores.blobs !== undefined + if (hasArtifacts !== hasBlobs) { + throw new Error( + 'Generation artifact persistence requires both stores.artifacts and stores.blobs.', + ) + } } /** 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..ebce9d867 --- /dev/null +++ b/packages/ai-persistence/tests/generation-artifacts.test.ts @@ -0,0 +1,499 @@ +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() + // `runs` is present so the required-runs rule passes and the artifacts / + // blobs pairing rule is the one under test. + const persistence: AIPersistence = defineAIPersistence({ + stores: { + runs: full.stores.runs, + artifacts: full.stores.artifacts, + }, + }) + + expect(() => withGenerationPersistence(persistence)).toThrow( + /artifact persistence requires both stores\.artifacts and stores\.blobs/i, + ) + }) + + it('persists transcription structured JSON output', async () => { + const persistence = memoryPersistence() + + const result = (await generateTranscription({ + adapter: transcriptionAdapter(), + audio: 'aW5wdXQtYXVkaW8=', + responseFormat: 'verbose_json', + threadId: 'thread-transcription', + runId: 'run-transcription', + middleware: [withGenerationPersistence(persistence)], + } as TranscriptionGenerateOptions)) as TranscriptionResult + + expect(result.artifacts?.map((artifact) => artifact.role)).toEqual([ + 'input', + 'output', + ]) + const structured = result.artifacts?.find( + (artifact) => artifact.source.mediaType === 'json', + ) as PersistedArtifactRef | undefined + expect(structured).toMatchObject({ + role: 'output', + mimeType: 'application/json', + source: { + activity: 'transcription', + path: 'transcription', + mediaType: 'json', + }, + }) + const blob = await persistence.stores.blobs!.get( + `artifacts/run-transcription/${structured!.artifactId}`, + ) + await expect(blob?.text()).resolves.toContain('"segments"') + }) +}) diff --git a/packages/ai-persistence/tests/memory.test.ts b/packages/ai-persistence/tests/memory.test.ts index 95820b9a8..811c3ce51 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 (no locks)', () => { expect(Object.keys(memoryPersistence().stores).sort()).toEqual([ + 'artifacts', + 'blobs', 'interrupts', 'messages', 'metadata', diff --git a/packages/ai-persistence/tests/persistence-types.test-d.ts b/packages/ai-persistence/tests/persistence-types.test-d.ts index 60f040141..2c5d209aa 100644 --- a/packages/ai-persistence/tests/persistence-types.test-d.ts +++ b/packages/ai-persistence/tests/persistence-types.test-d.ts @@ -14,6 +14,8 @@ import { InMemoryLockStore, withLocks } from '@tanstack/ai/locks' import type { LockStore } from '@tanstack/ai/locks' import type { AIPersistence, + ArtifactStore, + BlobStore, ChatPersistence, ChatTranscriptPersistence, ChatTranscriptStores, @@ -123,8 +125,18 @@ const uncertainInherited = composePersistence(base, { }) expectTypeOf(uncertainInherited.stores.messages).toEqualTypeOf() -// Named shapes -expectTypeOf(memoryPersistence()).toEqualTypeOf() +// Named shapes. memoryPersistence() is the chat shape widened with the +// generation media stores, so it still satisfies every chat entrypoint while +// also opting into the byte path of withGenerationPersistence. +const memory = memoryPersistence() +expectTypeOf(memory.stores.messages).toEqualTypeOf() +expectTypeOf(memory.stores.runs).toEqualTypeOf() +expectTypeOf(memory.stores.interrupts).toEqualTypeOf() +expectTypeOf(memory.stores.metadata).toEqualTypeOf() +expectTypeOf(memory.stores.artifacts).toEqualTypeOf() +expectTypeOf(memory.stores.blobs).toEqualTypeOf() +withPersistence(memory) +withGenerationPersistence(memory) const transcript: ChatTranscriptPersistence = messagesOnly void transcript declare const fullChat: ChatPersistence 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/src/activities/generateAudio/index.ts b/packages/ai/src/activities/generateAudio/index.ts index ec2e72890..cb25e2a0b 100644 --- a/packages/ai/src/activities/generateAudio/index.ts +++ b/packages/ai/src/activities/generateAudio/index.ts @@ -9,6 +9,7 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { streamGenerationResult } from '../stream-generation-result.js' import { resolveDebugOption } from '../../logger/resolve' import { + applyGenerationResultTransforms, createGenerationContext, runGenerationError, runGenerationFinish, @@ -84,6 +85,10 @@ export interface AudioActivityOptions< * `GenerationMiddleware` contract for a custom backend. */ middleware?: Array + /** Stable conversation/thread id for correlating this run when persisted. */ + threadId?: string + /** Stable run id for correlating this run when persisted. */ + runId?: string } // =========================== @@ -134,8 +139,9 @@ export function generateAudio< options: AudioActivityOptions, ): AudioActivityResult { if (options.stream) { - return streamGenerationResult(() => - runGenerateAudio(options), + return streamGenerationResult( + (resolved) => runGenerateAudio({ ...options, ...resolved }), + options, ) as AudioActivityResult } return runGenerateAudio(options) as AudioActivityResult @@ -154,6 +160,8 @@ async function runGenerateAudio< stream: _stream, debug: _debug, middleware, + threadId, + runId, ...rest } = options const model = adapter.model @@ -171,6 +179,9 @@ async function runGenerateAudio< provider: adapter.name, model, modelOptions: rest.modelOptions, + threadId, + runId, + artifactInputs: { prompt: rest.prompt, duration: rest.duration }, createId, }) @@ -192,7 +203,8 @@ async function runGenerateAudio< }) try { - const result = await adapter.generateAudio({ ...rest, model, logger }) + const rawResult = await adapter.generateAudio({ ...rest, model, logger }) + const result = await applyGenerationResultTransforms(mwCtx, rawResult) const elapsedMs = Date.now() - startTime aiEventClient.emit('audio:request:completed', { diff --git a/packages/ai/src/activities/generateImage/index.ts b/packages/ai/src/activities/generateImage/index.ts index 8021e0ce2..5e49baad1 100644 --- a/packages/ai/src/activities/generateImage/index.ts +++ b/packages/ai/src/activities/generateImage/index.ts @@ -9,6 +9,7 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { streamGenerationResult } from '../stream-generation-result.js' import { resolveDebugOption } from '../../logger/resolve' import { + applyGenerationResultTransforms, createGenerationContext, runGenerationError, runGenerationFinish, @@ -137,6 +138,10 @@ export type ImageActivityOptions< * `GenerationMiddleware` contract for a custom backend. */ middleware?: Array + /** Stable conversation/thread id for correlating this run when persisted. */ + threadId?: string + /** Stable run id for correlating this run when persisted. */ + runId?: string } & ({} extends ImageProviderOptionsForModel ? { /** Provider-specific options for image generation */ modelOptions?: ImageProviderOptionsForModel< @@ -225,8 +230,9 @@ export function generateImage< options: ImageActivityOptions, ): ImageActivityResult { if (options.stream) { - return streamGenerationResult(() => - runGenerateImage(options), + return streamGenerationResult( + (resolved) => runGenerateImage({ ...options, ...resolved }), + options, ) as ImageActivityResult } @@ -247,6 +253,8 @@ async function runGenerateImage< stream: _stream, debug: _debug, middleware, + threadId, + runId, ...rest } = options const model = adapter.model @@ -260,6 +268,9 @@ async function runGenerateImage< provider: adapter.name, model, modelOptions: rest.modelOptions, + threadId, + runId, + artifactInputs: { prompt: rest.prompt }, createId, }) @@ -295,7 +306,8 @@ async function runGenerateImage< }) try { - const result = await adapter.generateImages({ ...rest, model, logger }) + const rawResult = await adapter.generateImages({ ...rest, model, logger }) + const result = await applyGenerationResultTransforms(mwCtx, rawResult) const duration = Date.now() - startTime aiEventClient.emit('image:request:completed', { diff --git a/packages/ai/src/activities/generateSpeech/index.ts b/packages/ai/src/activities/generateSpeech/index.ts index 06ac193c8..0a134f4cf 100644 --- a/packages/ai/src/activities/generateSpeech/index.ts +++ b/packages/ai/src/activities/generateSpeech/index.ts @@ -9,6 +9,7 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { streamGenerationResult } from '../stream-generation-result.js' import { resolveDebugOption } from '../../logger/resolve' import { + applyGenerationResultTransforms, createGenerationContext, runGenerationError, runGenerationFinish, @@ -87,6 +88,10 @@ export interface TTSActivityOptions< * `GenerationMiddleware` contract for a custom backend. */ middleware?: Array + /** Stable conversation/thread id for correlating this run when persisted. */ + threadId?: string + /** Stable run id for correlating this run when persisted. */ + runId?: string } // =========================== @@ -144,8 +149,9 @@ export function generateSpeech< TStream extends boolean = false, >(options: TTSActivityOptions): TTSActivityResult { if (options.stream) { - return streamGenerationResult(() => - runGenerateSpeech(options), + return streamGenerationResult( + (resolved) => runGenerateSpeech({ ...options, ...resolved }), + options, ) as TTSActivityResult } return runGenerateSpeech(options) as TTSActivityResult @@ -162,6 +168,8 @@ async function runGenerateSpeech< stream: _stream, debug: _debug, middleware, + threadId, + runId, ...rest } = options const model = adapter.model @@ -179,6 +187,14 @@ async function runGenerateSpeech< provider: adapter.name, model, modelOptions: rest.modelOptions, + artifactInputs: { + text: rest.text, + voice: rest.voice, + format: rest.format, + speed: rest.speed, + }, + threadId, + runId, createId, }) @@ -202,7 +218,8 @@ async function runGenerateSpeech< }) try { - const result = await adapter.generateSpeech({ ...rest, model, logger }) + const rawResult = await adapter.generateSpeech({ ...rest, model, logger }) + const result = await applyGenerationResultTransforms(mwCtx, rawResult) const duration = Date.now() - startTime aiEventClient.emit('speech:request:completed', { diff --git a/packages/ai/src/activities/generateTranscription/index.ts b/packages/ai/src/activities/generateTranscription/index.ts index 03c6a9ed1..ccf85f235 100644 --- a/packages/ai/src/activities/generateTranscription/index.ts +++ b/packages/ai/src/activities/generateTranscription/index.ts @@ -9,6 +9,7 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { streamGenerationResult } from '../stream-generation-result.js' import { resolveDebugOption } from '../../logger/resolve' import { + applyGenerationResultTransforms, createGenerationContext, runGenerationError, runGenerationFinish, @@ -94,6 +95,10 @@ export interface TranscriptionActivityOptions< * `GenerationMiddleware` contract for a custom backend. */ middleware?: Array + /** Stable conversation/thread id for correlating this run when persisted. */ + threadId?: string + /** Stable run id for correlating this run when persisted. */ + runId?: string } // =========================== @@ -171,8 +176,9 @@ export function generateTranscription< options: TranscriptionActivityOptions, ): TranscriptionActivityResult { if (options.stream) { - return streamGenerationResult(() => - runGenerateTranscription(options), + return streamGenerationResult( + (resolved) => runGenerateTranscription({ ...options, ...resolved }), + options, ) as TranscriptionActivityResult } @@ -197,6 +203,8 @@ async function runGenerateTranscription< stream: _stream, debug: _debug, middleware, + threadId, + runId, ...rest } = options const model = adapter.model @@ -214,6 +222,14 @@ async function runGenerateTranscription< provider: adapter.name, model, modelOptions: rest.modelOptions, + artifactInputs: { + audio: rest.audio, + language: rest.language, + prompt: rest.prompt, + responseFormat: rest.responseFormat, + }, + threadId, + runId, createId, }) @@ -236,7 +252,8 @@ async function runGenerateTranscription< }) try { - const result = await adapter.transcribe({ ...rest, model, logger }) + const rawResult = await adapter.transcribe({ ...rest, model, logger }) + const result = await applyGenerationResultTransforms(mwCtx, rawResult) const duration = Date.now() - startTime aiEventClient.emit('transcription:request:completed', { diff --git a/packages/ai/src/activities/middleware/index.ts b/packages/ai/src/activities/middleware/index.ts index 79454ac5c..d5123374d 100644 --- a/packages/ai/src/activities/middleware/index.ts +++ b/packages/ai/src/activities/middleware/index.ts @@ -9,6 +9,8 @@ export type { GenerationAbortInfo, GenerationErrorInfo, AnyGenerationMiddleware, + GenerationResultTransform, + GenerationResultTransformContext, } from './types' export { createGenerationContext, diff --git a/packages/ai/src/activities/middleware/run.ts b/packages/ai/src/activities/middleware/run.ts index 6983b1e55..a1793cda9 100644 --- a/packages/ai/src/activities/middleware/run.ts +++ b/packages/ai/src/activities/middleware/run.ts @@ -4,6 +4,7 @@ import type { GenerationFinishInfo, GenerationMiddleware, GenerationMiddlewareContext, + GenerationResultTransformContext, GenerationUsageInfo, } from './types' @@ -19,6 +20,9 @@ export function createGenerationContext(args: { provider: string model: string modelOptions?: unknown + threadId?: string + runId?: string + artifactInputs?: unknown createId: (prefix: string) => string }): GenerationMiddlewareContext { return { @@ -27,9 +31,13 @@ export function createGenerationContext(args: { provider: args.provider, model: args.model, modelOptions: args.modelOptions, + threadId: args.threadId, + runId: args.runId, source: 'server', createId: args.createId, context: undefined, + resultTransforms: [], + artifactInputs: args.artifactInputs, } } @@ -86,3 +94,26 @@ export function runGenerationError( ): Promise { return run(middleware, (mw) => mw.onError?.(ctx, info)) } + +/** + * Apply the result transforms middleware registered on the context, in order, + * to the raw adapter result. Each transform may return a replacement result or + * `undefined` to leave it unchanged. Runs after the adapter result exists and + * before the final result is returned or streamed. + */ +export async function applyGenerationResultTransforms( + ctx: GenerationMiddlewareContext, + result: TResult, +): Promise { + let current = result + const transformCtx: GenerationResultTransformContext = { middleware: ctx } + + for (const transform of ctx.resultTransforms ?? []) { + const transformed = await transform(current, transformCtx) + if (transformed !== undefined) { + current = transformed as TResult + } + } + + return current +} diff --git a/packages/ai/src/activities/middleware/types.ts b/packages/ai/src/activities/middleware/types.ts index 99ae3e44e..ef1136413 100644 --- a/packages/ai/src/activities/middleware/types.ts +++ b/packages/ai/src/activities/middleware/types.ts @@ -60,6 +60,10 @@ export interface GenerationMiddlewareContext { provider: string /** Model id. Emitted as `gen_ai.request.model`. */ model: string + /** Stable conversation/thread id, when supplied by the caller. */ + threadId?: string + /** Stable run id, when supplied by the caller. */ + runId?: string /** * Provider-specific options passed to the activity, if any. Typed `unknown` * because each activity's options are strongly typed per model; a supertype @@ -72,8 +76,36 @@ export interface GenerationMiddlewareContext { createId: (prefix: string) => string /** Runtime context provided by the activity options, if any. */ context: TContext + /** + * Result transforms registered by middleware during this activity call. + * Transforms run after the raw adapter result exists and before the final + * result is returned or streamed. Push multiple transforms to run them in + * registration order. + */ + resultTransforms?: Array> + /** + * Activity inputs captured for middleware that needs to transform or persist + * the result together with reconstructable request metadata. + */ + artifactInputs?: unknown } +/** Stable context handed to each {@link GenerationResultTransform}. */ +export interface GenerationResultTransformContext { + /** The activity call being transformed. */ + middleware: GenerationMiddlewareContext +} + +/** + * A transform middleware registers on `ctx.resultTransforms` to rewrite the raw + * adapter result before it is returned or streamed. Return a new result to + * replace it, or `undefined` to leave it unchanged. + */ +export type GenerationResultTransform = ( + result: TResult, + ctx: GenerationResultTransformContext, +) => TResult | undefined | Promise + // =========================== // Hook payloads // =========================== diff --git a/packages/ai/src/activities/stream-generation-result.ts b/packages/ai/src/activities/stream-generation-result.ts index 1bb530179..721b7a349 100644 --- a/packages/ai/src/activities/stream-generation-result.ts +++ b/packages/ai/src/activities/stream-generation-result.ts @@ -12,6 +12,19 @@ function createId(prefix: string): string { return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 9)}` } +/** + * Persisted artifact refs a middleware may have attached to the result. Read + * defensively: the result shape is activity-specific and `artifacts` is only + * present when generation persistence is wired with an artifact + blob store. + */ +function artifactsFromResult(result: unknown): Array | undefined { + if (typeof result !== 'object' || result === null) return undefined + const artifacts = (result as { artifacts?: unknown }).artifacts + return Array.isArray(artifacts) && artifacts.length > 0 + ? artifacts + : undefined +} + /** * Wrap a one-shot generation result as a StreamChunk async iterable. * @@ -23,7 +36,10 @@ function createId(prefix: string): string { * @returns An AsyncIterable of StreamChunks with RUN_STARTED, CUSTOM(generation:result), and RUN_FINISHED events on success, or RUN_STARTED and RUN_ERROR on failure */ export async function* streamGenerationResult( - generator: () => Promise, + generator: (resolved: { + runId: string + threadId: string + }) => Promise, options?: { runId?: string; threadId?: string }, ): AsyncIterable { const runId = options?.runId ?? createId('run') @@ -37,7 +53,19 @@ export async function* streamGenerationResult( } try { - const result = await generator() + const result = await generator({ runId, threadId }) + + // Emit persisted artifact refs (if a middleware attached any) before the + // result, so the client records them as the run streams. + const artifacts = artifactsFromResult(result) + if (artifacts) { + yield { + type: EventType.CUSTOM, + name: 'generation:artifacts', + value: artifacts, + timestamp: Date.now(), + } + } yield { type: EventType.CUSTOM, diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index c903e3e94..c263b85d3 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -202,6 +202,8 @@ export type { GenerationAbortInfo, GenerationErrorInfo, AnyGenerationMiddleware, + GenerationResultTransform, + GenerationResultTransformContext, } from './activities/middleware/index' // Capability primitives + middleware builder export { diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index caec892a8..cee0068c4 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -2075,6 +2075,10 @@ importers: version: 4.3.6 packages/ai-persistence: + dependencies: + '@tanstack/ai-utils': + specifier: workspace:* + version: link:../ai-utils devDependencies: '@tanstack/ai': specifier: workspace:* @@ -32223,8 +32227,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: {} @@ -32374,7 +32378,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