From 721c0235f649210107cde717d6fbc51c656a78b2 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Thu, 9 Jul 2026 18:48:53 +0200 Subject: [PATCH 1/9] =?UTF-8?q?feat(ai)!:=20split=20state=20vs=20delivery?= =?UTF-8?q?=20durability=20=E2=80=94=20StreamDurability=20transport,=20rem?= =?UTF-8?q?ove=20event-log/cursor/resume=20seam?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- packages/ai/src/activities/chat/index.ts | 55 --- .../src/activities/chat/middleware/types.ts | 1 - .../ai/src/activities/generateAudio/index.ts | 16 +- .../ai/src/activities/generateImage/index.ts | 14 +- .../ai/src/activities/generateSpeech/index.ts | 16 +- .../activities/generateTranscription/index.ts | 14 +- .../ai/src/activities/generateVideo/index.ts | 44 +-- packages/ai/src/activities/generation-run.ts | 26 -- packages/ai/src/activities/middleware/run.ts | 2 - .../ai/src/activities/middleware/types.ts | 8 - .../activities/stream-generation-result.ts | 12 +- packages/ai/src/client.ts | 7 +- packages/ai/src/custom-events.ts | 2 +- packages/ai/src/index.ts | 16 +- packages/ai/src/resume.ts | 46 --- packages/ai/src/stream-durability.ts | 329 ++++++++++++++++++ packages/ai/src/stream-to-response.ts | 220 +++++++++++- packages/ai/src/types.ts | 20 +- packages/ai/src/utilities/chat-params.ts | 25 -- packages/ai/tests/chat-params.test.ts | 12 +- packages/ai/tests/generation-params.test.ts | 13 +- .../generation-run-identity-replay.test.ts | 34 +- packages/ai/tests/resume-seam.test.ts | 194 ----------- packages/ai/tests/stream-durability.test.ts | 270 ++++++++++++++ .../stream-to-response-durability.test.ts | 194 +++++++++++ 25 files changed, 1071 insertions(+), 519 deletions(-) delete mode 100644 packages/ai/src/resume.ts create mode 100644 packages/ai/src/stream-durability.ts delete mode 100644 packages/ai/tests/resume-seam.test.ts create mode 100644 packages/ai/tests/stream-durability.test.ts create mode 100644 packages/ai/tests/stream-to-response-durability.test.ts diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index b530d3b46..b8d9341ef 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -11,7 +11,6 @@ import { stripToSpecMiddleware } from '../../strip-to-spec-middleware' import { streamToText } from '../../stream-to-response.js' import { resolveDebugOption } from '../../logger/resolve' import { EventType } from '../../types' -import { getResumeSource } from '../../resume' import { normalizeToolResult } from '../../utilities/tool-result' import { isProviderExecutedToolCall } from '../../utilities/provider-executed' import { LazyToolManager } from './tools/lazy-tool-manager' @@ -256,12 +255,6 @@ export interface TextActivityOptions< runId?: TextOptions['runId'] /** Parent run ID for AG-UI protocol nested run correlation. */ parentRunId?: TextOptions['parentRunId'] - /** - * Resume cursor. When provided with a resume source (e.g. via - * `withChatPersistence`), the engine replays persisted events after this cursor - * instead of running the adapter. A no-op when no resume source is present. - */ - cursor?: TextOptions['cursor'] /** * AG-UI interrupt resume responses. Persistence middleware validates these * before accepting new input on a thread with pending interrupts. @@ -558,8 +551,6 @@ class TextEngine< private readonly threadId: string private readonly runIdOverride?: string private readonly parentRunIdOverride?: string - /** Resume cursor supplied by the caller; drives the resume seam (no-op without a resume source). */ - private readonly cursorInput?: string // Middleware support private readonly middlewareRunner: MiddlewareRunner @@ -658,7 +649,6 @@ class TextEngine< this.createId('thread') this.runIdOverride = config.params.runId this.parentRunIdOverride = config.params.parentRunId - this.cursorInput = config.params.cursor // Initialize middleware — devtools first, strip-to-spec always last. // handleStreamChunk processes raw chunks BEFORE middleware, so internal @@ -798,41 +788,6 @@ class TextEngine< return this.finalizationError } - /** - * Resume seam. When the caller supplied a `cursor` and a `ResumeSource` was - * provided by middleware (e.g. `withChatPersistence`) for this run, replay the - * persisted event tail after the cursor and report that the run was handled - * by replay. Returns `false` (a no-op) when there is no cursor, no resume - * source, or no persisted run — so a normal run proceeds unchanged. - * - * Phase 1 is replay-only: it yields the persisted tail and ends. Live - * re-attach for harness adapters (continuing an in-sandbox process) layers on - * top of this in a later phase. - */ - private async *maybeResume(): AsyncGenerator { - const cursor = this.cursorInput - if (cursor === undefined) { - return false - } - if (this.params.resume !== undefined) { - return false - } - const source = getResumeSource(this.middlewareCtx, { optional: true }) - if (!source) { - return false - } - const runId = this.middlewareCtx.runId - if (!(await source.hasRun(runId))) { - return false - } - for await (const chunk of source.replay(runId, cursor)) { - yield chunk - } - // Durable resume is replay-only in core chat: replay public events and - // stop. Live harness/process re-entry is handled outside this seam. - return true - } - async *run(): AsyncGenerator { this.beforeRun() this.logger.agentLoop('run started', { @@ -855,15 +810,6 @@ class TextEngine< // Run onStart (devtools middleware emits text:request:started and initial messages here) await this.middlewareRunner.runOnStart(this.middlewareCtx) - // Resume seam: when a cursor was supplied AND a resume source is available - // for this run, replay the persisted event tail instead of running the - // adapter. No-op (falls through) when no resume source is provided, so a - // normal run is unaffected. - const resumed = yield* this.maybeResume() - if (resumed) { - return - } - const pendingPhase = yield* this.checkForPendingToolCalls() if (pendingPhase === 'wait') { return @@ -2614,7 +2560,6 @@ class TextEngine< messages: this.messages, systemPrompts: [...this.systemPrompts], tools: [...this.tools], - cursor: this.params.cursor, resume: this.params.resume, resumeToolState: { approvals: this.resumeApprovals, diff --git a/packages/ai/src/activities/chat/middleware/types.ts b/packages/ai/src/activities/chat/middleware/types.ts index 0f03e725e..741d2df45 100644 --- a/packages/ai/src/activities/chat/middleware/types.ts +++ b/packages/ai/src/activities/chat/middleware/types.ts @@ -209,7 +209,6 @@ export interface ChatMiddlewareConfig { messages: Array systemPrompts: Array tools: Array - cursor?: string | undefined resume?: Array | undefined resumeToolState?: ChatResumeToolState | undefined metadata?: Record | undefined diff --git a/packages/ai/src/activities/generateAudio/index.ts b/packages/ai/src/activities/generateAudio/index.ts index ee378e874..95eb28c05 100644 --- a/packages/ai/src/activities/generateAudio/index.ts +++ b/packages/ai/src/activities/generateAudio/index.ts @@ -8,8 +8,7 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { streamGenerationResult } from '../stream-generation-result.js' import { - createGenerationCursor, - generationEventFields, + generationIdentityFields, rejectEventsOnlyReplay, } from '../generation-run' import { resolveDebugOption } from '../../logger/resolve' @@ -169,7 +168,6 @@ async function runGenerateAudio< middleware, threadId, runId, - cursor, replay: _replay, ...rest } = options @@ -181,8 +179,7 @@ async function runGenerateAudio< (adapter as { name?: string; provider?: string }).provider ?? (adapter as { name?: string }).name ?? 'unknown' - const nextCursor = createGenerationCursor(cursor) - const identity = { threadId, runId, cursor } + const identity = { threadId, runId } const mwCtx = createGenerationContext({ requestId, @@ -192,7 +189,6 @@ async function runGenerateAudio< modelOptions: rest.modelOptions, threadId, runId, - cursor, artifactInputs: { prompt: rest.prompt, duration: rest.duration }, createId, }) @@ -201,7 +197,7 @@ async function runGenerateAudio< aiEventClient.emit('audio:request:started', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), provider: adapter.name, model, prompt: rest.prompt, @@ -222,7 +218,7 @@ async function runGenerateAudio< aiEventClient.emit('audio:request:completed', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), provider: adapter.name, model, audio: result.audio, @@ -234,7 +230,7 @@ async function runGenerateAudio< if (result.usage) { aiEventClient.emit('audio:usage', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), model, usage: result.usage, modelOptions: rest.modelOptions as Record | undefined, @@ -259,7 +255,7 @@ async function runGenerateAudio< const err = error as Error aiEventClient.emit('audio:request:error', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), provider: adapter.name, model, error: { message: err.message, name: err.name }, diff --git a/packages/ai/src/activities/generateImage/index.ts b/packages/ai/src/activities/generateImage/index.ts index 7be6af2aa..eecf00961 100644 --- a/packages/ai/src/activities/generateImage/index.ts +++ b/packages/ai/src/activities/generateImage/index.ts @@ -8,8 +8,7 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { streamGenerationResult } from '../stream-generation-result.js' import { - createGenerationCursor, - generationEventFields, + generationIdentityFields, rejectEventsOnlyReplay, } from '../generation-run' import { resolveDebugOption } from '../../logger/resolve' @@ -263,7 +262,6 @@ async function runGenerateImage< middleware, threadId, runId, - cursor, replay: _replay, ...rest } = options @@ -271,8 +269,7 @@ async function runGenerateImage< const requestId = createId('image') const startTime = Date.now() const logger: InternalLogger = resolveDebugOption(options.debug) - const nextCursor = createGenerationCursor(cursor) - const identity = { threadId, runId, cursor } + const identity = { threadId, runId } const mwCtx = createGenerationContext({ requestId, @@ -282,7 +279,6 @@ async function runGenerateImage< modelOptions: rest.modelOptions, threadId, runId, - cursor, artifactInputs: { prompt: rest.prompt }, createId, }) @@ -295,7 +291,7 @@ async function runGenerateImage< aiEventClient.emit('image:request:started', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), provider: adapter.name, model, prompt: resolved.text, @@ -326,7 +322,7 @@ async function runGenerateImage< aiEventClient.emit('image:request:completed', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), provider: adapter.name, model, // GeneratedImage is a discriminated `{ url } | { b64Json }` union, but the @@ -346,7 +342,7 @@ async function runGenerateImage< if (result.usage) { aiEventClient.emit('image:usage', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), model, usage: result.usage, modelOptions: rest.modelOptions, diff --git a/packages/ai/src/activities/generateSpeech/index.ts b/packages/ai/src/activities/generateSpeech/index.ts index 0aa4fc4fa..d4aa018bf 100644 --- a/packages/ai/src/activities/generateSpeech/index.ts +++ b/packages/ai/src/activities/generateSpeech/index.ts @@ -8,8 +8,7 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { streamGenerationResult } from '../stream-generation-result.js' import { - createGenerationCursor, - generationEventFields, + generationIdentityFields, rejectEventsOnlyReplay, } from '../generation-run' import { resolveDebugOption } from '../../logger/resolve' @@ -178,7 +177,6 @@ async function runGenerateSpeech< middleware, threadId, runId, - cursor, replay: _replay, ...rest } = options @@ -190,8 +188,7 @@ async function runGenerateSpeech< (adapter as { name?: string; provider?: string }).provider ?? (adapter as { name?: string }).name ?? 'unknown' - const nextCursor = createGenerationCursor(cursor) - const identity = { threadId, runId, cursor } + const identity = { threadId, runId } const mwCtx = createGenerationContext({ requestId, @@ -201,7 +198,6 @@ async function runGenerateSpeech< modelOptions: rest.modelOptions, threadId, runId, - cursor, artifactInputs: { text: rest.text, voice: rest.voice, @@ -215,7 +211,7 @@ async function runGenerateSpeech< aiEventClient.emit('speech:request:started', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), provider: adapter.name, model, text: rest.text, @@ -238,7 +234,7 @@ async function runGenerateSpeech< aiEventClient.emit('speech:request:completed', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), provider: adapter.name, model, audio: result.audio, @@ -253,7 +249,7 @@ async function runGenerateSpeech< if (result.usage) { aiEventClient.emit('speech:usage', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), model, usage: result.usage, modelOptions: rest.modelOptions as Record | undefined, @@ -278,7 +274,7 @@ async function runGenerateSpeech< const err = error as Error aiEventClient.emit('speech:request:error', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), provider: adapter.name, model, error: { message: err.message, name: err.name }, diff --git a/packages/ai/src/activities/generateTranscription/index.ts b/packages/ai/src/activities/generateTranscription/index.ts index fa50b4da6..0c3d3299c 100644 --- a/packages/ai/src/activities/generateTranscription/index.ts +++ b/packages/ai/src/activities/generateTranscription/index.ts @@ -8,8 +8,7 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { streamGenerationResult } from '../stream-generation-result.js' import { - createGenerationCursor, - generationEventFields, + generationIdentityFields, rejectEventsOnlyReplay, } from '../generation-run' import { resolveDebugOption } from '../../logger/resolve' @@ -213,7 +212,6 @@ async function runGenerateTranscription< middleware, threadId, runId, - cursor, replay: _replay, ...rest } = options @@ -225,8 +223,7 @@ async function runGenerateTranscription< (adapter as { name?: string; provider?: string }).provider ?? (adapter as { name?: string }).name ?? 'unknown' - const nextCursor = createGenerationCursor(cursor) - const identity = { threadId, runId, cursor } + const identity = { threadId, runId } const mwCtx = createGenerationContext({ requestId, @@ -236,7 +233,6 @@ async function runGenerateTranscription< modelOptions: rest.modelOptions, threadId, runId, - cursor, artifactInputs: { audio: rest.audio, language: rest.language, @@ -250,7 +246,7 @@ async function runGenerateTranscription< aiEventClient.emit('transcription:request:started', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), provider: adapter.name, model, language: rest.language, @@ -272,7 +268,7 @@ async function runGenerateTranscription< aiEventClient.emit('transcription:request:completed', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), provider: adapter.name, model, text: result.text, @@ -299,7 +295,7 @@ async function runGenerateTranscription< const err = error as Error aiEventClient.emit('transcription:request:error', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), provider: adapter.name, model, error: { message: err.message, name: err.name }, diff --git a/packages/ai/src/activities/generateVideo/index.ts b/packages/ai/src/activities/generateVideo/index.ts index bd4dff685..b6efd0f94 100644 --- a/packages/ai/src/activities/generateVideo/index.ts +++ b/packages/ai/src/activities/generateVideo/index.ts @@ -10,8 +10,7 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { toRunErrorPayload } from '../error-payload' import { - createGenerationCursor, - generationEventFields, + generationIdentityFields, rejectEventsOnlyReplay, replayGenerationEvents, } from '../generation-run' @@ -341,7 +340,6 @@ async function runCreateVideoJob< middleware, threadId, runId, - cursor, } = options const model = adapter.model const requestId = createId('video') @@ -360,7 +358,6 @@ async function runCreateVideoJob< modelOptions, threadId, runId, - cursor, artifactInputs: { prompt, size, duration }, createId, }) @@ -435,14 +432,13 @@ async function* runStreamingVideoGeneration< 'unknown' const threadId = options.threadId ?? createId('thread') - const nextCursor = createGenerationCursor(options.cursor) - const identity = { threadId, runId, cursor: options.cursor } + const identity = { threadId, runId } yield { type: 'RUN_STARTED', runId, threadId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), timestamp: Date.now(), } as StreamChunk @@ -455,7 +451,7 @@ async function* runStreamingVideoGeneration< type: 'CUSTOM', name: 'generation:artifacts', value: replayResult.artifacts, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), timestamp: Date.now(), } as StreamChunk } @@ -464,7 +460,7 @@ async function* runStreamingVideoGeneration< type: 'CUSTOM', name: 'generation:result', value: options.replay.result, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), timestamp: Date.now(), } as StreamChunk @@ -473,7 +469,7 @@ async function* runStreamingVideoGeneration< runId, threadId, finishReason: 'stop', - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), timestamp: Date.now(), } as StreamChunk return @@ -487,7 +483,6 @@ async function* runStreamingVideoGeneration< modelOptions, threadId, runId, - cursor: options.cursor, artifactInputs: { prompt, size, duration }, createId, }) @@ -520,7 +515,7 @@ async function* runStreamingVideoGeneration< type: 'CUSTOM', name: 'video:job:created', value: { jobId: jobResult.jobId }, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), timestamp: Date.now(), } as StreamChunk @@ -540,7 +535,7 @@ async function* runStreamingVideoGeneration< progress: statusResult.progress, error: statusResult.error, }, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), timestamp: Date.now(), } as StreamChunk @@ -576,7 +571,7 @@ async function* runStreamingVideoGeneration< type: 'CUSTOM', name: 'generation:artifacts', value: urlResult.artifacts, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), timestamp: Date.now(), } as StreamChunk } @@ -588,7 +583,7 @@ async function* runStreamingVideoGeneration< ...urlResult, status: 'completed', }, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), timestamp: Date.now(), } as StreamChunk @@ -597,7 +592,7 @@ async function* runStreamingVideoGeneration< runId, threadId, finishReason: 'stop', - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), timestamp: Date.now(), } as StreamChunk return @@ -628,7 +623,7 @@ async function* runStreamingVideoGeneration< type: 'RUN_ERROR', runId, threadId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), message: payload.message, code: payload.code, error: payload, @@ -687,15 +682,14 @@ export async function getVideoJobStatus< error?: string usage?: TokenUsage }> { - const { adapter, jobId, threadId, runId, cursor } = options + const { adapter, jobId, threadId, runId } = options const requestId = createId('video-status') const startTime = Date.now() - const nextCursor = createGenerationCursor(cursor) - const identity = { threadId, runId, cursor } + const identity = { threadId, runId } aiEventClient.emit('video:request:started', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), provider: adapter.name, model: adapter.model, requestType: 'status', @@ -712,7 +706,7 @@ export async function getVideoJobStatus< const urlResult = await adapter.getVideoUrl(jobId) aiEventClient.emit('video:request:completed', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), provider: adapter.name, model: adapter.model, requestType: 'status', @@ -726,7 +720,7 @@ export async function getVideoJobStatus< if (urlResult.usage) { aiEventClient.emit('video:usage', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), model: adapter.model, usage: urlResult.usage, timestamp: Date.now(), @@ -743,7 +737,7 @@ export async function getVideoJobStatus< error instanceof Error ? error.message : 'Failed to get video URL' aiEventClient.emit('video:request:completed', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), provider: adapter.name, model: adapter.model, requestType: 'status', @@ -765,7 +759,7 @@ export async function getVideoJobStatus< aiEventClient.emit('video:request:completed', { requestId, - ...generationEventFields(identity, nextCursor), + ...generationIdentityFields(identity), provider: adapter.name, model: adapter.model, requestType: 'status', diff --git a/packages/ai/src/activities/generation-run.ts b/packages/ai/src/activities/generation-run.ts index fa6c61149..328363f52 100644 --- a/packages/ai/src/activities/generation-run.ts +++ b/packages/ai/src/activities/generation-run.ts @@ -11,17 +11,6 @@ export function rejectEventsOnlyReplay( } } -export function createGenerationCursor( - baseCursor: string | undefined, -): () => string | undefined { - let index = 0 - return () => { - if (baseCursor === undefined) return undefined - index += 1 - return `${baseCursor}:${index}` - } -} - export function generationIdentityFields(identity: GenerationRunIdentity): { threadId?: string runId?: string @@ -32,21 +21,6 @@ export function generationIdentityFields(identity: GenerationRunIdentity): { } } -export function generationEventFields( - identity: GenerationRunIdentity, - nextCursor: () => string | undefined, -): { - threadId?: string - runId?: string - cursor?: string -} { - const cursor = nextCursor() - return { - ...generationIdentityFields(identity), - ...(cursor !== undefined ? { cursor } : {}), - } -} - export async function* replayGenerationEvents( replay: GenerationReplayInput, ): AsyncIterable { diff --git a/packages/ai/src/activities/middleware/run.ts b/packages/ai/src/activities/middleware/run.ts index f799e7473..4ae63db3e 100644 --- a/packages/ai/src/activities/middleware/run.ts +++ b/packages/ai/src/activities/middleware/run.ts @@ -22,7 +22,6 @@ export function createGenerationContext(args: { modelOptions?: unknown threadId?: string runId?: string - cursor?: string artifactInputs?: unknown createId: (prefix: string) => string }): GenerationMiddlewareContext { @@ -34,7 +33,6 @@ export function createGenerationContext(args: { modelOptions: args.modelOptions, threadId: args.threadId, runId: args.runId, - cursor: args.cursor, source: 'server', createId: args.createId, context: undefined, diff --git a/packages/ai/src/activities/middleware/types.ts b/packages/ai/src/activities/middleware/types.ts index df9d0ae48..7f5b9ffc4 100644 --- a/packages/ai/src/activities/middleware/types.ts +++ b/packages/ai/src/activities/middleware/types.ts @@ -45,12 +45,6 @@ export interface GenerationRunIdentity { threadId?: string /** Stable run id for correlating generation events. */ runId?: string - /** - * Opaque resume cursor supplied by the caller. Generation activities use this - * as the base for per-event cursors when no persistence middleware stamps - * cursors itself. - */ - cursor?: string } export interface GenerationReplayInput { @@ -100,8 +94,6 @@ export interface GenerationMiddlewareContext { threadId?: string /** Stable run id, when supplied by the caller. */ runId?: string - /** Caller-supplied resume cursor, when supplied. */ - cursor?: string /** * Provider-specific options passed to the activity, if any. Typed `unknown` * because each activity's options are strongly typed per model; a supertype diff --git a/packages/ai/src/activities/stream-generation-result.ts b/packages/ai/src/activities/stream-generation-result.ts index 44bca527a..ea01ba2d9 100644 --- a/packages/ai/src/activities/stream-generation-result.ts +++ b/packages/ai/src/activities/stream-generation-result.ts @@ -7,8 +7,6 @@ import { EventType } from '@ag-ui/core' import { toRunErrorPayload } from './error-payload' import { - createGenerationCursor, - generationEventFields, generationIdentityFields, replayGenerationEvents, } from './generation-run' @@ -47,13 +45,12 @@ export async function* streamGenerationResult( const threadId = options?.threadId ?? createId('thread') const identity = { runId, threadId } const resolvedOptions = { ...options, ...identity } - const nextCursor = createGenerationCursor(options?.cursor) yield { type: EventType.RUN_STARTED, runId, threadId, - ...generationEventFields(resolvedOptions, nextCursor), + ...generationIdentityFields(resolvedOptions), timestamp: Date.now(), } @@ -69,7 +66,7 @@ export async function* streamGenerationResult( type: EventType.CUSTOM, name: 'generation:artifacts', value: artifacts, - ...generationEventFields(resolvedOptions, nextCursor), + ...generationIdentityFields(resolvedOptions), timestamp: Date.now(), } } @@ -78,7 +75,7 @@ export async function* streamGenerationResult( type: EventType.CUSTOM, name: 'generation:result', value: result as unknown, - ...generationEventFields(resolvedOptions, nextCursor), + ...generationIdentityFields(resolvedOptions), timestamp: Date.now(), } @@ -87,7 +84,7 @@ export async function* streamGenerationResult( runId, threadId, finishReason: 'stop', - ...generationEventFields(resolvedOptions, nextCursor), + ...generationIdentityFields(resolvedOptions), timestamp: Date.now(), } } catch (error: unknown) { @@ -103,7 +100,6 @@ export async function* streamGenerationResult( ...generationIdentityFields(identity), message: payload.message, ...codeFields, - ...generationEventFields(resolvedOptions, nextCursor), // Deprecated nested form for backward compatibility error: { message: payload.message, diff --git a/packages/ai/src/client.ts b/packages/ai/src/client.ts index 895fbb0c7..1a9d8e623 100644 --- a/packages/ai/src/client.ts +++ b/packages/ai/src/client.ts @@ -26,7 +26,6 @@ export interface GenerationParams { forwardedProps: Record threadId?: string runId?: string - cursor?: string } const generationKinds = [ @@ -103,7 +102,7 @@ function forwardedPropsFromEnvelope( function optionalStringField( envelope: Record, - key: 'threadId' | 'runId' | 'cursor', + key: 'threadId' | 'runId', ): string | undefined { if (!hasOwnKey(envelope, key)) { return undefined @@ -120,20 +119,16 @@ function optionalStringField( function generationIdentityFields(envelope: Record): { threadId?: string runId?: string - cursor?: string } { const identity: { threadId?: string runId?: string - cursor?: string } = {} const threadId = optionalStringField(envelope, 'threadId') const runId = optionalStringField(envelope, 'runId') - const cursor = optionalStringField(envelope, 'cursor') if (threadId !== undefined) identity.threadId = threadId if (runId !== undefined) identity.runId = runId - if (cursor !== undefined) identity.cursor = cursor return identity } diff --git a/packages/ai/src/custom-events.ts b/packages/ai/src/custom-events.ts index 7829924ab..78905434d 100644 --- a/packages/ai/src/custom-events.ts +++ b/packages/ai/src/custom-events.ts @@ -103,5 +103,5 @@ export function isCustomEvent( chunk: StreamChunk, name: TName, ): chunk is WellKnownCustomEvent { - return chunk.type === EventType.CUSTOM && (chunk as CustomEvent).name === name + return chunk.type === EventType.CUSTOM && chunk.name === name } diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 58819880e..085d9139d 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -79,6 +79,14 @@ export { toHttpResponse, } from './stream-to-response' +// Delivery durability (transport layer) +export { + memoryStream, + encodeOffset, + decodeOffset, +} from './stream-durability' +export type { StreamDurability } from './stream-durability' + // Tool call management export { ToolCallManager } from './activities/chat/tools/tool-calls' @@ -172,14 +180,6 @@ export { } from './locks' export type { LockStore } from './locks' -// Resume-source capability (consumed by chat()'s resume seam; provided by persistence) -export { - ResumeSourceCapability, - getResumeSource, - provideResumeSource, -} from './resume' -export type { ResumeSource, RunStatus } from './resume' - // Well-known AG-UI CUSTOM event catalog (agent activity rides on CUSTOM events) export { CUSTOM_EVENT, isCustomEvent } from './custom-events' export type { diff --git a/packages/ai/src/resume.ts b/packages/ai/src/resume.ts deleted file mode 100644 index 9b4030e4f..000000000 --- a/packages/ai/src/resume.ts +++ /dev/null @@ -1,46 +0,0 @@ -/** - * Resume-source capability — the core seam that lets a reconnecting client pick - * up a run where it left off. - * - * The contract lives in core (not in `@tanstack/ai-persistence`) because the - * `chat()` engine itself owns the resume decision: when a `cursor` is supplied - * and a `ResumeSource` has been provided by middleware, the engine replays the - * persisted event tail instead of running the adapter again. `withChatPersistence` - * PROVIDES a `ResumeSource` backed by its event log + run store; without it, the - * `cursor` option is a silent no-op (a normal run is unaffected). - * - * Core depends only on this small read contract — never on the persistence - * package — mirroring how {@link LocksCapability} lives in core. - */ -import { createCapability } from './activities/chat/middleware/capabilities' -import type { StreamChunk } from './types' - -/** Lifecycle status of a persisted run. */ -export type RunStatus = 'running' | 'completed' | 'failed' | 'interrupted' - -/** - * A read-only view of persisted run history sufficient for the engine to resume. - * Provided by `withChatPersistence`; consumed by the `chat()` resume seam. - */ -export interface ResumeSource { - /** Whether any events have been persisted for `runId`. */ - hasRun: (runId: string) => Promise - /** - * Replay persisted chunks for `runId` whose cursor is strictly after - * `afterCursor` (or from the beginning when omitted). Each yielded chunk - * carries its stamped `cursor`. - */ - replay: (runId: string, afterCursor?: string) => AsyncIterable - /** Current status of `runId`, or null when unknown. */ - getStatus: (runId: string) => Promise -} - -/** - * The resume capability. PROVIDED by `withChatPersistence`; OPTIONALLY consumed by - * the chat engine. When absent, a supplied `cursor` is ignored. - */ -export const ResumeSourceCapability = - createCapability()('resume-source') - -/** Destructured accessors: `getResumeSource(ctx)` / `provideResumeSource(ctx, src)`. */ -export const [getResumeSource, provideResumeSource] = ResumeSourceCapability diff --git a/packages/ai/src/stream-durability.ts b/packages/ai/src/stream-durability.ts new file mode 100644 index 000000000..6847619b9 --- /dev/null +++ b/packages/ai/src/stream-durability.ts @@ -0,0 +1,329 @@ +import type { StreamChunk } from './types' + +/** + * A pluggable sink for **delivery durability** — the transport-layer concern of + * letting a client disconnect, reload, or open a second tab and still receive + * the full, ordered run stream exactly once. + * + * This is deliberately thin: append / read / resume-detect and nothing else. It + * stores no run status, no messages, no interrupts — those are *state* + * durability (the middleware layer). Keeping this interface small is what stops + * delivery durability drifting back into "a transport wearing a store costume". + * + * Adapters close over the incoming `Request` (that is why they are constructed + * as `memoryStream(request)` / `durableStream(request, opts)`), so `request` + * never pollutes the transport-helper signature. + * + * Offsets are opaque strings of the form `runId@seq` (see {@link encodeOffset}). + * The sentinels `-1` (from the start) and `now` (tail — only future writes) may + * appear in the seq position. + */ +export interface StreamDurability { + /** + * The offset to resume from, read off the captured request — the + * `Last-Event-ID` header (native `EventSource` reconnect) or, failing that, + * the `?offset` query param. Returns `null` for a fresh run (produce path). + */ + resumeFrom: () => string | null + /** + * The stable run/stream id for this request. Minted (`crypto.randomUUID()`) + * for a fresh run; parsed from the resume offset / `?runId` param otherwise. + * Memoized so repeated calls return the same id. + */ + runId: () => string + /** + * Append a batch of chunks. The first chunk takes sequence number `startSeq`, + * the next `startSeq + 1`, and so on (sequence numbers are 1-based). + * + * Resolves to a per-chunk array of the **resume offsets the transport should + * tag each chunk with** (its client-facing SSE `id:`), aligned position-for- + * position with `chunks`. An element is `undefined` when that chunk carries + * no resumable offset. + * + * The offsets MUST be the offsets the backend will actually accept on a + * `read(offset)` resume — never a transport-local counter — so a reconnect + * resumes against the real backend key space: + * + * - {@link memoryStream} tags **every** chunk (`runId@seq`), because its log + * is per-chunk addressable. + * - `durableStream` also tags **every** chunk: it POSTs chunks one-by-one and + * reads back each chunk's own backend offset, so the returned array is fully + * populated (no `undefined`) and a mid-batch reconnect resumes at the exact + * chunk it dropped on — exactly-once at any batch size. + * + * A backend that could only learn a coarser (e.g. per-batch) offset would + * return `undefined` for the untagged chunks; the client then relies on + * id-keyed de-dup for a partial replay. Prefer per-chunk tagging so exactly- + * once does not depend on client de-dup. + */ + append: ( + chunks: Array, + startSeq: number, + ) => Promise> + /** + * Replay the chunks strictly after `offset`. Pass `-1` (or an offset whose + * seq is `-1`) to replay from the start. Yields `{ seq, chunk }` in order. + * + * An optional `signal` bounds a live-tailing read: when it aborts, the read + * stops waiting for further appends and ends. Without it, a join to a runId + * with no in-process producer (a crash, a cross-instance run, or a bogus id) + * would park forever. + */ + read: ( + offset: string, + signal?: AbortSignal, + ) => AsyncIterable<{ seq: number; chunk: StreamChunk }> + /** + * Mark the run terminated so a concurrently live-tailing {@link read} (a + * mid-stream join / reconnect) stops waiting for further appends and returns. + * Called by the producer when the source stream ends — whether it ended with + * a terminal event, a thrown error, or simply ran dry. Optional: backends + * that detect completion from the stored terminal event (or their own + * close signal) need not implement it. + */ + markComplete?: () => void +} + +/** + * Encode a durability offset. The format is `runId@seq`; the seq may be a + * 1-based sequence number, or the sentinel `-1` (from start) / `now` (tail). + */ +export function encodeOffset(runId: string, seq: number | 'now'): string { + return `${runId}@${seq}` +} + +/** + * The seq token of an offset — the substring after the LAST `@`, or the whole + * string when there is no `@` (a bare seq token). + */ +function seqTokenOf(offset: string): string { + return offset.includes('@') + ? offset.slice(offset.lastIndexOf('@') + 1) + : offset +} + +/** + * Parse a seq token into its numeric value. The sentinels agree with + * {@link seqThreshold}: `-1` / `''` (from start) → `-1`, `now` (tail) → + * `+Infinity`. Throws only on a genuinely malformed (non-sentinel, + * non-numeric) token — never on a sentinel. + */ +function parseSeqToken(token: string): number { + if (token === '' || token === '-1') return -1 + if (token === 'now') return Number.POSITIVE_INFINITY + const n = Number(token) + if (!Number.isFinite(n)) { + throw new Error(`Invalid durability offset seq: ${token}`) + } + return n +} + +/** + * The run id of a `runId@seq` offset — everything before the LAST `@`, so run + * ids that themselves contain `@` round-trip. Deliberately does NOT parse the + * seq, so the sentinels `now` / `-1` (and any future seq form) never make run + * id extraction throw. When there is no `@`, the whole string is the run id. + */ +export function runIdOf(offset: string): string { + const at = offset.lastIndexOf('@') + return at === -1 ? offset : offset.slice(0, at) +} + +/** + * Decode a `runId@seq` offset. Splits on the LAST `@` so run ids that + * themselves contain `@` round-trip. The seq is parsed via the shared sentinel + * logic, so `runId@now` (→ `+Infinity`) and `runId@-1` (→ `-1`) decode without + * throwing; only a missing `@` or a non-sentinel non-numeric seq throws. + */ +export function decodeOffset(offset: string): { runId: string; seq: number } { + const at = offset.lastIndexOf('@') + if (at === -1) { + throw new Error(`Invalid durability offset (missing @): ${offset}`) + } + return { runId: offset.slice(0, at), seq: parseSeqToken(offset.slice(at + 1)) } +} + +/** + * Resolve the seq threshold a `read(offset)` should replay strictly after. + * Understands full `runId@seq` offsets and bare seq tokens, plus the sentinels + * `-1` (from start; threshold `-1` so every 1-based seq qualifies) and `now` + * (tail → `+Infinity`, so no already-stored chunk qualifies). + */ +function seqThreshold(offset: string): number { + return parseSeqToken(seqTokenOf(offset)) +} + +/** + * Reject a run id that would break SSE framing if it flowed into an `id:` line. + * A caller-supplied `?runId` (or a reconnect `Last-Event-ID`) is untrusted; a + * CR/LF in it could inject or split SSE events once it is emitted as + * `id: {runId}@{seq}`. Throws rather than silently sanitizing. + */ +function assertValidRunId(runId: string): string { + if (/[\r\n]/.test(runId)) { + throw new Error( + `Invalid runId (must not contain CR/LF): ${JSON.stringify(runId)}`, + ) + } + return runId +} + +function readResumeOffset(request: Request): string | null { + const header = request.headers.get('Last-Event-ID') + if (header) return header + try { + const param = new URL(request.url).searchParams.get('offset') + if (param) return param + } catch { + // A relative / opaque URL has no query string to parse — treat as fresh. + } + return null +} + +function runIdFor(request: Request, resumeOffset: string | null): string { + if (resumeOffset && resumeOffset.includes('@')) { + // Extract without parsing the seq, so a `runId@now` / `runId@-1` sentinel + // resume offset yields the run id instead of throwing. + return assertValidRunId(runIdOf(resumeOffset)) + } + try { + const param = new URL(request.url).searchParams.get('runId') + if (param) return assertValidRunId(param) + } catch (err) { + // Re-throw a CR/LF rejection; swallow only the URL-parse failure (a + // relative / opaque URL has no query string) and fall through to minting. + if (err instanceof Error && err.message.startsWith('Invalid runId')) { + throw err + } + } + return crypto.randomUUID() +} + +/** Terminal chunk types after which a run's delivery log is complete. */ +function isTerminalChunk(chunk: StreamChunk): boolean { + return chunk.type === 'RUN_FINISHED' || chunk.type === 'RUN_ERROR' +} + +/** + * One run's in-process delivery log: the ordered entries, a `complete` flag + * (set once the run terminates), and the waiters a live-tailing `read` parks on + * until the next append or completion. + */ +interface MemoryLog { + entries: Array<{ seq: number; chunk: StreamChunk }> + complete: boolean + waiters: Array<() => void> +} + +/** + * Process-local delivery log, shared across requests so a second tab / reconnect + * hitting the same server instance can replay an in-flight or finished run. + * + * ponytail: process-local map, dev/test backend only. It is NOT durable across + * restarts and NOT shared across instances — production uses `durableStream` + * (the durable-streams protocol adapter) instead. + */ +const memoryLogs = new Map() + +function getOrCreateLog(id: string): MemoryLog { + let log = memoryLogs.get(id) + if (!log) { + log = { entries: [], complete: false, waiters: [] } + memoryLogs.set(id, log) + } + return log +} + +/** Wake every waiter parked on a live-tailing read of this log. */ +function wakeWaiters(log: MemoryLog): void { + const waiters = log.waiters + log.waiters = [] + for (const wake of waiters) wake() +} + +/** + * The zero-infrastructure {@link StreamDurability} backend: an in-process log. + * + * Perfect for local dev and tests — a reconnect or second tab that reaches the + * same process replays the ordered stream. It does NOT survive a restart and is + * NOT shared across server instances; swap in `durableStream` from + * `@tanstack/ai-durable-stream` for production. + */ +export function memoryStream(request: Request): StreamDurability { + const resumeOffset = readResumeOffset(request) + let cachedRunId: string | undefined + + const runId = (): string => { + if (cachedRunId === undefined) { + cachedRunId = runIdFor(request, resumeOffset) + } + return cachedRunId + } + + return { + resumeFrom: () => resumeOffset, + runId, + append: async (chunks, startSeq) => { + const id = runId() + const log = getOrCreateLog(id) + // Every chunk is per-chunk addressable in this log's own key space, so + // tag each with its `runId@seq` offset — that IS the offset a resumer + // passes back to `read`. + const offsets = chunks.map((chunk, i) => { + const seq = startSeq + i + log.entries.push({ seq, chunk }) + if (isTerminalChunk(chunk)) log.complete = true + return encodeOffset(id, seq) + }) + wakeWaiters(log) + return offsets + }, + markComplete: () => { + const log = getOrCreateLog(runId()) + log.complete = true + wakeWaiters(log) + }, + read: async function* (offset, signal) { + const id = runId() + const threshold = seqThreshold(offset) + const log = getOrCreateLog(id) + // Live-tail: drain everything currently past the threshold, then park on + // the next append until a terminal chunk is read or the run is marked + // complete. Without this, a mid-stream join returns at the current tail + // BEFORE the run finishes, truncating the joiner's view. + let index = 0 + for (;;) { + while (index < log.entries.length) { + const entry = log.entries[index] + index += 1 + if (entry && entry.seq > threshold) { + yield entry + if (isTerminalChunk(entry.chunk)) return + } + } + if (log.complete) return + // A consumer that already aborted must not park. + if (signal?.aborted) return + // Park until the next append / completion, OR the consumer aborts. + // Threading the abort wake is what stops a join to a runId with no + // in-process producer (crash, cross-instance, or bogus id) — whose log + // is created empty here and never gets a `markComplete` — from parking + // forever: the caller bounds it with its request/consumer signal. + await new Promise((resolve) => { + const onAbort = () => { + const i = log.waiters.indexOf(wake) + if (i !== -1) log.waiters.splice(i, 1) + resolve() + } + const wake = () => { + signal?.removeEventListener('abort', onAbort) + resolve() + } + log.waiters.push(wake) + signal?.addEventListener('abort', onAbort) + }) + if (signal?.aborted) return + } + }, + } +} diff --git a/packages/ai/src/stream-to-response.ts b/packages/ai/src/stream-to-response.ts index 9850f4d60..a1a10c41e 100644 --- a/packages/ai/src/stream-to-response.ts +++ b/packages/ai/src/stream-to-response.ts @@ -1,4 +1,6 @@ import { toRunErrorPayload } from './activities/error-payload' +import { encodeOffset } from './stream-durability' +import type { StreamDurability } from './stream-durability' import type { StreamChunk } from './types' /** @@ -50,21 +52,31 @@ export async function streamToText( export function toServerSentEventsStream( stream: AsyncIterable, abortController?: AbortController, + getId?: (chunk: StreamChunk, index: number) => string | undefined, ): ReadableStream { const encoder = new TextEncoder() return new ReadableStream({ async start(controller) { try { + let index = 0 for await (const chunk of stream) { // Check if stream was cancelled/aborted if (abortController?.signal.aborted) { break } + // Tag the event with an `id:` line when a resolver is supplied + // (delivery durability offsets). Native `EventSource` echoes the + // last `id:` back as `Last-Event-ID` on reconnect, which is how + // resume works with zero client-side cursor state. + const id = getId?.(chunk, index) + index += 1 + const idLine = id !== undefined ? `id: ${id}\n` : '' + // Send each chunk as Server-Sent Events format controller.enqueue( - encoder.encode(`data: ${JSON.stringify(chunk)}\n\n`), + encoder.encode(`${idLine}data: ${JSON.stringify(chunk)}\n\n`), ) } @@ -99,6 +111,148 @@ export function toServerSentEventsStream( }) } +/** Default number of chunks buffered before a durability `append`. */ +const DEFAULT_DURABILITY_BATCH = 32 + +/** + * Resolve and validate the durability batch size. A non-positive-integer (0, + * negative, fractional, or `NaN`) is rejected rather than clamped: silently + * `Math.max(1, …)`-ing a `NaN` used to disable size-based flushing entirely + * (`length >= NaN` is always false), which is a subtle footgun. + */ +function resolveBatchSize(batch: number | undefined): number { + if (batch === undefined) return DEFAULT_DURABILITY_BATCH + if (!Number.isInteger(batch) || batch <= 0) { + throw new Error( + `Invalid durability batch size: ${batch}. Must be a positive integer.`, + ) + } + return batch +} + +/** + * Boundaries at which the batching producer flushes early, regardless of the + * batch size — terminal events and tool-call ends. Flushing here keeps the + * durability log promptly consistent at semantically meaningful points. + */ +function isDurabilityFlushBoundary(chunk: StreamChunk): boolean { + return ( + chunk.type === 'RUN_FINISHED' || + chunk.type === 'RUN_ERROR' || + chunk.type === 'TOOL_CALL_END' + ) +} + +/** + * Build the delivery-durable source iterable for a transport helper. + * + * - **Resume** (`resumeFrom()` non-null): replay strictly after the offset, + * reading only from the durability log. The input `stream` is NEVER iterated, + * so `chat()`'s lazy iterator never fires the provider — the untouched + * generator is simply GC'd. This is what makes resume free of re-invocation. + * - **Fresh** (`resumeFrom()` null): iterate `stream`, buffering up to `batch` + * chunks (flushing early at terminal / tool-call boundaries), `append` each + * batch to the log, then forward. Appending BEFORE forwarding guarantees a + * reconnecting client can always replay exactly what it already saw. + * + * The returned `getId` maps each forwarded chunk to its `runId@seq` offset, for + * the SSE `id:` line (ignored by the ndjson HTTP helper). + */ +function durableStreamSource( + stream: AsyncIterable, + durability: StreamDurability, + options: { abortController?: AbortController; batch?: number }, +): { + source: AsyncIterable + getId: (chunk: StreamChunk) => string | undefined +} { + const runId = durability.runId() + const resumeOffset = durability.resumeFrom() + const batchSize = resolveBatchSize(options.batch) + const abortController = options.abortController + const idByChunk = new WeakMap() + const getId = (chunk: StreamChunk): string | undefined => + idByChunk.get(chunk) + + async function* produce(): AsyncIterable { + let seq = 1 + let batch: Array = [] + let batchStart = 1 + + async function* flush(): AsyncIterable { + if (batch.length === 0) return + const toForward = batch + batch = [] + // Tag each forwarded chunk's client-facing SSE `id:` with the offset the + // BACKEND returned for it — never a transport-local counter. On resume + // the backend derives seq from this id, so the two must be the same key + // space. Both bundled backends tag every chunk (`memoryStream` per-chunk; + // `durableStream` POSTs chunks one-by-one for a per-chunk offset), so a + // mid-batch reconnect resumes exactly-once. A backend that returned + // `undefined` for some chunks would leave those without an `id:`. + const offsets = await durability.append(toForward, batchStart) + toForward.forEach((chunk, i) => { + const offset = offsets[i] + if (offset !== undefined) idByChunk.set(chunk, offset) + }) + for (const chunk of toForward) yield chunk + } + + try { + for await (const chunk of stream) { + if (abortController?.signal.aborted) break + if (batch.length === 0) batchStart = seq + batch.push(chunk) + seq += 1 + if (batch.length >= batchSize || isDurabilityFlushBoundary(chunk)) { + yield* flush() + } + } + yield* flush() + } catch (error) { + // The provider stream threw. Persist a terminal RUN_ERROR to the + // durability log so a resumer / joiner learns the run failed (otherwise + // the log ends with no terminal and they wait forever). Flush any + // buffered chunks first, then append the terminal WITHOUT forwarding it + // live — the transport layer synthesizes the live RUN_ERROR on rethrow, + // so forwarding here too would double-emit. + yield* flush() + if (!abortController?.signal.aborted) { + const errorChunk = { + type: 'RUN_ERROR', + timestamp: Date.now(), + error: toRunErrorPayload(error), + } as StreamChunk + await durability.append([errorChunk], seq) + } + throw error + } finally { + // Unblock any live-tailing join, however the run ended (terminal, thrown, + // or a source that simply ran dry with no terminal event). + durability.markComplete?.() + } + } + + async function* replay(offset: string): AsyncIterable { + // Thread the consumer's abort signal into the read so a live-tailing join + // (a mid-stream reconnect) that is aborted — or that hit a runId with no + // in-process producer — stops parking and ends instead of hanging forever. + for await (const { seq, chunk } of durability.read( + offset, + abortController?.signal, + )) { + if (abortController?.signal.aborted) break + idByChunk.set(chunk, encodeOffset(runId, seq)) + yield chunk + } + } + + return { + source: resumeOffset !== null ? replay(resumeOffset) : produce(), + getId, + } +} + /** * Convert a StreamChunk async iterable to a Response in Server-Sent Events format * @@ -107,21 +261,32 @@ export function toServerSentEventsStream( * - Each chunk is followed by "\n\n" * - Stream ends when the underlying iterable is exhausted (RUN_FINISHED is the terminal event) * + * Pass a `durability` sink (`memoryStream(request)` / `durableStream(request)`) + * to make the stream resumable: fresh runs are appended to the log and each SSE + * event is tagged with an `id:` offset; a reconnect (native `Last-Event-ID`) or + * a `?offset` join replays from the log without re-running the producer. `batch` + * controls how many chunks are buffered per `append` (default 32). + * * @param stream - AsyncIterable of StreamChunks from chat() - * @param init - Optional Response initialization options (including `abortController`) + * @param init - Optional Response initialization options (including `abortController`, `durability`, `batch`) * @returns Response in Server-Sent Events format * * @example * ```typescript - * const stream = chat({ adapter: openaiText(), model: "gpt-4o", messages: [...] }); - * return toServerSentEventsResponse(stream, { abortController }); + * const stream = chat({ adapter: openaiText(), model: "gpt-5.5", messages: [...] }); + * return toServerSentEventsResponse(stream, { durability: memoryStream(request) }); * ``` */ export function toServerSentEventsResponse( stream: AsyncIterable, - init?: ResponseInit & { abortController?: AbortController }, + init?: ResponseInit & { + abortController?: AbortController + durability?: StreamDurability + batch?: number + }, ): Response { - const { headers, abortController, ...responseInit } = init ?? {} + const { headers, abortController, durability, batch, ...responseInit } = + init ?? {} // Start with default SSE headers const mergedHeaders = new Headers({ @@ -139,7 +304,18 @@ export function toServerSentEventsResponse( }) } - return new Response(toServerSentEventsStream(stream, abortController), { + let body: ReadableStream + if (durability) { + const { source, getId } = durableStreamSource(stream, durability, { + abortController, + batch, + }) + body = toServerSentEventsStream(source, abortController, getId) + } else { + body = toServerSentEventsStream(stream, abortController) + } + + return new Response(body, { ...responseInit, headers: mergedHeaders, }) @@ -227,21 +403,39 @@ export function toHttpStream( * * This format is compatible with `fetchHttpStream` connection adapter. * + * Pass a `durability` sink to make the stream resumable (same semantics as + * {@link toServerSentEventsResponse}); ndjson carries no `id:` line, so this + * relies on the `?offset` query param rather than native `Last-Event-ID`. + * * @param stream - AsyncIterable of StreamChunks from chat() - * @param init - Optional Response initialization options (including `abortController`) + * @param init - Optional Response initialization options (including `abortController`, `durability`, `batch`) * @returns Response in HTTP stream format (newline-delimited JSON) * * @example * ```typescript - * const stream = chat({ adapter: openaiText(), model: "gpt-4o", messages: [...] }); - * return toHttpResponse(stream, { abortController }); + * const stream = chat({ adapter: openaiText(), model: "gpt-5.5", messages: [...] }); + * return toHttpResponse(stream, { durability: memoryStream(request) }); * ``` */ export function toHttpResponse( stream: AsyncIterable, - init?: ResponseInit & { abortController?: AbortController }, + init?: ResponseInit & { + abortController?: AbortController + durability?: StreamDurability + batch?: number + }, ): Response { - return new Response(toHttpStream(stream, init?.abortController), { - ...init, + const { abortController, durability, batch, ...responseInit } = init ?? {} + + const body = durability + ? toHttpStream( + durableStreamSource(stream, durability, { abortController, batch }) + .source, + abortController, + ) + : toHttpStream(stream, abortController) + + return new Response(body, { + ...responseInit, }) } diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index a25b00da9..cbff017b5 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -973,16 +973,6 @@ export interface TextOptions< */ parentRunId?: string - /** - * Resume cursor for persistence-backed runs. When provided alongside `runId`, - * and a resume source has been provided by middleware (e.g. `withChatPersistence`), - * the engine replays persisted events after this cursor instead of running the - * adapter fresh. Opaque string; generated/consumed by `@tanstack/ai-persistence`. - * Ignored (a no-op) when no resume source is present — so a normal run is - * unaffected. See {@link StreamChunk.cursor}. - */ - cursor?: string - /** * AG-UI interrupt resume responses supplied by the client on a follow-up run. * Threaded through request parsing now so later runtime behavior can resolve @@ -1650,16 +1640,8 @@ export type AGUIEvent = /** * Chunk returned by the SDK during streaming chat completions. * Uses the AG-UI protocol event format. - * - * `cursor` is an OPTIONAL, opaque per-event resume cursor. It is absent on a - * normal (non-persisted) run and stamped by `withChatPersistence` when persistence - * is active, so a reconnecting client can resume with the last cursor it saw. - * It rides in-band so both the SSE and ndjson transports carry it unchanged. */ -export type StreamChunk = AGUIEvent & { - /** Opaque per-event resume cursor (present only on persisted runs). */ - cursor?: string -} +export type StreamChunk = AGUIEvent // Simple streaming format for basic text completions // Converted to StreamChunk format by convertTextCompletionStream() diff --git a/packages/ai/src/utilities/chat-params.ts b/packages/ai/src/utilities/chat-params.ts index 1002f24c7..2370d0efd 100644 --- a/packages/ai/src/utilities/chat-params.ts +++ b/packages/ai/src/utilities/chat-params.ts @@ -30,23 +30,6 @@ function isValidParts(value: unknown): value is Array<{ type: string }> { return true } -function getRequestCursor(body: unknown): string | undefined { - if (!body || typeof body !== 'object' || !('cursor' in body)) { - return undefined - } - - const cursor = (body as { cursor?: unknown }).cursor - if (cursor === undefined) { - return undefined - } - - if (typeof cursor !== 'string') { - throw new AGUIError('Request body cursor must be a string when provided.') - } - - return cursor -} - /** * Parse and validate an HTTP request body as an AG-UI `RunAgentInput`. * @@ -67,7 +50,6 @@ export function chatParamsFromRequestBody(body: unknown): Promise<{ tools: Array<{ name: string; description: string; parameters: JSONSchema }> forwardedProps: Record state: unknown - cursor?: string resume?: Array /** * @deprecated Use `aguiContext` instead. This alias will be removed in a @@ -89,12 +71,6 @@ export function chatParamsFromRequestBody(body: unknown): Promise<{ } const parsed = parseResult.data - let cursor: string | undefined - try { - cursor = getRequestCursor(body) - } catch (cause) { - return Promise.reject(cause) - } const aguiContext = parsed.context // AG-UI Zod uses `.strip()` so extra fields like `parts` on messages are @@ -127,7 +103,6 @@ export function chatParamsFromRequestBody(body: unknown): Promise<{ }>, forwardedProps: (parsed.forwardedProps ?? {}) as Record, state: parsed.state, - cursor, resume: parsed.resume, context: aguiContext, aguiContext, diff --git a/packages/ai/tests/chat-params.test.ts b/packages/ai/tests/chat-params.test.ts index 7db1fa14a..d9f35166f 100644 --- a/packages/ai/tests/chat-params.test.ts +++ b/packages/ai/tests/chat-params.test.ts @@ -22,7 +22,6 @@ describe('chatParamsFromRequestBody', () => { tools: [], context: [], forwardedProps: { temperature: 0.7 }, - cursor: 'cursor-1', resume: [ { interruptId: 'interrupt-1', @@ -45,10 +44,19 @@ describe('chatParamsFromRequestBody', () => { expect(result.forwardedProps).toEqual({ temperature: 0.7 }) expect(result.aguiContext).toEqual([]) expect(result.context).toBe(result.aguiContext) - expect(result.cursor).toBe('cursor-1') + // Delivery cursor is gone — request parsing never surfaces a `cursor`. + expect('cursor' in result).toBe(false) expect(result.resume).toEqual(validBody.resume) }) + it('ignores any `cursor` field on the body (delivery cursor removed)', async () => { + const result = await chatParamsFromRequestBody({ + ...validBody, + cursor: 'cursor-1', + }) + expect('cursor' in result).toBe(false) + }) + it('preserves the `parts` field on messages (AG-UI strip mode tolerates extras in raw JSON)', async () => { const result = await chatParamsFromRequestBody(validBody) const m = result.messages[0] as { parts?: unknown } diff --git a/packages/ai/tests/generation-params.test.ts b/packages/ai/tests/generation-params.test.ts index d10b50bb3..0e8924153 100644 --- a/packages/ai/tests/generation-params.test.ts +++ b/packages/ai/tests/generation-params.test.ts @@ -11,7 +11,6 @@ describe('generation params helpers', () => { forwardedProps: { tenantId: 'acme' }, threadId: 'thread-1', runId: 'run-1', - cursor: 'cursor-1', }) expect(params).toEqual({ @@ -19,7 +18,6 @@ describe('generation params helpers', () => { forwardedProps: { tenantId: 'acme' }, threadId: 'thread-1', runId: 'run-1', - cursor: 'cursor-1', }) }) @@ -74,20 +72,17 @@ describe('generation params helpers', () => { const raw = generationParamsFromBody('video', { prompt: 'launch', runId: 'raw-run', - cursor: 'raw-cursor', }) const envelope = generationParamsFromBody('video', { data: { prompt: 'launch' }, runId: 'envelope-run', - cursor: 'envelope-cursor', }) expect(raw).toEqual({ - input: { prompt: 'launch', runId: 'raw-run', cursor: 'raw-cursor' }, + input: { prompt: 'launch', runId: 'raw-run' }, forwardedProps: {}, }) expect(envelope.runId).toBe('envelope-run') - expect(envelope.cursor).toBe('envelope-cursor') }) it('throws when envelope forwarded props are not an object', () => { @@ -112,12 +107,6 @@ describe('generation params helpers', () => { runId: 123, }), ).toThrow(/runId/) - expect(() => - generationParamsFromBody('image', { - data: { prompt: 'draw a mountain' }, - cursor: 123, - }), - ).toThrow(/cursor/) }) it('parses request JSON before extracting params', async () => { diff --git a/packages/ai/tests/generation-run-identity-replay.test.ts b/packages/ai/tests/generation-run-identity-replay.test.ts index 46d1cd5ed..b56ab862b 100644 --- a/packages/ai/tests/generation-run-identity-replay.test.ts +++ b/packages/ai/tests/generation-run-identity-replay.test.ts @@ -31,7 +31,6 @@ function chunkSummary(chunk: StreamChunk) { type: chunk.type, runId: 'runId' in chunk ? chunk.runId : undefined, threadId: 'threadId' in chunk ? chunk.threadId : undefined, - cursor: chunk.cursor, } } @@ -98,7 +97,7 @@ function createMockVideoAdapter(overrides?: { } describe('generation run identity and replay', () => { - it('passes identity and advancing cursors through one-shot generation stream events', async () => { + it('passes identity through one-shot generation stream events', async () => { const adapter: ImageAdapter = { kind: 'image', name: 'test-image', @@ -118,7 +117,6 @@ describe('generation run identity and replay', () => { stream: true, threadId: 'thread-1', runId: 'run-1', - cursor: 'cursor-0', }), ) @@ -127,24 +125,21 @@ describe('generation run identity and replay', () => { type: 'RUN_STARTED', runId: 'run-1', threadId: 'thread-1', - cursor: 'cursor-0:1', }, { type: 'CUSTOM', runId: 'run-1', threadId: 'thread-1', - cursor: 'cursor-0:2', }, { type: 'RUN_FINISHED', runId: 'run-1', threadId: 'thread-1', - cursor: 'cursor-0:3', }, ]) }) - it('adds identity and cursor metadata to direct generation devtools events', async () => { + it('adds identity metadata to direct generation devtools events', async () => { const emit = vi.spyOn(aiEventClient, 'emit').mockImplementation(() => {}) const adapter: AudioAdapter = { kind: 'audio', @@ -163,7 +158,6 @@ describe('generation run identity and replay', () => { prompt: 'drums', threadId: 'thread-a', runId: 'run-a', - cursor: 'cursor-a', }) expect(emit).toHaveBeenCalledWith( @@ -171,7 +165,6 @@ describe('generation run identity and replay', () => { expect.objectContaining({ threadId: 'thread-a', runId: 'run-a', - cursor: 'cursor-a:1', }), ) expect(emit).toHaveBeenCalledWith( @@ -179,12 +172,11 @@ describe('generation run identity and replay', () => { expect.objectContaining({ threadId: 'thread-a', runId: 'run-a', - cursor: 'cursor-a:2', }), ) }) - it('passes identity and cursors through speech, transcription, and video streams', async () => { + it('passes identity through speech, transcription, and video streams', async () => { const speech = await collect( generateSpeech({ adapter: { @@ -203,7 +195,6 @@ describe('generation run identity and replay', () => { stream: true, threadId: 'thread-s', runId: 'run-s', - cursor: 'cursor-s', }), ) @@ -224,7 +215,6 @@ describe('generation run identity and replay', () => { stream: true, threadId: 'thread-t', runId: 'run-t', - cursor: 'cursor-t', }), ) @@ -238,32 +228,23 @@ describe('generation run identity and replay', () => { pollingInterval: 0, threadId: 'thread-v', runId: 'run-v', - cursor: 'cursor-v', }), ) - expect(speech.map((chunk) => chunk.cursor)).toEqual([ - 'cursor-s:1', - 'cursor-s:2', - 'cursor-s:3', - ]) expect(speech.map(chunkSummary)).toContainEqual({ type: 'CUSTOM', runId: 'run-s', threadId: 'thread-s', - cursor: 'cursor-s:2', }) expect(transcription.map(chunkSummary)).toContainEqual({ type: 'CUSTOM', runId: 'run-t', threadId: 'thread-t', - cursor: 'cursor-t:2', }) expect(video.map(chunkSummary)).toContainEqual({ type: 'RUN_FINISHED', runId: 'run-v', threadId: 'thread-v', - cursor: 'cursor-v:5', }) }) @@ -288,14 +269,12 @@ describe('generation run identity and replay', () => { }, runId: 'run-replay', threadId: 'thread-replay', - cursor: 'cursor-replay:2', }, { type: EventType.RUN_FINISHED, runId: 'run-replay', threadId: 'thread-replay', finishReason: 'stop', - cursor: 'cursor-replay:3', }, ] satisfies Array @@ -414,7 +393,6 @@ describe('generation run identity and replay', () => { value: { id: 'persisted', model: 'test-model' }, runId: 'run-replay', threadId: 'thread-replay', - cursor: 'cursor-replay:2', }, ] satisfies Array const expectedMessage = @@ -505,7 +483,7 @@ describe('generation run identity and replay', () => { expect(createVideoJob).not.toHaveBeenCalled() }) - it('adds identity and advancing cursors to video status devtools events', async () => { + it('adds identity to video status devtools events', async () => { const emit = vi.spyOn(aiEventClient, 'emit').mockImplementation(() => {}) const adapter = createMockVideoAdapter({ createVideoJob: vi.fn(), @@ -533,7 +511,6 @@ describe('generation run identity and replay', () => { jobId: 'job-status', threadId: 'thread-status', runId: 'run-status', - cursor: 'cursor-status', }) expect(emit).toHaveBeenCalledWith( @@ -541,7 +518,6 @@ describe('generation run identity and replay', () => { expect.objectContaining({ threadId: 'thread-status', runId: 'run-status', - cursor: 'cursor-status:1', }), ) expect(emit).toHaveBeenCalledWith( @@ -549,7 +525,6 @@ describe('generation run identity and replay', () => { expect.objectContaining({ threadId: 'thread-status', runId: 'run-status', - cursor: 'cursor-status:2', }), ) expect(emit).toHaveBeenCalledWith( @@ -557,7 +532,6 @@ describe('generation run identity and replay', () => { expect.objectContaining({ threadId: 'thread-status', runId: 'run-status', - cursor: 'cursor-status:3', }), ) }) diff --git a/packages/ai/tests/resume-seam.test.ts b/packages/ai/tests/resume-seam.test.ts deleted file mode 100644 index f46848886..000000000 --- a/packages/ai/tests/resume-seam.test.ts +++ /dev/null @@ -1,194 +0,0 @@ -import { describe, expect, it } from 'vitest' -import { chat } from '../src/activities/chat/index' -import { defineChatMiddleware } from '../src/activities/chat/middleware/index' -import { ResumeSourceCapability, provideResumeSource } from '../src/resume' -import type { ResumeSource } from '../src/resume' -import type { StreamChunk } from '../src/types' -import { ev, createMockAdapter, collectChunks, getDeltas } from './test-utils' - -/** Middleware that provides a fixed ResumeSource (stands in for withChatPersistence). */ -function withFakeResumeSource(source: ResumeSource) { - return defineChatMiddleware({ - name: 'fake-resume-source', - provides: [ResumeSourceCapability], - setup(ctx) { - provideResumeSource(ctx, source) - }, - }) -} - -describe('chat() resume seam', () => { - it('ignores `cursor` when no resume source is provided (no-op invariant)', async () => { - const { adapter, calls } = createMockAdapter({ - iterations: [ - [ev.runStarted(), ev.textContent('hello'), ev.runFinished('stop')], - ], - }) - - const stream = chat({ - adapter, - messages: [{ role: 'user', content: 'Hi' }], - runId: 'run-1', - cursor: 'some-cursor', - }) - - const chunks = await collectChunks(stream as AsyncIterable) - // Adapter ran normally — cursor was ignored. - expect(calls.length).toBe(1) - expect(getDeltas(chunks).join('')).toBe('hello') - }) - - it('replays persisted events after the cursor and does NOT run the adapter', async () => { - const replayed: Array = [ - { ...ev.textContent(' world'), cursor: '3' }, - { ...ev.runFinished('stop'), cursor: '4' }, - ] - const source: ResumeSource = { - hasRun: async (runId) => runId === 'run-1', - replay: async function* (runId, afterCursor) { - expect(runId).toBe('run-1') - expect(afterCursor).toBe('2') - for (const c of replayed) yield c - }, - getStatus: async () => 'completed', - } - - const { adapter, calls } = createMockAdapter({ - iterations: [ - [ev.runStarted(), ev.textContent('SHOULD NOT RUN'), ev.runFinished()], - ], - }) - - const stream = chat({ - adapter, - messages: [], - runId: 'run-1', - cursor: '2', - middleware: [withFakeResumeSource(source)], - }) - - const chunks = await collectChunks(stream as AsyncIterable) - // The adapter must never run on a resume. - expect(calls.length).toBe(0) - // Output is exactly the replayed tail, cursors intact. - expect(getDeltas(chunks).join('')).toBe(' world') - expect(chunks.map((c) => c.cursor)).toEqual(['3', '4']) - }) - - it('does not replay when cursor is paired with interrupt resume entries', async () => { - let replayCalled = false - const source: ResumeSource = { - hasRun: async (runId) => runId === 'run-1', - replay: async function* () { - replayCalled = true - yield { ...ev.textContent('SHOULD NOT REPLAY'), cursor: '3' } - }, - getStatus: async () => 'interrupted', - } - - const { adapter, calls } = createMockAdapter({ - iterations: [ - [ev.runStarted(), ev.textContent('continued'), ev.runFinished('stop')], - ], - }) - - const stream = chat({ - adapter, - messages: [], - runId: 'run-1', - cursor: '2', - resume: [{ interruptId: 'interrupt-1', status: 'resolved' }], - middleware: [withFakeResumeSource(source)], - }) - - const chunks = await collectChunks(stream as AsyncIterable) - expect(replayCalled).toBe(false) - expect(calls.length).toBe(1) - expect(getDeltas(chunks).join('')).toBe('continued') - }) - - it('replays public events and does NOT re-attach even when the run is still running and the adapter supports it', async () => { - const source: ResumeSource = { - hasRun: async () => true, - replay: async function* () { - yield { ...ev.textContent('replayed'), cursor: '2' } - }, - getStatus: async () => 'running', // still in flight - } - - const { adapter, calls } = createMockAdapter({ - iterations: [ - [ev.runStarted(), ev.textContent(' live'), ev.runFinished('stop')], - ], - }) - // Mark the adapter as re-attach capable (a harness adapter would). - ;(adapter as { supportsReattach?: boolean }).supportsReattach = true - - const stream = chat({ - adapter, - messages: [], - runId: 'run-1', - cursor: '1', - middleware: [withFakeResumeSource(source)], - }) - - const chunks = await collectChunks(stream as AsyncIterable) - // Durable resume is replay-only; live adapter continuation is not part of - // the core chat resume seam. - expect(calls.length).toBe(0) - expect(getDeltas(chunks).join('')).toBe('replayed') - }) - - it('does NOT re-attach for a finished run even if the adapter supports it', async () => { - const source: ResumeSource = { - hasRun: async () => true, - replay: async function* () { - yield { ...ev.textContent('replayed'), cursor: '2' } - }, - getStatus: async () => 'completed', - } - const { adapter, calls } = createMockAdapter({ - iterations: [[ev.runStarted(), ev.textContent('NOPE'), ev.runFinished()]], - }) - ;(adapter as { supportsReattach?: boolean }).supportsReattach = true - - const stream = chat({ - adapter, - messages: [], - runId: 'run-1', - cursor: '1', - middleware: [withFakeResumeSource(source)], - }) - const chunks = await collectChunks(stream as AsyncIterable) - expect(calls.length).toBe(0) - expect(getDeltas(chunks).join('')).toBe('replayed') - }) - - it('falls through to a normal run when the resume source has no such run', async () => { - const source: ResumeSource = { - hasRun: async () => false, - replay: async function* () { - throw new Error('replay should not be called') - }, - getStatus: async () => null, - } - - const { adapter, calls } = createMockAdapter({ - iterations: [ - [ev.runStarted(), ev.textContent('fresh'), ev.runFinished('stop')], - ], - }) - - const stream = chat({ - adapter, - messages: [{ role: 'user', content: 'Hi' }], - runId: 'run-unknown', - cursor: '5', - middleware: [withFakeResumeSource(source)], - }) - - const chunks = await collectChunks(stream as AsyncIterable) - expect(calls.length).toBe(1) - expect(getDeltas(chunks).join('')).toBe('fresh') - }) -}) diff --git a/packages/ai/tests/stream-durability.test.ts b/packages/ai/tests/stream-durability.test.ts new file mode 100644 index 000000000..8bc1043fb --- /dev/null +++ b/packages/ai/tests/stream-durability.test.ts @@ -0,0 +1,270 @@ +import { describe, expect, it } from 'vitest' +import { + decodeOffset, + encodeOffset, + memoryStream, + runIdOf, +} from '../src/stream-durability' +import type { StreamChunk } from '../src/types' + +function chunk(delta: string): StreamChunk { + return { + type: 'TEXT_MESSAGE_CONTENT', + delta, + timestamp: 0, + } as StreamChunk +} + +function terminal(): StreamChunk { + return { type: 'RUN_FINISHED', timestamp: 0 } as StreamChunk +} + +describe('encodeOffset / decodeOffset', () => { + it('round-trips a runId and seq', () => { + const offset = encodeOffset('run-abc', 7) + expect(offset).toBe('run-abc@7') + expect(decodeOffset(offset)).toEqual({ runId: 'run-abc', seq: 7 }) + }) + + it('round-trips runIds that contain @', () => { + const offset = encodeOffset('tenant@run-1', 3) + expect(decodeOffset(offset)).toEqual({ runId: 'tenant@run-1', seq: 3 }) + }) + + // Finding 2: the `now` / `-1` sentinels must not make decode / runId + // extraction throw (`Number('now')` → NaN used to throw). + it('decodes the `now` sentinel without throwing (tail → +Infinity)', () => { + const offset = encodeOffset('run-x', 'now') + expect(offset).toBe('run-x@now') + expect(decodeOffset(offset)).toEqual({ + runId: 'run-x', + seq: Number.POSITIVE_INFINITY, + }) + }) + + it('decodes the `-1` sentinel without throwing (from start)', () => { + expect(decodeOffset('run-x@-1')).toEqual({ runId: 'run-x', seq: -1 }) + }) + + it('extracts the runId from a sentinel offset without parsing the seq', () => { + expect(runIdOf('run-x@now')).toBe('run-x') + expect(runIdOf('run-x@-1')).toBe('run-x') + expect(runIdOf('tenant@run-1@now')).toBe('tenant@run-1') + expect(runIdOf('bare-run')).toBe('bare-run') + }) + + it('resolves a resume runId carried in Last-Event-ID with a `now` seq', () => { + // A tail-join reconnect carries `runId@now`; runId() must not throw. + const durability = memoryStream( + new Request('https://example.test/api/chat', { + method: 'POST', + headers: { 'Last-Event-ID': encodeOffset('run-now', 'now') }, + }), + ) + expect(durability.resumeFrom()).toBe('run-now@now') + expect(durability.runId()).toBe('run-now') + }) + + it('still throws on a genuinely malformed offset', () => { + expect(() => decodeOffset('no-at-sign')).toThrow(/missing @/) + expect(() => decodeOffset('run@not-a-number')).toThrow( + /Invalid durability offset seq/, + ) + }) +}) + +describe('memoryStream', () => { + it('treats a request with no offset as a fresh run and replays from the start', async () => { + const request = new Request('https://example.test/api/chat', { + method: 'POST', + }) + const durability = memoryStream(request) + + expect(durability.resumeFrom()).toBeNull() + const runId = durability.runId() + expect(runId).toBeTruthy() + + await durability.append([chunk('a'), chunk('b'), chunk('c')], 1) + // Mark complete so the live-tailing read returns after draining. + durability.markComplete!() + + const read: Array<{ seq: number; delta: string }> = [] + for await (const { seq, chunk: c } of durability.read('-1')) { + read.push({ seq, delta: (c as { delta: string }).delta }) + } + expect(read).toEqual([ + { seq: 1, delta: 'a' }, + { seq: 2, delta: 'b' }, + { seq: 3, delta: 'c' }, + ]) + }) + + it('append returns the per-chunk backend offsets to tag with', async () => { + const durability = memoryStream( + new Request('https://example.test/api/chat?runId=run-off', { + method: 'POST', + }), + ) + const offsets = await durability.append( + [chunk('a'), chunk('b'), chunk('c')], + 1, + ) + expect(offsets).toEqual([ + encodeOffset('run-off', 1), + encodeOffset('run-off', 2), + encodeOffset('run-off', 3), + ]) + }) + + it('resumes strictly after the offset carried in Last-Event-ID', async () => { + // First, a producer writes three chunks under a known run. + const producer = memoryStream( + new Request('https://example.test/api/chat?runId=run-resume', { + method: 'POST', + }), + ) + expect(producer.runId()).toBe('run-resume') + const offsets = await producer.append( + [chunk('a'), chunk('b'), chunk('c')], + 1, + ) + expect(offsets).toEqual([ + encodeOffset('run-resume', 1), + encodeOffset('run-resume', 2), + encodeOffset('run-resume', 3), + ]) + producer.markComplete!() + + // A reconnect arrives carrying Last-Event-ID at seq 2. + const reconnect = memoryStream( + new Request('https://example.test/api/chat', { + method: 'POST', + headers: { 'Last-Event-ID': encodeOffset('run-resume', 2) }, + }), + ) + expect(reconnect.resumeFrom()).toBe(encodeOffset('run-resume', 2)) + expect(reconnect.runId()).toBe('run-resume') + + const read: Array = [] + for await (const { chunk: c } of reconnect.read(reconnect.resumeFrom()!)) { + read.push((c as { delta: string }).delta) + } + expect(read).toEqual(['c']) + }) + + it('reads the offset from the ?offset query param when no header is present', async () => { + const producer = memoryStream( + new Request('https://example.test/api/chat?runId=run-q', { + method: 'POST', + }), + ) + await producer.append([chunk('x'), chunk('y')], 1) + producer.markComplete!() + + const joiner = memoryStream( + new Request( + `https://example.test/api/chat?offset=${encodeURIComponent( + encodeOffset('run-q', -1), + )}`, + { method: 'POST' }, + ), + ) + expect(joiner.resumeFrom()).toBe(encodeOffset('run-q', -1)) + expect(joiner.runId()).toBe('run-q') + + const read: Array = [] + for await (const { chunk: c } of joiner.read(joiner.resumeFrom()!)) { + read.push((c as { delta: string }).delta) + } + expect(read).toEqual(['x', 'y']) + }) + + // Finding 3: a mid-stream join must live-tail — receive future appends up to + // the terminal, not return at the snapshot tail before the run finishes. + it('live-tails a join on a still-producing run through the terminal', async () => { + const runId = 'run-live' + const producer = memoryStream( + new Request(`https://example.test/api/chat?runId=${runId}`, { + method: 'POST', + }), + ) + await producer.append([chunk('a'), chunk('b')], 1) + + // A joiner attaches mid-stream and reads from the start. + const joiner = memoryStream( + new Request( + `https://example.test/api/chat?offset=${encodeURIComponent( + encodeOffset(runId, -1), + )}`, + { method: 'POST' }, + ), + ) + + const received: Array = [] + const done = (async () => { + for await (const { chunk: c } of joiner.read(joiner.resumeFrom()!)) { + received.push((c as { delta?: string }).delta ?? `[${c.type}]`) + } + })() + + // Give the joiner a tick to drain the initial snapshot and park. + await new Promise((r) => setTimeout(r, 10)) + expect(received).toEqual(['a', 'b']) + + // The producer emits more, then the terminal. + await producer.append([chunk('c'), chunk('d')], 3) + await producer.append([terminal()], 5) + + await done + expect(received).toEqual(['a', 'b', 'c', 'd', '[RUN_FINISHED]']) + }) + + // Finding 2 (hang guard): a join to a runId with no in-process producer + // (crash / cross-instance / bogus id) creates an empty log and parks — no + // `markComplete` ever fires. Threading the consumer's AbortSignal must wake + // the parked waiter so the iteration ends instead of hanging forever. + it('does not hang: an aborted join to a never-produced runId terminates', async () => { + const controller = new AbortController() + const joiner = memoryStream( + new Request('https://example.test/api/chat?runId=never-produced-run', { + method: 'POST', + }), + ) + + const received: Array = [] + const iterated = (async () => { + for await (const { chunk: c } of joiner.read('-1', controller.signal)) { + received.push((c as { delta?: string }).delta ?? `[${c.type}]`) + } + })() + + // Nothing is ever produced under this runId, so the read parks. Give it a + // tick to create its empty log and park on the (never-fired) next append. + await new Promise((r) => setTimeout(r, 10)) + expect(received).toEqual([]) + + // The abort must wake the parked waiter and end the iteration. Race against + // a deadline so a regression (parking forever) fails loudly rather than + // hanging the whole suite. + controller.abort() + const timedOut = Symbol('timed-out') + const outcome = await Promise.race([ + iterated.then(() => 'done' as const), + new Promise((r) => setTimeout(() => r(timedOut), 1000)), + ]) + expect(outcome).toBe('done') + expect(received).toEqual([]) + }) + + it('rejects a caller-supplied ?runId containing CR/LF (SSE injection)', () => { + const durability = memoryStream( + new Request( + `https://example.test/api/chat?runId=${encodeURIComponent( + 'evil\ninjected', + )}`, + { method: 'POST' }, + ), + ) + expect(() => durability.runId()).toThrow(/Invalid runId/) + }) +}) diff --git a/packages/ai/tests/stream-to-response-durability.test.ts b/packages/ai/tests/stream-to-response-durability.test.ts new file mode 100644 index 000000000..221e5be3c --- /dev/null +++ b/packages/ai/tests/stream-to-response-durability.test.ts @@ -0,0 +1,194 @@ +import { describe, expect, it, vi } from 'vitest' +import { encodeOffset, memoryStream } from '../src/stream-durability' +import { toServerSentEventsResponse } from '../src/stream-to-response' +import type { StreamChunk } from '../src/types' + +function textChunk(delta: string): StreamChunk { + return { type: 'TEXT_MESSAGE_CONTENT', delta, timestamp: 0 } as StreamChunk +} + +function fiveChunkStream(): { stream: AsyncIterable; iterated: () => boolean } { + let started = false + const stream: AsyncIterable = { + async *[Symbol.asyncIterator]() { + started = true + for (const d of ['1', '2', '3', '4', '5']) { + yield textChunk(d) + } + }, + } + return { stream, iterated: () => started } +} + +async function readBody(res: Response): Promise { + const reader = res.body!.getReader() + const decoder = new TextDecoder() + let out = '' + for (;;) { + const { done, value } = await reader.read() + if (done) break + out += decoder.decode(value) + } + return out +} + +function parseSseEvents( + body: string, +): Array<{ id?: string; data: string }> { + return body + .split('\n\n') + .filter((block) => block.trim().length > 0) + .map((block) => { + const lines = block.split('\n') + const idLine = lines.find((l) => l.startsWith('id: ')) + const dataLine = lines.find((l) => l.startsWith('data: ')) + return { + ...(idLine ? { id: idLine.slice('id: '.length) } : {}), + data: dataLine ? dataLine.slice('data: '.length) : '', + } + }) +} + +describe('toServerSentEventsResponse with durability', () => { + it('appends + forwards a fresh run, tagging each event with an id offset', async () => { + const request = new Request('https://example.test/api/chat', { + method: 'POST', + }) + const durability = memoryStream(request) + const runId = durability.runId() + const { stream, iterated } = fiveChunkStream() + + const res = toServerSentEventsResponse(stream, { durability }) + const events = parseSseEvents(await readBody(res)) + + expect(iterated()).toBe(true) + expect(events).toHaveLength(5) + expect(events.map((e) => e.id)).toEqual([ + encodeOffset(runId, 1), + encodeOffset(runId, 2), + encodeOffset(runId, 3), + encodeOffset(runId, 4), + encodeOffset(runId, 5), + ]) + + // The durability log now holds all 5 chunks. + const logged: Array = [] + for await (const { chunk } of durability.read('-1')) { + logged.push((chunk as { delta: string }).delta) + } + expect(logged).toEqual(['1', '2', '3', '4', '5']) + }) + + it('replays from the log on resume and never iterates the input stream', async () => { + // Produce a run under a known id. + const producerReq = new Request( + 'https://example.test/api/chat?runId=run-x', + { method: 'POST' }, + ) + const { stream } = fiveChunkStream() + await readBody( + toServerSentEventsResponse(stream, { durability: memoryStream(producerReq) }), + ) + + // Reconnect carrying Last-Event-ID at seq 2. + const reconnectReq = new Request('https://example.test/api/chat', { + method: 'POST', + headers: { 'Last-Event-ID': encodeOffset('run-x', 2) }, + }) + const exploding: AsyncIterable = { + [Symbol.asyncIterator]() { + return { + next() { + throw new Error('input stream must not be iterated on resume') + }, + } + }, + } + + const res = toServerSentEventsResponse(exploding, { + durability: memoryStream(reconnectReq), + }) + const events = parseSseEvents(await readBody(res)) + + expect(events.map((e) => e.id)).toEqual([ + encodeOffset('run-x', 3), + encodeOffset('run-x', 4), + encodeOffset('run-x', 5), + ]) + expect(events.map((e) => JSON.parse(e.data).delta)).toEqual(['3', '4', '5']) + }) + + it('batches appends to at most `batch` chunks', async () => { + const request = new Request('https://example.test/api/chat?runId=run-b', { + method: 'POST', + }) + const durability = memoryStream(request) + const appendSpy = vi.spyOn(durability, 'append') + const { stream } = fiveChunkStream() + + await readBody(toServerSentEventsResponse(stream, { durability, batch: 2 })) + + const batchSizes = appendSpy.mock.calls.map(([chunks]) => chunks.length) + expect(batchSizes.every((n) => n <= 2)).toBe(true) + expect(batchSizes.reduce((a, b) => a + b, 0)).toBe(5) + }) + + // Finding 7a: a non-positive-integer batch is rejected loudly (a NaN used to + // silently disable size-based flushing: `length >= NaN` is always false). + it('rejects a non-positive-integer batch size', () => { + const durability = memoryStream( + new Request('https://example.test/api/chat?runId=run-bad-batch', { + method: 'POST', + }), + ) + const { stream } = fiveChunkStream() + for (const bad of [0, -1, 1.5, NaN]) { + expect(() => + toServerSentEventsResponse(stream, { durability, batch: bad }), + ).toThrow(/Invalid durability batch size/) + } + // A valid positive integer is accepted. + expect(() => + toServerSentEventsResponse(stream, { durability, batch: 4 }), + ).not.toThrow() + }) + + // Finding 4: when the provider stream throws, a terminal RUN_ERROR must be + // persisted to the durability log, so a later reader / join learns the run + // failed instead of finding a log with no terminal. + it('persists a terminal RUN_ERROR to the log when the stream throws', async () => { + const request = new Request('https://example.test/api/chat?runId=run-err', { + method: 'POST', + }) + const durability = memoryStream(request) + + const throwing: AsyncIterable = { + async *[Symbol.asyncIterator]() { + yield textChunk('1') + throw new Error('provider exploded') + }, + } + + // The live consumer still sees a RUN_ERROR (emitted by the transport). + const liveEvents = parseSseEvents( + await readBody(toServerSentEventsResponse(throwing, { durability })), + ) + const liveTypes = liveEvents.map((e) => JSON.parse(e.data).type) + expect(liveTypes).toContain('RUN_ERROR') + + // A second reader / joiner replaying the log sees the persisted terminal. + const joiner = memoryStream( + new Request('https://example.test/api/chat?runId=run-err', { + method: 'POST', + }), + ) + const logged: Array = [] + for await (const { chunk } of joiner.read('-1')) { + logged.push(chunk) + } + expect(logged.map((c) => c.type)).toEqual([ + 'TEXT_MESSAGE_CONTENT', + 'RUN_ERROR', + ]) + }) +}) From 0ce5d8927f73a484456df5e895e64f4f0af1f73f Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Thu, 9 Jul 2026 18:49:37 +0200 Subject: [PATCH 2/9] feat(ai-durable-stream): durable-streams protocol StreamDurability adapter --- packages/ai-durable-stream/package.json | 50 ++ .../ai-durable-stream/src/durable-stream.ts | 328 +++++++++++ packages/ai-durable-stream/src/index.ts | 6 + .../tests/durable-stream.test.ts | 513 ++++++++++++++++++ packages/ai-durable-stream/tsconfig.json | 8 + packages/ai-durable-stream/vite.config.ts | 37 ++ 6 files changed, 942 insertions(+) create mode 100644 packages/ai-durable-stream/package.json create mode 100644 packages/ai-durable-stream/src/durable-stream.ts create mode 100644 packages/ai-durable-stream/src/index.ts create mode 100644 packages/ai-durable-stream/tests/durable-stream.test.ts create mode 100644 packages/ai-durable-stream/tsconfig.json create mode 100644 packages/ai-durable-stream/vite.config.ts diff --git a/packages/ai-durable-stream/package.json b/packages/ai-durable-stream/package.json new file mode 100644 index 000000000..c3337b5a6 --- /dev/null +++ b/packages/ai-durable-stream/package.json @@ -0,0 +1,50 @@ +{ + "name": "@tanstack/ai-durable-stream", + "version": "0.1.0", + "description": "Delivery durability for TanStack AI over the durable-streams HTTP protocol — a resumable StreamDurability transport sink (append/read/resume) that stores zero delivery events itself.", + "author": "", + "license": "MIT", + "repository": { + "type": "git", + "url": "git+https://github.com/TanStack/ai.git", + "directory": "packages/ai-durable-stream" + }, + "keywords": [ + "ai", + "tanstack", + "durable-streams", + "delivery-durability", + "resumable", + "sse" + ], + "type": "module", + "module": "./dist/esm/index.js", + "types": "./dist/esm/index.d.ts", + "exports": { + ".": { + "types": "./dist/esm/index.d.ts", + "import": "./dist/esm/index.js" + } + }, + "files": [ + "dist", + "src" + ], + "scripts": { + "build": "vite build", + "clean": "premove ./build ./dist", + "lint:fix": "eslint ./src --fix", + "test:build": "publint --strict", + "test:eslint": "eslint ./src", + "test:lib": "vitest", + "test:lib:dev": "pnpm test:lib --watch", + "test:types": "tsc" + }, + "peerDependencies": { + "@tanstack/ai": "workspace:^" + }, + "devDependencies": { + "@tanstack/ai": "workspace:*", + "@vitest/coverage-v8": "4.0.14" + } +} diff --git a/packages/ai-durable-stream/src/durable-stream.ts b/packages/ai-durable-stream/src/durable-stream.ts new file mode 100644 index 000000000..ed75c5811 --- /dev/null +++ b/packages/ai-durable-stream/src/durable-stream.ts @@ -0,0 +1,328 @@ +import { decodeOffset, encodeOffset } from '@tanstack/ai' +import type { StreamChunk, StreamDurability } from '@tanstack/ai' + +/** + * Options for {@link durableStream}. + */ +export interface DurableStreamOptions { + /** Base URL of the durable-streams server (no trailing slash needed). */ + server: string + /** + * Prefix for the stream name (`{prefix}/{runId}`). Defaults to `runs`, so a + * run lands at `runs/{runId}`. + */ + streamPrefix?: string + /** + * `fetch` implementation to use. Defaults to the global `fetch`. Injectable + * for testing and for wiring a custom transport / auth layer. + */ + fetch?: typeof globalThis.fetch +} + +/** One parsed SSE event from a durable-streams read. */ +interface SseEvent { + id?: string + event?: string + data?: string +} + +/** Read a byte `ReadableStream` as `\n`-delimited lines (blank lines included). */ +async function* readLines( + body: ReadableStream, +): AsyncGenerator { + const reader = body.getReader() + const decoder = new TextDecoder() + let buffer = '' + try { + for (;;) { + const { done, value } = await reader.read() + if (done) break + buffer += decoder.decode(value, { stream: true }) + const parts = buffer.split('\n') + buffer = parts.pop() ?? '' + for (const raw of parts) { + yield raw.endsWith('\r') ? raw.slice(0, -1) : raw + } + } + if (buffer) { + yield buffer.endsWith('\r') ? buffer.slice(0, -1) : buffer + } + } finally { + reader.releaseLock() + } +} + +/** + * Parse an SSE `ReadableStream` into discrete events. Events are separated by a + * blank line; `id:` / `event:` / `data:` fields are collected per event. `:` + * comment lines (proxy/CDN keepalives) are skipped. A final event not + * terminated by a trailing blank line is still flushed. + */ +async function* parseSseEvents( + body: ReadableStream, +): AsyncGenerator { + let current: SseEvent = {} + let hasField = false + + for await (const line of readLines(body)) { + if (line === '') { + if (hasField) { + yield current + current = {} + hasField = false + } + continue + } + if (line.startsWith(':')) continue + const colon = line.indexOf(':') + const field = colon === -1 ? line : line.slice(0, colon) + let value = colon === -1 ? '' : line.slice(colon + 1) + if (value.startsWith(' ')) value = value.slice(1) + if (field === 'id') current.id = value + else if (field === 'event') current.event = value + else if (field === 'data') { + current.data = + current.data === undefined ? value : `${current.data}\n${value}` + } + hasField = true + } + if (hasField) yield current +} + +/** Extract the seq token (`-1` / `now` / a number) from a `runId@seq` offset. */ +function offsetToken(offset: string): string { + return offset.includes('@') + ? offset.slice(offset.lastIndexOf('@') + 1) + : offset +} + +/** + * Interpret a live-read `event: control` frame's `data`. A rollover frame + * carries `{ offset }` (resume the read from there across a live-window + * boundary); a close frame carries `{ closed: true }` (the server finalized the + * stream). Anything else — absent data, invalid JSON, or an object with neither + * signal — yields `{}`, which the caller treats as malformed. + */ +function parseControlFrame(data: string | undefined): { + offset?: string + closed?: boolean +} { + if (data === undefined) return {} + let parsed: unknown + try { + parsed = JSON.parse(data) + } catch { + return {} + } + if (typeof parsed !== 'object' || parsed === null) return {} + const result: { offset?: string; closed?: boolean } = {} + const offset = (parsed as { offset?: unknown }).offset + if (typeof offset === 'string') result.offset = offset + const closed = (parsed as { closed?: unknown }).closed + if (closed === true) result.closed = true + return result +} + +function safeSearchParam(request: Request, key: string): string | null { + try { + return new URL(request.url).searchParams.get(key) + } catch { + return null + } +} + +/** + * A {@link StreamDurability} backed by the + * [durable-streams](https://durablestreams.com) HTTP protocol. We own **zero** + * delivery-event storage — the durable-streams server owns the bytes (its own + * WAL / group-commit), and this adapter is a thin append/read/resume shim. + * + * - `runId()` — fresh: minted; resume: parsed from `Last-Event-ID` / `?offset`. + * - stream name — `${streamPrefix ?? 'runs'}/${runId}`. + * - `append` — `PUT` the stream once (idempotent create), then `POST` each chunk + * INDIVIDUALLY (sequential awaited POSTs preserve order); every chunk is tagged + * with its OWN backend offset (from that POST's `Stream-Next-Offset`), so the + * returned offset array is fully populated (one resume offset per chunk) and + * exactly-once holds at any batch size. durable-streams is built for + * high-throughput per-append writes with server-side group-commit, so + * one-POST-per-chunk is the intended usage. + * - `read` — `GET ?offset=&live=sse`, parsing `data` events into `{ seq, chunk }` + * and reconnecting across the CDN live-window boundary via the last offset. The + * `id:` the server replays for chunk i equals the offset `append` tagged chunk + * i with on the fresh pass — the invariant that makes resume exactly-once. + * + * @example + * ```ts + * import { toServerSentEventsResponse } from '@tanstack/ai' + * import { durableStream } from '@tanstack/ai-durable-stream' + * + * export async function POST(request: Request) { + * const stream = chat({ adapter, model: 'gpt-5.5', messages }) + * return toServerSentEventsResponse(stream, { + * durability: durableStream(request, { server: process.env.DS_URL! }), + * }) + * } + * ``` + */ +export function durableStream( + request: Request, + options: DurableStreamOptions, +): StreamDurability { + const fetchFn = options.fetch ?? globalThis.fetch + const prefix = options.streamPrefix ?? 'runs' + // Validate the server as an absolute URL at construction, so a malformed + // base fails loudly here rather than surfacing later as an opaque fetch + // error deep inside `append` / `read`. + try { + void new URL(options.server) + } catch { + throw new Error( + `durableStream: invalid server URL: ${JSON.stringify(options.server)}`, + ) + } + const server = options.server.replace(/\/$/, '') + + const resumeOffset = + request.headers.get('Last-Event-ID') ?? safeSearchParam(request, 'offset') + + let cachedRunId: string | undefined + const runId = (): string => { + if (cachedRunId === undefined) { + if (resumeOffset && resumeOffset.includes('@')) { + cachedRunId = decodeOffset(resumeOffset).runId + } else { + cachedRunId = safeSearchParam(request, 'runId') ?? crypto.randomUUID() + } + } + return cachedRunId + } + + const streamUrl = (): string => + `${server}/streams/${encodeURIComponent(`${prefix}/${runId()}`)}` + + let created = false + const ensureCreated = async (): Promise => { + if (created) return + // Idempotent create. A repeat PUT (e.g. a second server instance for the + // same run) is a server-side no-op. + const res = await fetchFn(streamUrl(), { method: 'PUT' }) + if (!res.ok) { + throw new Error( + `durableStream: failed to create stream (${res.status} ${res.statusText})`, + ) + } + created = true + } + + return { + resumeFrom: () => resumeOffset, + runId, + append: async (chunks, startSeq) => { + await ensureCreated() + // POST each chunk INDIVIDUALLY (sequential awaited POSTs preserve order) + // and capture EACH message's own backend offset from its response's + // `Stream-Next-Offset`. This yields a fully-populated per-chunk offset + // array — every chunk carries a backend-addressable resume id — so a + // mid-batch reconnect resumes at the exact chunk it dropped on and + // exactly-once holds at ANY batch size (the transport `batch` option only + // controls how many chunks are handed to one `append` call; we POST them + // one-by-one internally). The tagged offset MUST equal the `id:` the + // server will emit for that same chunk on a GET replay, so the fresh path + // and the resume path share one key space. + const id = runId() + const offsets: Array = [] + for (let i = 0; i < chunks.length; i++) { + const res = await fetchFn(streamUrl(), { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify([chunks[i]]), + }) + if (!res.ok) { + throw new Error( + `durableStream: failed to append (${res.status} ${res.statusText})`, + ) + } + const next = res.headers.get('Stream-Next-Offset') + const seq = + next !== null && Number.isFinite(Number(next)) + ? Number(next) + : startSeq + i + offsets.push(encodeOffset(id, seq)) + } + return offsets + }, + read: async function* (offset, signal) { + let cursor = offsetToken(offset) + // Reconnect loop: the CDN live window closes the SSE response + // periodically (~60s); resume from the last delivered offset. + reconnect: for (;;) { + if (signal?.aborted) return + const url = `${streamUrl()}?offset=${encodeURIComponent( + cursor, + )}&live=sse` + const res = await fetchFn(url, { method: 'GET', signal }) + if (!res.ok) { + throw new Error( + `durableStream: failed to read (${res.status} ${res.statusText})`, + ) + } + if (!res.body) { + throw new Error('durableStream: read response had no body') + } + // A read ends legitimately only when it saw the run's terminal event + // or the server flagged the stream closed (`Stream-Closed` header / + // a `{ closed: true }` control frame). A bare EOF without either means + // the live window was cut mid-run — surface it rather than silently + // truncating a still-in-flight read. + const streamClosed = res.headers.get('Stream-Closed') === 'true' + let genuineEnd = streamClosed + for await (const evt of parseSseEvents(res.body)) { + if (evt.event === 'control') { + const parsed = parseControlFrame(evt.data) + if (parsed.offset !== undefined) { + cursor = parsed.offset + continue reconnect + } + if (parsed.closed) { + genuineEnd = true + continue + } + // A control frame we cannot act on (unparseable, or neither a + // rollover offset nor a close signal) — surface rather than + // silently ending the read. + throw new Error( + 'durableStream: malformed control frame (no offset or close signal) in read stream', + ) + } + if (evt.data === undefined) continue + const chunk = JSON.parse(evt.data) as StreamChunk + // Every data event MUST carry a numeric offset id; a missing / + // non-numeric id would yield a `NaN` seq that corrupts the resume + // cursor, so reject the frame instead. + if (evt.id === undefined) { + throw new Error( + 'durableStream: read event missing id (cannot derive resume offset)', + ) + } + const seq = Number(evt.id) + if (!Number.isFinite(seq)) { + throw new Error( + `durableStream: read event has non-numeric id: ${JSON.stringify( + evt.id, + )}`, + ) + } + cursor = evt.id + if (chunk.type === 'RUN_FINISHED' || chunk.type === 'RUN_ERROR') { + genuineEnd = true + } + yield { seq, chunk } + } + if (genuineEnd) return + throw new Error( + 'durableStream: read stream ended without a terminal event or close signal — connection was likely cut short.', + ) + } + }, + } +} diff --git a/packages/ai-durable-stream/src/index.ts b/packages/ai-durable-stream/src/index.ts new file mode 100644 index 000000000..74bb8d4ad --- /dev/null +++ b/packages/ai-durable-stream/src/index.ts @@ -0,0 +1,6 @@ +/** + * `@tanstack/ai-durable-stream` — a {@link StreamDurability} backend over the + * durable-streams HTTP protocol, for production-grade delivery durability. + */ +export { durableStream } from './durable-stream' +export type { DurableStreamOptions } from './durable-stream' diff --git a/packages/ai-durable-stream/tests/durable-stream.test.ts b/packages/ai-durable-stream/tests/durable-stream.test.ts new file mode 100644 index 000000000..f1280f415 --- /dev/null +++ b/packages/ai-durable-stream/tests/durable-stream.test.ts @@ -0,0 +1,513 @@ +import { describe, expect, it, vi } from 'vitest' +import { encodeOffset, toServerSentEventsResponse } from '@tanstack/ai' +import { durableStream } from '../src' +import type { StreamChunk } from '@tanstack/ai' + +function textChunk(delta: string): StreamChunk { + return { type: 'TEXT_MESSAGE_CONTENT', delta, timestamp: 0 } as StreamChunk +} + +async function readBody(res: Response): Promise { + const reader = res.body!.getReader() + const decoder = new TextDecoder() + let out = '' + for (;;) { + const { done, value } = await reader.read() + if (done) break + out += decoder.decode(value) + } + return out +} + +function parseSseEvents(body: string): Array<{ id?: string; data: string }> { + return body + .split('\n\n') + .filter((block) => block.trim().length > 0) + .map((block) => { + const lines = block.split('\n') + const idLine = lines.find((l) => l.startsWith('id: ')) + const dataLine = lines.find((l) => l.startsWith('data: ')) + return { + ...(idLine ? { id: idLine.slice('id: '.length) } : {}), + data: dataLine ? dataLine.slice('data: '.length) : '', + } + }) +} + +/** + * A fake DS server that assigns NON-1-based, non-contiguous offsets (1000, + * 1007, 1014, …) so a test can prove the client-facing ids are the BACKEND's + * offsets rather than a transport-local 1-based counter. + */ +function makeNon1BasedDsServer() { + const entries: Array<{ offset: number; data: string }> = [] + let nextOffset = 1000 + const fetchStub = vi.fn( + async (input: string | URL | Request, init?: RequestInit) => { + const url = new URL(typeof input === 'string' ? input : input.toString()) + const method = (init?.method ?? 'GET').toUpperCase() + if (method === 'PUT') return new Response(null, { status: 201 }) + if (method === 'POST') { + const body = JSON.parse(String(init?.body)) as Array + let last = nextOffset + for (const c of body) { + last = nextOffset + entries.push({ offset: nextOffset, data: JSON.stringify(c) }) + nextOffset += 7 + } + return new Response(null, { + status: 200, + headers: { 'Stream-Next-Offset': String(last) }, + }) + } + // GET (read strictly after ?offset) + const offsetParam = url.searchParams.get('offset') ?? '-1' + const from = + offsetParam === '-1' || offsetParam === '' + ? Number.NEGATIVE_INFINITY + : Number(offsetParam) + let sse = '' + for (const e of entries) { + if (e.offset > from) sse += `id: ${e.offset}\ndata: ${e.data}\n\n` + } + return new Response(sse, { + status: 200, + headers: { + 'Content-Type': 'text/event-stream', + 'Stream-Closed': 'true', + }, + }) + }, + ) + return { fetchStub, entries } +} + +/** + * An in-memory fake of the durable-streams HTTP protocol: + * - PUT /streams/{name} create (idempotent) + * - POST /streams/{name} append JSON array; sets `Stream-Next-Offset` + * - GET /streams/{name}?offset= read SSE from strictly after `offset` + */ +function makeFakeDsServer() { + const streams = new Map>() + const puts: Array<{ name: string; created: boolean }> = [] + + const fetchStub = vi.fn( + async (input: string | URL | Request, init?: RequestInit) => { + const url = new URL(typeof input === 'string' ? input : input.toString()) + const method = (init?.method ?? 'GET').toUpperCase() + const name = decodeURIComponent( + url.pathname.replace(/^\/streams\//, ''), + ) + + if (method === 'PUT') { + const created = !streams.has(name) + if (created) streams.set(name, []) + puts.push({ name, created }) + return new Response(null, { status: created ? 201 : 200 }) + } + + if (method === 'POST') { + const body = JSON.parse(String(init?.body)) as Array + const arr = streams.get(name) ?? [] + if (!streams.has(name)) streams.set(name, arr) + for (const c of body) arr.push(JSON.stringify(c)) + return new Response(null, { + status: 200, + headers: { 'Stream-Next-Offset': String(arr.length) }, + }) + } + + // GET (read) + const arr = streams.get(name) ?? [] + const offsetParam = url.searchParams.get('offset') ?? '-1' + const from = + offsetParam === '-1' || offsetParam === '' + ? 0 + : offsetParam === 'now' + ? arr.length + : Number(offsetParam) + let sse = '' + for (let i = from; i < arr.length; i++) { + const pos = i + 1 + sse += `id: ${pos}\ndata: ${arr[i]}\n\n` + } + // This fake returns the finalized stored range then EOF, so flag the + // stream closed — a genuine end (not a cut-short live window). + return new Response(sse, { + status: 200, + headers: { + 'Content-Type': 'text/event-stream', + 'Stream-Closed': 'true', + }, + }) + }, + ) + + return { fetchStub, streams, puts } +} + +describe('durableStream', () => { + it('parses runId and stream identity from the request', () => { + const { fetchStub } = makeFakeDsServer() + const d = durableStream( + new Request('https://app.test/api/chat?runId=run-1', { method: 'POST' }), + { server: 'https://ds.test', fetch: fetchStub }, + ) + expect(d.resumeFrom()).toBeNull() + expect(d.runId()).toBe('run-1') + }) + + it('creates the stream once (PUT idempotent) and advances the offset on append', async () => { + const { fetchStub, puts } = makeFakeDsServer() + const d = durableStream( + new Request('https://app.test/api/chat?runId=run-2', { method: 'POST' }), + { server: 'https://ds.test', fetch: fetchStub }, + ) + + // append POSTs each chunk INDIVIDUALLY and reads back its own backend + // offset, so EVERY chunk carries a resume offset (fully populated array, + // no `undefined`) — this is what makes a mid-batch reconnect exactly-once. + const first = await d.append([textChunk('a'), textChunk('b'), textChunk('c')], 1) + expect(first).toEqual([ + encodeOffset('run-2', 1), + encodeOffset('run-2', 2), + encodeOffset('run-2', 3), + ]) + + const second = await d.append([textChunk('d'), textChunk('e')], 4) + expect(second).toEqual([ + encodeOffset('run-2', 4), + encodeOffset('run-2', 5), + ]) + + // The stream is created exactly once, on the first append (PUT-once); + // per-chunk POSTs never re-PUT. The server-side PUT stays idempotent. + expect(puts).toHaveLength(1) + expect(puts[0]).toEqual({ name: 'runs/run-2', created: true }) + }) + + it('reads the full stream from the start', async () => { + const { fetchStub } = makeFakeDsServer() + const d = durableStream( + new Request('https://app.test/api/chat?runId=run-3', { method: 'POST' }), + { server: 'https://ds.test', fetch: fetchStub }, + ) + await d.append([textChunk('a'), textChunk('b'), textChunk('c')], 1) + + const read: Array<{ seq: number; delta: string }> = [] + for await (const { seq, chunk } of d.read('-1')) { + read.push({ seq, delta: (chunk as { delta: string }).delta }) + } + expect(read).toEqual([ + { seq: 1, delta: 'a' }, + { seq: 2, delta: 'b' }, + { seq: 3, delta: 'c' }, + ]) + }) + + it('reads exactly the tail strictly after a resume offset', async () => { + const { fetchStub } = makeFakeDsServer() + // Produce under a known run. + const producer = durableStream( + new Request('https://app.test/api/chat?runId=run-4', { method: 'POST' }), + { server: 'https://ds.test', fetch: fetchStub }, + ) + await producer.append( + [textChunk('a'), textChunk('b'), textChunk('c'), textChunk('d')], + 1, + ) + + // Reconnect carrying Last-Event-ID at seq 2. + const reconnect = durableStream( + new Request('https://app.test/api/chat', { + method: 'POST', + headers: { 'Last-Event-ID': encodeOffset('run-4', 2) }, + }), + { server: 'https://ds.test', fetch: fetchStub }, + ) + expect(reconnect.resumeFrom()).toBe(encodeOffset('run-4', 2)) + expect(reconnect.runId()).toBe('run-4') + + const read: Array<{ seq: number; delta: string }> = [] + for await (const { seq, chunk } of reconnect.read(reconnect.resumeFrom()!)) { + read.push({ seq, delta: (chunk as { delta: string }).delta }) + } + expect(read).toEqual([ + { seq: 3, delta: 'c' }, + { seq: 4, delta: 'd' }, + ]) + }) + + // Finding 7b: a malformed server URL is rejected at construction. + it('rejects a malformed server URL at construction', () => { + expect(() => + durableStream( + new Request('https://app.test/api/chat', { method: 'POST' }), + { server: 'not a url' }, + ), + ).toThrow(/invalid server URL/) + }) + + // Finding 7c: a read data event without a usable numeric id must throw rather + // than yield a `NaN` seq that would corrupt the resume cursor. + it('throws on a read data event missing its id', async () => { + const fetchStub = vi.fn(async (_input: unknown, init?: RequestInit) => { + const method = (init?.method ?? 'GET').toUpperCase() + if (method === 'PUT') return new Response(null, { status: 201 }) + // A data event with no `id:` line. + return new Response(`data: ${JSON.stringify(textChunk('x'))}\n\n`, { + status: 200, + headers: { 'Content-Type': 'text/event-stream' }, + }) + }) + const d = durableStream( + new Request('https://app.test/api/chat?runId=run-noid', { + method: 'POST', + }), + { server: 'https://ds.test', fetch: fetchStub as unknown as typeof fetch }, + ) + await expect(async () => { + for await (const _ of d.read('-1')) { + // drain + } + }).rejects.toThrow(/missing id/) + }) + + it('throws on a read data event with a non-numeric id', async () => { + const fetchStub = vi.fn(async (_input: unknown, init?: RequestInit) => { + const method = (init?.method ?? 'GET').toUpperCase() + if (method === 'PUT') return new Response(null, { status: 201 }) + return new Response( + `id: abc\ndata: ${JSON.stringify(textChunk('x'))}\n\n`, + { + status: 200, + headers: { 'Content-Type': 'text/event-stream' }, + }, + ) + }) + const d = durableStream( + new Request('https://app.test/api/chat?runId=run-badid', { + method: 'POST', + }), + { server: 'https://ds.test', fetch: fetchStub as unknown as typeof fetch }, + ) + await expect(async () => { + for await (const _ of d.read('-1')) { + // drain + } + }).rejects.toThrow(/non-numeric id/) + }) + + // Finding 5: a read that ends without a terminal / close signal (a cut-short + // live window) must surface rather than silently truncating. + it('surfaces a bare-EOF end of a still-in-flight read', async () => { + const fetchStub = vi.fn(async (_input: unknown, init?: RequestInit) => { + const method = (init?.method ?? 'GET').toUpperCase() + if (method === 'PUT') return new Response(null, { status: 201 }) + // Two data events, then EOF — no terminal chunk, no Stream-Closed header, + // no rollover control frame. + return new Response( + `id: 1\ndata: ${JSON.stringify(textChunk('a'))}\n\n` + + `id: 2\ndata: ${JSON.stringify(textChunk('b'))}\n\n`, + { status: 200, headers: { 'Content-Type': 'text/event-stream' } }, + ) + }) + const d = durableStream( + new Request('https://app.test/api/chat?runId=run-trunc', { + method: 'POST', + }), + { server: 'https://ds.test', fetch: fetchStub as unknown as typeof fetch }, + ) + const seen: Array = [] + await expect(async () => { + for await (const { chunk } of d.read('-1')) { + seen.push((chunk as { delta: string }).delta) + } + }).rejects.toThrow(/ended without a terminal event or close signal/) + // It still delivered what it saw before surfacing the truncation. + expect(seen).toEqual(['a', 'b']) + }) + + it('surfaces a malformed control frame instead of silently ending', async () => { + const fetchStub = vi.fn(async (_input: unknown, init?: RequestInit) => { + const method = (init?.method ?? 'GET').toUpperCase() + if (method === 'PUT') return new Response(null, { status: 201 }) + // A control frame with unparseable data (neither a rollover offset nor a + // close signal). + return new Response(`event: control\ndata: not-json\n\n`, { + status: 200, + headers: { 'Content-Type': 'text/event-stream' }, + }) + }) + const d = durableStream( + new Request('https://app.test/api/chat?runId=run-ctrl', { + method: 'POST', + }), + { server: 'https://ds.test', fetch: fetchStub as unknown as typeof fetch }, + ) + await expect(async () => { + for await (const _ of d.read('-1')) { + // drain + } + }).rejects.toThrow(/malformed control frame/) + }) +}) + +// Finding 1: the client-facing SSE ids on a FRESH run must be the offsets the +// backend actually assigns (here, non-1-based 1000/1007/1014…), and a reconnect +// carrying one of those backend offsets must resume correctly. +describe('durableStream + toServerSentEventsResponse (backend-offset ids)', () => { + function textStream(deltas: Array): AsyncIterable { + return { + async *[Symbol.asyncIterator]() { + for (const d of deltas) yield textChunk(d) + }, + } + } + + async function sseIds(res: Response): Promise> { + return parseSseEvents(await readBody(res)).map((e) => e.id) + } + + it('tags fresh-run events with the backend offsets (not a 1-based counter)', async () => { + const { fetchStub } = makeNon1BasedDsServer() + const durability = durableStream( + new Request('https://app.test/api/chat?runId=run-x', { method: 'POST' }), + { server: 'https://ds.test', fetch: fetchStub }, + ) + // batch:1 so every chunk flushes as its own append and carries a backend id. + const res = toServerSentEventsResponse(textStream(['1', '2', '3']), { + durability, + batch: 1, + }) + expect(await sseIds(res)).toEqual([ + encodeOffset('run-x', 1000), + encodeOffset('run-x', 1007), + encodeOffset('run-x', 1014), + ]) + }) + + it('resumes from a backend offset carried in Last-Event-ID', async () => { + const { fetchStub } = makeNon1BasedDsServer() + // Produce the full run first. + await readBody( + toServerSentEventsResponse(textStream(['1', '2', '3']), { + durability: durableStream( + new Request('https://app.test/api/chat?runId=run-x', { + method: 'POST', + }), + { server: 'https://ds.test', fetch: fetchStub }, + ), + batch: 1, + }), + ) + + // Reconnect at the backend offset 1007; the input stream must NOT be + // iterated on resume. + const exploding: AsyncIterable = { + [Symbol.asyncIterator]() { + return { + next() { + throw new Error('input stream must not be iterated on resume') + }, + } + }, + } + const reconnect = durableStream( + new Request('https://app.test/api/chat', { + method: 'POST', + headers: { 'Last-Event-ID': encodeOffset('run-x', 1007) }, + }), + { server: 'https://ds.test', fetch: fetchStub }, + ) + const res = toServerSentEventsResponse(exploding, { durability: reconnect }) + const events = parseSseEvents(await readBody(res)) + expect(events.map((e) => e.id)).toEqual([encodeOffset('run-x', 1014)]) + expect(events.map((e) => JSON.parse(e.data).delta)).toEqual(['3']) + }) +}) + +// Finding 1 (exactly-once): with the DEFAULT batch (>1), an entire short run is +// one `append` call. Per-chunk POSTing tags EVERY chunk with its own backend +// offset, so a mid-batch drop (including a drop within the very first batch) +// resumes at the exact chunk it dropped on — each chunk is delivered exactly +// once, no dup and no skip. The old per-batch tagging left every non-last chunk +// id-less, so this fails against it (and it must NOT be weakened to batch:1). +describe('durableStream — exactly-once across a mid-batch reconnect', () => { + function textStream(deltas: Array): AsyncIterable { + return { + async *[Symbol.asyncIterator]() { + for (const d of deltas) yield textChunk(d) + }, + } + } + + it('delivers each chunk exactly once when the socket drops mid-first-batch', async () => { + const { fetchStub } = makeNon1BasedDsServer() + const full = ['a', 'b', 'c', 'd', 'e', 'f'] + + // Produce the whole run with the DEFAULT batch (no `batch` option). All 6 + // chunks (< 32) buffer into a single `append`; the fix POSTs them one-by-one + // so each gets its own backend offset (1000, 1007, 1014, …). + const produced = parseSseEvents( + await readBody( + toServerSentEventsResponse(textStream(full), { + durability: durableStream( + new Request('https://app.test/api/chat?runId=run-eo', { + method: 'POST', + }), + { server: 'https://ds.test', fetch: fetchStub }, + ), + }), + ), + ) + + // Core of the fix: EVERY chunk in the batch carries a resume id. Under the + // old per-batch tagging only the last chunk did (the rest were id-less), so + // this assertion fails against the pre-fix code. + expect(produced).toHaveLength(6) + expect(produced.every((e) => e.id !== undefined)).toBe(true) + expect(produced.map((e) => JSON.parse(e.data).delta)).toEqual(full) + + // Simulate a socket drop MID-BATCH after receiving only the first 2 chunks + // (well within the single 6-chunk first batch). The client's last-seen id is + // the 2nd chunk's backend offset — defined only because of per-chunk tagging. + const beforeDrop = produced.slice(0, 2) + const lastSeenId = beforeDrop.at(-1)?.id + expect(lastSeenId).toBeDefined() + + // Reconnect carrying that id as Last-Event-ID; the input stream must NOT be + // re-iterated on resume (replay reads only from the backend). + const exploding: AsyncIterable = { + [Symbol.asyncIterator]() { + return { + next() { + throw new Error('input stream must not be iterated on resume') + }, + } + }, + } + const afterDrop = parseSseEvents( + await readBody( + toServerSentEventsResponse(exploding, { + durability: durableStream( + new Request('https://app.test/api/chat', { + method: 'POST', + headers: { 'Last-Event-ID': lastSeenId! }, + }), + { server: 'https://ds.test', fetch: fetchStub }, + ), + }), + ), + ) + + // What the joiner actually saw = pre-drop prefix + post-reconnect tail. The + // backend replays strictly AFTER the offset, so there is no overlap: each + // chunk appears exactly once, in order, with none skipped. + const received = [...beforeDrop, ...afterDrop].map( + (e) => JSON.parse(e.data).delta as string, + ) + expect(received).toEqual(full) + }) +}) diff --git a/packages/ai-durable-stream/tsconfig.json b/packages/ai-durable-stream/tsconfig.json new file mode 100644 index 000000000..c38689f4e --- /dev/null +++ b/packages/ai-durable-stream/tsconfig.json @@ -0,0 +1,8 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "outDir": "dist" + }, + "include": ["src", "tests"], + "exclude": ["node_modules", "dist"] +} diff --git a/packages/ai-durable-stream/vite.config.ts b/packages/ai-durable-stream/vite.config.ts new file mode 100644 index 000000000..11f5b20b7 --- /dev/null +++ b/packages/ai-durable-stream/vite.config.ts @@ -0,0 +1,37 @@ +import { defineConfig, mergeConfig } from 'vitest/config' +import { tanstackViteConfig } from '@tanstack/vite-config' +import packageJson from './package.json' + +const config = defineConfig({ + test: { + name: packageJson.name, + dir: './', + watch: false, + + globals: true, + environment: 'node', + include: ['tests/**/*.test.ts'], + coverage: { + provider: 'v8', + reporter: ['text', 'json', 'html', 'lcov'], + exclude: [ + 'node_modules/', + 'dist/', + 'tests/', + '**/*.test.ts', + '**/*.config.ts', + '**/types.ts', + ], + include: ['src/**/*.ts'], + }, + }, +}) + +export default mergeConfig( + config, + tanstackViteConfig({ + entry: ['./src/index.ts'], + srcDir: './src', + cjs: false, + }), +) From 6d3fc4c409e818f46fee40c2e517950d06951bea Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Thu, 9 Jul 2026 18:50:20 +0200 Subject: [PATCH 3/9] feat(persistence)!: state-only middleware + memory backend + shared conformance suite; remove delivery stores --- packages/ai-persistence/package.json | 15 +- .../ai-persistence/src/approval-controller.ts | 71 ---- packages/ai-persistence/src/capabilities.ts | 38 +- packages/ai-persistence/src/cursor.ts | 84 ---- packages/ai-persistence/src/history.ts | 25 -- packages/ai-persistence/src/index.ts | 39 +- packages/ai-persistence/src/memory.ts | 163 +------- packages/ai-persistence/src/middleware.ts | 199 ++-------- packages/ai-persistence/src/resume-source.ts | 78 ---- .../ai-persistence/src/testkit/conformance.ts | 352 +++++++++++++++++ packages/ai-persistence/src/types.ts | 107 +----- packages/ai-persistence/tests/cursor.test.ts | 65 ---- .../tests/generation-artifacts.test.ts | 1 - packages/ai-persistence/tests/history.test.ts | 63 --- .../ai-persistence/tests/interrupts.test.ts | 170 +-------- .../tests/memory.conformance.test.ts | 4 + packages/ai-persistence/tests/memory.test.ts | 128 +------ .../ai-persistence/tests/state-only.test.ts | 105 +++++ .../tests/with-persistence.test.ts | 360 +----------------- packages/ai-persistence/vite.config.ts | 6 +- 20 files changed, 576 insertions(+), 1497 deletions(-) delete mode 100644 packages/ai-persistence/src/approval-controller.ts delete mode 100644 packages/ai-persistence/src/cursor.ts delete mode 100644 packages/ai-persistence/src/history.ts delete mode 100644 packages/ai-persistence/src/resume-source.ts create mode 100644 packages/ai-persistence/src/testkit/conformance.ts delete mode 100644 packages/ai-persistence/tests/cursor.test.ts delete mode 100644 packages/ai-persistence/tests/history.test.ts create mode 100644 packages/ai-persistence/tests/memory.conformance.test.ts create mode 100644 packages/ai-persistence/tests/state-only.test.ts diff --git a/packages/ai-persistence/package.json b/packages/ai-persistence/package.json index 202416e8d..e76fc9f7f 100644 --- a/packages/ai-persistence/package.json +++ b/packages/ai-persistence/package.json @@ -28,6 +28,10 @@ ".": { "types": "./dist/esm/index.d.ts", "import": "./dist/esm/index.js" + }, + "./testkit": { + "types": "./dist/esm/testkit/conformance.d.ts", + "import": "./dist/esm/testkit/conformance.js" } }, "files": [ @@ -48,10 +52,17 @@ "@tanstack/ai-utils": "workspace:*" }, "peerDependencies": { - "@tanstack/ai": "workspace:^" + "@tanstack/ai": "workspace:^", + "vitest": "^4.0.14" + }, + "peerDependenciesMeta": { + "vitest": { + "optional": true + } }, "devDependencies": { "@tanstack/ai": "workspace:*", - "@vitest/coverage-v8": "4.0.14" + "@vitest/coverage-v8": "4.0.14", + "vitest": "^4.0.14" } } diff --git a/packages/ai-persistence/src/approval-controller.ts b/packages/ai-persistence/src/approval-controller.ts deleted file mode 100644 index 828ecc00c..000000000 --- a/packages/ai-persistence/src/approval-controller.ts +++ /dev/null @@ -1,71 +0,0 @@ -import type { ApprovalRecord, ApprovalStore, InterruptStore } from './types' - -/** @deprecated Use InterruptController. */ -export interface ApprovalController { - resolve: (approvalId: string, granted: boolean) => Promise - request: ( - record: Omit, - ) => Promise - decisionsForThread: (threadId: string) => Promise> -} - -/** @deprecated Use createInterruptController. */ -export function createApprovalController(opts: { - store: InterruptStore -}): ApprovalController -/** @deprecated Use createInterruptController. */ -export function createApprovalController(opts: { - store: ApprovalStore -}): ApprovalController -export function createApprovalController(opts: { - store: ApprovalStore | InterruptStore -}): ApprovalController { - const { store } = opts - if ('decisionsForThread' in store) { - return { - resolve: (approvalId, granted) => store.resolve(approvalId, granted), - request: (record) => - store.create({ - ...record, - status: 'pending', - }), - decisionsForThread: (threadId) => store.decisionsForThread(threadId), - } - } - - return { - async resolve(approvalId, granted) { - if (granted) { - await store.resolve(approvalId, { granted: true }) - } else { - await store.cancel(approvalId) - } - }, - async request(record) { - await store.create({ - interruptId: record.approvalId, - runId: record.runId, - threadId: record.threadId, - status: 'pending', - requestedAt: record.requestedAt, - payload: { - ...record.payload, - approvalId: record.approvalId, - compatibility: 'approval', - }, - }) - }, - async decisionsForThread(threadId) { - const decisions = new Map() - for (const record of await store.list(threadId)) { - if (record.payload.compatibility !== 'approval') continue - if (record.status === 'resolved') { - decisions.set(record.interruptId, true) - } else if (record.status === 'cancelled') { - decisions.set(record.interruptId, false) - } - } - return decisions - }, - } -} diff --git a/packages/ai-persistence/src/capabilities.ts b/packages/ai-persistence/src/capabilities.ts index db1260f7e..b1673a4ab 100644 --- a/packages/ai-persistence/src/capabilities.ts +++ b/packages/ai-persistence/src/capabilities.ts @@ -2,52 +2,22 @@ * Persistence capability tokens. * * `withChatPersistence` PROVIDES these so later middleware (and harness adapters) - * can read durable state. `LocksCapability` and `ResumeSourceCapability` are - * re-exported from core (`@tanstack/ai`) — they are shared, single-owner tokens - * (locks with the sandbox layer; the resume source with the chat engine's - * resume seam). + * can read durable state. `LocksCapability` is re-exported from core + * (`@tanstack/ai`) — a shared, single-owner token owned with the sandbox layer. */ import { createCapability } from '@tanstack/ai' -import type { AIPersistence, InterruptStore, PublicEventStore } from './types' +import type { AIPersistence, InterruptStore } from './types' export const PersistenceCapability = createCapability()('persistence') -export const EventsCapability = createCapability()( - 'persistence.publicEvents', -) - export const InterruptsCapability = createCapability()( 'persistence.interrupts', ) export const [getPersistence, providePersistence] = PersistenceCapability -export const [getEvents, provideEvents] = EventsCapability export const [getInterrupts, provideInterrupts] = InterruptsCapability -/** - * @deprecated Use InterruptsCapability. - * @alias - */ -export const ApprovalsCapability = InterruptsCapability -/** - * @deprecated Use getInterrupts. - * @alias - */ -export const getApprovals = getInterrupts -/** - * @deprecated Use provideInterrupts. - * @alias - */ -export const provideApprovals = provideInterrupts - // Shared, single-owner tokens live in core; re-export so consumers import // everything persistence-related from this package. -export { - LocksCapability, - getLocks, - provideLocks, - ResumeSourceCapability, - getResumeSource, - provideResumeSource, -} from '@tanstack/ai' +export { LocksCapability, getLocks, provideLocks } from '@tanstack/ai' diff --git a/packages/ai-persistence/src/cursor.ts b/packages/ai-persistence/src/cursor.ts deleted file mode 100644 index 8a79d7dc2..000000000 --- a/packages/ai-persistence/src/cursor.ts +++ /dev/null @@ -1,84 +0,0 @@ -/** - * Resume cursors. - * - * A cursor is an OPAQUE string the client echoes back to resume a run. It - * encodes a `(runId, seq)` pair where `seq` is a per-run monotonic sequence - * assigned to each persisted event. The opacity keeps the public contract free - * of any "cursor is an integer" assumption — backends compare by the decoded - * `seq`, and the wire format can evolve without breaking clients. - * - * Encoding: base64url of `":"`. The sequence comes first and has no - * colon, so the runId (which may contain colons) is everything after the first - * delimiter — robust regardless of runId content. - */ - -/** Decode the inner `":"` payload, or null if malformed. */ -function decodePayload(cursor: string): { runId: string; seq: number } | null { - let decoded: string - try { - decoded = Buffer.from(cursor, 'base64url').toString('utf8') - } catch { - return null - } - const delimiter = decoded.indexOf(':') - if (delimiter <= 0) { - return null - } - const seqText = decoded.slice(0, delimiter) - if (!/^\d+$/.test(seqText)) { - return null - } - return { seq: Number(seqText), runId: decoded.slice(delimiter + 1) } -} - -/** Encode a `(runId, seq)` pair into an opaque cursor string. */ -export function encodeCursor(runId: string, seq: number): string { - return Buffer.from(`${seq}:${runId}`, 'utf8').toString('base64url') -} - -/** Decode a cursor back to its `(runId, seq)` pair. Throws if malformed. */ -export function decodeCursor(cursor: string): { runId: string; seq: number } { - const payload = decodePayload(cursor) - if (payload === null) { - throw new Error(`Invalid resume cursor: ${cursor}`) - } - return payload -} - -/** Whether `value` is a well-formed cursor produced by {@link encodeCursor}. */ -export function isValidCursor(value: string): boolean { - return value.length > 0 && decodePayload(value) !== null -} - -/** - * Per-run monotonic sequence counter held by `withChatPersistence` for the lifetime - * of one run. `next()` assigns the next sequence to an event; on resume, - * construct it with the highest already-persisted sequence so new events keep - * climbing without colliding with replayed ones. - */ -export class RunSequence { - private seq: number - - constructor( - private readonly runId: string, - initialSeq = 0, - ) { - this.seq = initialSeq - } - - /** Assign and return the next sequence number. */ - next(): number { - this.seq += 1 - return this.seq - } - - /** The most recently assigned sequence number. */ - current(): number { - return this.seq - } - - /** Encode the current sequence as a cursor for this run. */ - toCursor(): string { - return encodeCursor(this.runId, this.seq) - } -} diff --git a/packages/ai-persistence/src/history.ts b/packages/ai-persistence/src/history.ts deleted file mode 100644 index 84118063b..000000000 --- a/packages/ai-persistence/src/history.ts +++ /dev/null @@ -1,25 +0,0 @@ -/** - * History projection — read a persisted run's events back as a `StreamChunk` - * timeline. This is the `events -> StreamChunk[]` projection devtools (and any - * client) consume to render or replay a PAST run through the exact same - * chunk-rendering path they use for live runs. Live devtools observation is - * unchanged; this adds the read-from-store side. - */ -import type { StreamChunk } from '@tanstack/ai' -import type { PublicEventStore } from './types' - -/** - * Collect a run's persisted events into an ordered `StreamChunk[]` timeline. - * `afterSeq` skips events up to and including that sequence (e.g. for paging). - */ -export async function loadRunHistory( - events: PublicEventStore, - runId: string, - opts?: { afterSeq?: number }, -): Promise> { - const chunks: Array = [] - for await (const { event } of events.read(runId, opts)) { - chunks.push(event) - } - return chunks -} diff --git a/packages/ai-persistence/src/index.ts b/packages/ai-persistence/src/index.ts index 02e49ec89..8cc98ebe4 100644 --- a/packages/ai-persistence/src/index.ts +++ b/packages/ai-persistence/src/index.ts @@ -1,6 +1,5 @@ // Store contracts + aggregate export { - AppendConflictError, defineAIPersistence, defineChatPersistence, validatePersistenceFeatures, @@ -8,21 +7,12 @@ export { export type { PersistenceMode, PersistenceFeature, - PersistedEvent, - PersistedPublicEvent, - PersistedInternalEvent, MessageStore, RunStatus, RunRecord, RunStore, - EventLog, - PublicEventStore, - InternalEventStore, - DurableRunStream, InterruptRecord, InterruptStore, - ApprovalRecord, - ApprovalStore, MetadataStore, ArtifactRecord, ArtifactStore, @@ -55,46 +45,19 @@ export type { // Reference in-memory implementation export { memoryPersistence } from './memory' -// Cursor utilities -export { - encodeCursor, - decodeCursor, - isValidCursor, - RunSequence, -} from './cursor' - -// Resume-source adapter (EventLog + RunStore -> core ResumeSource) -export { createResumeSource } from './resume-source' - -// History projection (events -> StreamChunk[] timeline for devtools / replay) -export { loadRunHistory } from './history' - // Interrupt controller export { createInterruptController } from './interrupts' export type { InterruptController } from './interrupts' -// Deprecated approval compatibility -export { createApprovalController } from './approval-controller' -export type { ApprovalController } from './approval-controller' - -// Capabilities (incl. re-exported core Locks/ResumeSource tokens) +// Capabilities (incl. re-exported core Locks token) export { PersistenceCapability, - EventsCapability, InterruptsCapability, - ApprovalsCapability, getPersistence, providePersistence, - getEvents, - provideEvents, getInterrupts, provideInterrupts, - getApprovals, - provideApprovals, LocksCapability, getLocks, provideLocks, - ResumeSourceCapability, - getResumeSource, - provideResumeSource, } from './capabilities' diff --git a/packages/ai-persistence/src/memory.ts b/packages/ai-persistence/src/memory.ts index 9bea5015d..be711e535 100644 --- a/packages/ai-persistence/src/memory.ts +++ b/packages/ai-persistence/src/memory.ts @@ -1,7 +1,6 @@ import { InMemoryLockStore } from '@tanstack/ai' -import { encodeCursor } from './cursor' -import { AppendConflictError, defineAIPersistence } from './types' -import type { ModelMessage, StreamChunk } from '@tanstack/ai' +import { defineAIPersistence } from './types' +import type { ModelMessage } from '@tanstack/ai' import type { AIPersistence, ArtifactRecord, @@ -11,14 +10,10 @@ import type { BlobObject, BlobRecord, BlobStore, - InternalEventStore, InterruptRecord, InterruptStore, MessageStore, MetadataStore, - PersistedInternalEvent, - PersistedPublicEvent, - PublicEventStore, RunRecord, RunStore, } from './types' @@ -68,158 +63,6 @@ class MemoryRunStore implements RunStore { } } -function stableJsonValue(value: unknown): unknown { - if (Array.isArray(value)) { - return value.map((item) => stableJsonValue(item)) - } - if (value && typeof value === 'object') { - return Object.fromEntries( - Object.entries(value) - .sort(([a], [b]) => a.localeCompare(b)) - .map(([key, item]) => [key, stableJsonValue(item)]), - ) - } - return value -} - -function sameJsonValue(a: unknown, b: unknown): boolean { - return ( - JSON.stringify(stableJsonValue(a)) === JSON.stringify(stableJsonValue(b)) - ) -} - -class MemoryPublicEventStore implements PublicEventStore { - private readonly logs = new Map>() - append(input: { - runId: string - expectedSeq: number - event: StreamChunk - }): Promise { - const log = this.logs.get(input.runId) ?? [] - const targetSeq = input.expectedSeq + 1 - const existingAtTarget = log.find((e) => e.seq === targetSeq) - if (existingAtTarget) { - if (sameJsonValue(existingAtTarget.event, input.event)) { - return Promise.resolve(existingAtTarget) - } - return Promise.reject( - new AppendConflictError( - `Public event append conflict for run ${input.runId} at seq ${targetSeq}`, - ), - ) - } - - const latest = log.at(-1)?.seq ?? 0 - if (latest !== input.expectedSeq) { - return Promise.reject( - new AppendConflictError( - `Public event append conflict for run ${input.runId}: expected latest seq ${input.expectedSeq}, got ${latest}`, - ), - ) - } - - const persisted = { - seq: targetSeq, - event: input.event, - cursor: input.event.cursor ?? encodeCursor(input.runId, targetSeq), - } - log.push(persisted) - this.logs.set(input.runId, log) - return Promise.resolve(persisted) - } - read( - runId: string, - opts?: { afterSeq?: number }, - ): AsyncIterable { - const after = opts?.afterSeq ?? -Infinity - const events = (this.logs.get(runId) ?? []).filter((e) => e.seq > after) - return (async function* () { - await Promise.resolve() - for (const e of events) yield e - })() - } - hasRun(runId: string): Promise { - return Promise.resolve((this.logs.get(runId)?.length ?? 0) > 0) - } - latestSeq(runId: string): Promise { - return Promise.resolve(this.logs.get(runId)?.at(-1)?.seq ?? 0) - } -} - -class MemoryInternalEventStore implements InternalEventStore { - private readonly logs = new Map>() - append(input: { - runId: string - expectedSeq: number - namespace: string - type: string - payload: unknown - }): Promise { - const log = this.logs.get(input.runId) ?? [] - const targetSeq = input.expectedSeq + 1 - const existingAtTarget = log.find( - (e) => e.namespace === input.namespace && e.seq === targetSeq, - ) - if (existingAtTarget) { - if ( - existingAtTarget.type === input.type && - sameJsonValue(existingAtTarget.payload, input.payload) - ) { - return Promise.resolve(existingAtTarget) - } - return Promise.reject( - new AppendConflictError( - `Internal event append conflict for run ${input.runId} namespace ${input.namespace} at seq ${targetSeq}`, - ), - ) - } - - const latest = - log.filter((e) => e.namespace === input.namespace).at(-1)?.seq ?? 0 - if (latest !== input.expectedSeq) { - return Promise.reject( - new AppendConflictError( - `Internal event append conflict for run ${input.runId} namespace ${input.namespace}: expected latest seq ${input.expectedSeq}, got ${latest}`, - ), - ) - } - - const persisted = { - seq: targetSeq, - namespace: input.namespace, - type: input.type, - payload: input.payload, - cursor: encodeCursor(input.runId, targetSeq), - } - log.push(persisted) - this.logs.set(input.runId, log) - return Promise.resolve(persisted) - } - read( - runId: string, - opts?: { namespace?: string; afterSeq?: number }, - ): AsyncIterable { - const after = opts?.afterSeq ?? -Infinity - const events = (this.logs.get(runId) ?? []).filter( - (e) => - e.seq > after && - (opts?.namespace === undefined || e.namespace === opts.namespace), - ) - return (async function* () { - await Promise.resolve() - for (const e of events) yield e - })() - } - latestSeq(runId: string, namespace?: string): Promise { - const log = this.logs.get(runId) ?? [] - const filtered = - namespace === undefined - ? log - : log.filter((event) => event.namespace === namespace) - return Promise.resolve(filtered.at(-1)?.seq ?? 0) - } -} - class MemoryInterruptStore implements InterruptStore { private readonly interrupts = new Map() create(record: Omit): Promise { @@ -501,8 +344,6 @@ export function memoryPersistence(): AIPersistence { stores: { messages: new MemoryMessageStore(), runs: new MemoryRunStore(), - publicEvents: new MemoryPublicEventStore(), - internalEvents: new MemoryInternalEventStore(), interrupts: new MemoryInterruptStore(), metadata: new MemoryMetadataStore(), artifacts: new MemoryArtifactStore(), diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index 711916816..ea55ffb42 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -1,19 +1,13 @@ import { defineChatMiddleware } from '@tanstack/ai' import { base64ToUint8Array } from '@tanstack/ai-utils' import { - EventsCapability, InterruptsCapability, LocksCapability, PersistenceCapability, - ResumeSourceCapability, - provideEvents, provideInterrupts, provideLocks, providePersistence, - provideResumeSource, } from './capabilities' -import { RunSequence, decodeCursor, encodeCursor } from './cursor' -import { createResumeSource } from './resume-source' import { validatePersistenceFeatures } from './types' import type { AbortInfo, @@ -89,7 +83,7 @@ export interface GenerationArtifactNameInput { const runState = new WeakMap< object, - { seq: RunSequence; merged: boolean; interrupted: boolean } + { merged: boolean; interrupted: boolean } >() function defaultFeatures( @@ -97,17 +91,9 @@ function defaultFeatures( ): Array { const features: Array = [] if (persistence.stores.messages) features.push('messages') - if (persistence.stores.runs && persistence.stores.publicEvents) { - features.push('durable-replay') - } - if ( - persistence.stores.runs && - persistence.stores.publicEvents && - persistence.stores.interrupts - ) { + if (persistence.stores.runs && persistence.stores.interrupts) { features.push('interrupts') } - if (persistence.stores.internalEvents) features.push('internal-events') if (persistence.stores.metadata) features.push('metadata') if (persistence.stores.locks) features.push('locks') if (persistence.stores.artifacts && persistence.stores.blobs) { @@ -245,79 +231,6 @@ function interruptPayload(interrupt: unknown): Record { : { value: interrupt } } -async function publicEventExistsAtSeq( - publicEvents: NonNullable, - runId: string, - seq: number, -): Promise { - for await (const persisted of publicEvents.read(runId, { - afterSeq: seq - 1, - })) { - return persisted.seq === seq - } - return false -} - -async function validateReplayCursor( - ctx: ChatMiddlewareContext, - cursor: string | undefined, - persistence: AIPersistence, -): Promise { - if (!cursor) return false - if (!persistence.stores.runs || !persistence.stores.publicEvents) { - return false - } - - const decoded = decodeCursor(cursor) - if (decoded.seq < 1) { - throw new Error( - `Resume cursor sequence ${decoded.seq} is invalid; expected a persisted event sequence >= 1.`, - ) - } - if (decoded.runId !== ctx.runId) { - throw new Error( - `Resume cursor runId ${decoded.runId} does not match request runId ${ctx.runId}.`, - ) - } - - const run = await persistence.stores.runs.get(decoded.runId) - if (!run) { - throw new Error(`Resume cursor references unknown run ${decoded.runId}.`) - } - if (run.threadId !== ctx.threadId) { - throw new Error( - `Resume cursor run ${decoded.runId} belongs to thread ${run.threadId}, not request thread ${ctx.threadId}.`, - ) - } - - const latestSeq = await persistence.stores.publicEvents.latestSeq( - decoded.runId, - ) - if (latestSeq === 0) { - throw new Error( - `Resume cursor references run ${decoded.runId}, but no public events are persisted.`, - ) - } - if (decoded.seq > latestSeq) { - throw new Error( - `Resume cursor sequence ${decoded.seq} is beyond latest persisted sequence ${latestSeq} for run ${decoded.runId}.`, - ) - } - if ( - !(await publicEventExistsAtSeq( - persistence.stores.publicEvents, - decoded.runId, - decoded.seq, - )) - ) { - throw new Error( - `Resume cursor sequence ${decoded.seq} does not reference a persisted public event for run ${decoded.runId}.`, - ) - } - - return true -} - function isArtifactRef(value: unknown): value is PersistedArtifactRef { const record = objectValue(value) return !!record && typeof record.artifactId === 'string' @@ -766,12 +679,9 @@ async function persistGenerationArtifacts( interface PersistencePlan { features: Array wantsMessages: boolean - wantsReplay: boolean wantsInterrupts: boolean - wantsPublicEvents: boolean wantsLocks: boolean wantsArtifactPersistence: boolean - publicEvents: AIPersistence['stores']['publicEvents'] runs: AIPersistence['stores']['runs'] } @@ -783,9 +693,7 @@ function resolvePersistencePlan( validatePersistenceFeatures(persistence, features) const wantsMessages = features.includes('messages') - const wantsReplay = features.includes('durable-replay') const wantsInterrupts = features.includes('interrupts') - const wantsPublicEvents = wantsReplay || wantsInterrupts const wantsLocks = features.includes('locks') const wantsArtifactStore = features.includes('artifacts') const wantsBlobStore = features.includes('blobs') @@ -798,12 +706,9 @@ function resolvePersistencePlan( return { features, wantsMessages, - wantsReplay, wantsInterrupts, - wantsPublicEvents, wantsLocks, wantsArtifactPersistence: wantsArtifactStore && wantsBlobStore, - publicEvents: persistence.stores.publicEvents, runs: persistence.stores.runs, } } @@ -859,26 +764,20 @@ async function interruptRun( // --------------------------------------------------------------------------- /** - * Chat-only persistence middleware. Provides durable messages, public event - * replay, interrupts, and locks for `chat()`. + * Chat-only persistence middleware. Provides durable **state** for `chat()`: + * thread messages, run records, interrupts, and locks. Delivery durability + * (replaying a disconnected/reloaded stream) lives on the transport layer via + * `StreamDurability`, not here — this middleware never mutates the chunk stream. */ export function withChatPersistence( persistence: AIPersistence, opts?: WithPersistenceOptions, ): ChatMiddleware { const plan = resolvePersistencePlan(persistence, opts) - const { - wantsMessages, - wantsInterrupts, - wantsPublicEvents, - wantsLocks, - publicEvents, - runs, - } = plan + const { wantsMessages, wantsInterrupts, wantsLocks, runs } = plan const provides = [ PersistenceCapability, - ...(wantsPublicEvents ? [EventsCapability, ResumeSourceCapability] : []), ...(wantsInterrupts ? [InterruptsCapability] : []), ...(wantsLocks ? [LocksCapability] : []), ] @@ -886,23 +785,14 @@ export function withChatPersistence( return defineChatMiddleware({ name: 'chat-persistence', provides, - async setup(ctx: ChatMiddlewareContext) { + setup(ctx: ChatMiddlewareContext) { providePersistence(ctx, persistence) - const initialSeq = - wantsPublicEvents && publicEvents - ? await publicEvents.latestSeq(ctx.runId) - : 0 runState.set(ctx, { - seq: new RunSequence(ctx.runId, initialSeq), merged: false, interrupted: false, }) - if (wantsPublicEvents && publicEvents) { - provideEvents(ctx, publicEvents) - provideResumeSource(ctx, createResumeSource(publicEvents, runs)) - } if (wantsInterrupts && persistence.stores.interrupts) { provideInterrupts(ctx, persistence.stores.interrupts) } @@ -914,17 +804,9 @@ export function withChatPersistence( async onConfig(ctx: ChatMiddlewareContext, config: ChatMiddlewareConfig) { if (ctx.phase !== 'init') return - const hasResume = config.resume !== undefined - const isReplay = wantsPublicEvents - ? await validateReplayCursor(ctx, config.cursor, persistence) - : false let resumeToolState: ChatResumeToolState | undefined - if ( - wantsInterrupts && - persistence.stores.interrupts && - (!isReplay || hasResume) - ) { + if (wantsInterrupts && persistence.stores.interrupts) { const pending = await persistence.stores.interrupts.listPending( ctx.threadId, ) @@ -963,48 +845,37 @@ export function withChatPersistence( }, async onChunk(ctx: ChatMiddlewareContext, chunk: StreamChunk) { - if (!wantsPublicEvents || !publicEvents) return + // State-only: react to the interrupt boundary (create interrupt records, + // mark the run interrupted, snapshot thread messages). The chunk stream is + // never mutated — delivery durability is a transport-layer concern. + if ( + chunk.type !== 'RUN_FINISHED' || + chunk.outcome?.type !== 'interrupt' + ) { + return + } const state = runState.get(ctx) if (!state) return - const expectedSeq = state.seq.current() - const seq = state.seq.next() - const stamped: StreamChunk = { - ...chunk, - cursor: encodeCursor(ctx.runId, seq), - } - await publicEvents.append({ - runId: ctx.runId, - expectedSeq, - event: stamped, - }) - await persistence.stream?.publish(ctx.runId, seq, stamped) - if ( - stamped.type === 'RUN_FINISHED' && - stamped.outcome?.type === 'interrupt' - ) { - if (wantsInterrupts && persistence.stores.interrupts) { - for (const interrupt of stamped.outcome.interrupts) { - await persistence.stores.interrupts.create({ - interruptId: interrupt.id, - runId: ctx.runId, - threadId: ctx.threadId, - status: 'pending', - requestedAt: Date.now(), - payload: interruptPayload(interrupt), - }) - } - } - await interruptRun(runs, ctx.runId) - if (wantsMessages && persistence.stores.messages) { - await persistence.stores.messages.saveThread(ctx.threadId, [ - ...ctx.messages, - ]) + if (wantsInterrupts && persistence.stores.interrupts) { + for (const interrupt of chunk.outcome.interrupts) { + await persistence.stores.interrupts.create({ + interruptId: interrupt.id, + runId: ctx.runId, + threadId: ctx.threadId, + status: 'pending', + requestedAt: Date.now(), + payload: interruptPayload(interrupt), + }) } - state.interrupted = true } - - return stamped + await interruptRun(runs, ctx.runId) + if (wantsMessages && persistence.stores.messages) { + await persistence.stores.messages.saveThread(ctx.threadId, [ + ...ctx.messages, + ]) + } + state.interrupted = true }, async onFinish(ctx: ChatMiddlewareContext, info: FinishInfo) { diff --git a/packages/ai-persistence/src/resume-source.ts b/packages/ai-persistence/src/resume-source.ts deleted file mode 100644 index b9d51b8c6..000000000 --- a/packages/ai-persistence/src/resume-source.ts +++ /dev/null @@ -1,78 +0,0 @@ -/** - * Adapts the persistence {@link PublicEventStore} (+ optional {@link RunStore}) to the - * core `ResumeSource` contract the `chat()` resume seam consumes. Cursors are - * decoded to a per-run sequence so replay returns only events after the point - * the client last saw. - */ -import { decodeCursor } from './cursor' -import type { ResumeSource, RunStatus, StreamChunk } from '@tanstack/ai' -import type { PublicEventStore, RunStore } from './types' - -async function publicEventExistsAtSeq( - events: PublicEventStore, - runId: string, - seq: number, -): Promise { - for await (const persisted of events.read(runId, { afterSeq: seq - 1 })) { - return persisted.seq === seq - } - return false -} - -export function createResumeSource( - events: PublicEventStore, - runs?: RunStore, -): ResumeSource { - return { - async hasRun(runId) { - if (runs && !(await runs.get(runId))) return false - return events.hasRun(runId) - }, - async *replay(runId, afterCursor): AsyncIterable { - if (runs && !(await runs.get(runId))) { - throw new Error(`Resume replay references unknown run ${runId}.`) - } - let afterSeq: number | undefined - if (afterCursor) { - const decoded = decodeCursor(afterCursor) - if (decoded.seq < 1) { - throw new Error( - `Resume cursor sequence ${decoded.seq} is invalid; expected a persisted event sequence >= 1.`, - ) - } - if (decoded.runId !== runId) { - throw new Error( - `Resume cursor runId ${decoded.runId} does not match replay runId ${runId}.`, - ) - } - const latestSeq = await events.latestSeq(runId) - if (latestSeq === 0) { - throw new Error( - `Resume cursor references run ${runId}, but no public events are persisted.`, - ) - } - if (decoded.seq > latestSeq) { - throw new Error( - `Resume cursor sequence ${decoded.seq} is beyond latest persisted sequence ${latestSeq} for run ${runId}.`, - ) - } - if (!(await publicEventExistsAtSeq(events, runId, decoded.seq))) { - throw new Error( - `Resume cursor sequence ${decoded.seq} does not reference a persisted public event for run ${runId}.`, - ) - } - afterSeq = decoded.seq - } - for await (const { event } of events.read( - runId, - afterSeq === undefined ? undefined : { afterSeq }, - )) { - yield event - } - }, - async getStatus(runId): Promise { - const run = await runs?.get(runId) - return run ? run.status : null - }, - } -} diff --git a/packages/ai-persistence/src/testkit/conformance.ts b/packages/ai-persistence/src/testkit/conformance.ts new file mode 100644 index 000000000..6814bd4b3 --- /dev/null +++ b/packages/ai-persistence/src/testkit/conformance.ts @@ -0,0 +1,352 @@ +/** + * Shared conformance suite for the `AIPersistence` state contract. + * + * Every backend (memory, drizzle, prisma, …) runs this identical suite so that + * schema drift or an implementation gap fails immediately. It exercises every + * method of every store the persistence exposes; stores that are absent (a + * backend may not provide `locks`, for example) are skipped, so the suite is + * universal across backends. + * + * Artifacts are tested via `externalUrl` (not inline `bytes`) so the assertions + * hold identically for byte-storing (memory) and reference-only (SQL) backends. + */ +import { beforeAll, describe, expect, it } from 'vitest' +import type { AIPersistence } from '../types' + +type MakePersistence = () => Promise | AIPersistence + +/** + * Register a Vitest suite that validates `makePersistence()` against the full + * `AIPersistence` contract. + */ +export function runPersistenceConformance( + name: string, + makePersistence: MakePersistence, +): void { + describe(`AIPersistence conformance: ${name}`, () => { + let persistence: AIPersistence + + beforeAll(async () => { + persistence = await makePersistence() + }) + + describe('messages', () => { + it('round-trips a thread and returns [] for unknown threads', async () => { + const store = persistence.stores.messages + if (!store) return + + expect(await store.loadThread('thread-unknown')).toEqual([]) + + await store.saveThread('thread-msg', [ + { role: 'user', content: 'hi' }, + { role: 'assistant', content: 'hello' }, + ]) + expect(await store.loadThread('thread-msg')).toEqual([ + { role: 'user', content: 'hi' }, + { role: 'assistant', content: 'hello' }, + ]) + + // Overwrites, not appends. + await store.saveThread('thread-msg', [{ role: 'user', content: 'redo' }]) + expect(await store.loadThread('thread-msg')).toEqual([ + { role: 'user', content: 'redo' }, + ]) + }) + }) + + describe('runs', () => { + it('creates, resumes idempotently, updates, and gets', async () => { + const store = persistence.stores.runs + if (!store) return + + expect(await store.get('run-missing')).toBeNull() + + const created = await store.createOrResume({ + runId: 'run-1', + threadId: 'thread-1', + startedAt: 1000, + }) + expect(created).toMatchObject({ + runId: 'run-1', + threadId: 'thread-1', + status: 'running', + startedAt: 1000, + }) + + // createOrResume is idempotent: returns the existing record unchanged. + const resumed = await store.createOrResume({ + runId: 'run-1', + threadId: 'thread-different', + startedAt: 9999, + }) + expect(resumed).toMatchObject({ + runId: 'run-1', + threadId: 'thread-1', + startedAt: 1000, + }) + + await store.update('run-1', { + status: 'completed', + finishedAt: 2000, + usage: { promptTokens: 3, completionTokens: 4, totalTokens: 7 }, + }) + const done = await store.get('run-1') + expect(done).toMatchObject({ + runId: 'run-1', + status: 'completed', + finishedAt: 2000, + usage: { promptTokens: 3, completionTokens: 4, totalTokens: 7 }, + }) + + await store.update('run-1', { status: 'failed', error: 'boom' }) + const failed = await store.get('run-1') + expect(failed?.status).toBe('failed') + expect(failed?.error).toBe('boom') + + // Updating a missing run is a no-op (does not throw, does not create). + await store.update('run-absent', { status: 'completed' }) + expect(await store.get('run-absent')).toBeNull() + }) + }) + + describe('interrupts', () => { + it('creates, resolves, cancels, and lists by thread and run', async () => { + const store = persistence.stores.interrupts + if (!store) return + + expect(await store.get('int-missing')).toBeNull() + + await store.create({ + interruptId: 'int-1', + runId: 'run-i', + threadId: 'thread-i', + status: 'pending', + requestedAt: 10, + payload: { tool: 'search', args: { q: 'x' } }, + }) + await store.create({ + interruptId: 'int-2', + runId: 'run-i', + threadId: 'thread-i', + status: 'pending', + requestedAt: 20, + payload: { tool: 'write' }, + }) + await store.create({ + interruptId: 'int-3', + runId: 'run-other', + threadId: 'thread-i', + status: 'pending', + requestedAt: 30, + payload: {}, + }) + + const one = await store.get('int-1') + expect(one).toMatchObject({ + interruptId: 'int-1', + runId: 'run-i', + threadId: 'thread-i', + status: 'pending', + requestedAt: 10, + payload: { tool: 'search', args: { q: 'x' } }, + }) + + expect((await store.list('thread-i')).map((r) => r.interruptId)).toEqual( + ['int-1', 'int-2', 'int-3'], + ) + expect( + (await store.listByRun('run-i')).map((r) => r.interruptId), + ).toEqual(['int-1', 'int-2']) + expect( + (await store.listPending('thread-i')).map((r) => r.interruptId), + ).toEqual(['int-1', 'int-2', 'int-3']) + + await store.resolve('int-1', { ok: true }) + const resolved = await store.get('int-1') + expect(resolved?.status).toBe('resolved') + expect(resolved?.response).toEqual({ ok: true }) + expect(typeof resolved?.resolvedAt).toBe('number') + + await store.cancel('int-2') + const cancelled = await store.get('int-2') + expect(cancelled?.status).toBe('cancelled') + expect(typeof cancelled?.resolvedAt).toBe('number') + + expect( + (await store.listPending('thread-i')).map((r) => r.interruptId), + ).toEqual(['int-3']) + expect( + (await store.listPendingByRun('run-i')).map((r) => r.interruptId), + ).toEqual([]) + }) + }) + + describe('metadata', () => { + it('sets, gets, scopes, and deletes', async () => { + const store = persistence.stores.metadata + if (!store) return + + expect(await store.get('scope-a', 'k')).toBeNull() + + await store.set('scope-a', 'k', { n: 1 }) + await store.set('scope-b', 'k', { n: 2 }) + expect(await store.get('scope-a', 'k')).toEqual({ n: 1 }) + expect(await store.get('scope-b', 'k')).toEqual({ n: 2 }) + + await store.set('scope-a', 'k', { n: 3 }) + expect(await store.get('scope-a', 'k')).toEqual({ n: 3 }) + + await store.delete('scope-a', 'k') + expect(await store.get('scope-a', 'k')).toBeNull() + // Delete is scoped: scope-b untouched. + expect(await store.get('scope-b', 'k')).toEqual({ n: 2 }) + }) + }) + + describe('artifacts', () => { + it('saves, gets, lists by run, and deletes', async () => { + const store = persistence.stores.artifacts + if (!store) return + + expect(await store.get('art-missing')).toBeNull() + + await store.save({ + artifactId: 'art-1', + runId: 'run-art', + threadId: 'thread-art', + name: 'image.png', + mimeType: 'image/png', + size: 123, + externalUrl: 'https://example.com/image.png', + createdAt: 500, + }) + await store.save({ + artifactId: 'art-2', + runId: 'run-art', + threadId: 'thread-art', + name: 'clip.mp4', + mimeType: 'video/mp4', + size: 456, + externalUrl: 'https://example.com/clip.mp4', + createdAt: 600, + }) + + const got = await store.get('art-1') + expect(got).toMatchObject({ + artifactId: 'art-1', + runId: 'run-art', + threadId: 'thread-art', + name: 'image.png', + mimeType: 'image/png', + size: 123, + externalUrl: 'https://example.com/image.png', + createdAt: 500, + }) + + expect((await store.list('run-art')).map((a) => a.artifactId).sort()).toEqual( + ['art-1', 'art-2'], + ) + expect(await store.list('run-none')).toEqual([]) + + if (store.delete) { + await store.delete('art-1') + expect(await store.get('art-1')).toBeNull() + expect((await store.list('run-art')).map((a) => a.artifactId)).toEqual( + ['art-2'], + ) + } + if (store.deleteForRun) { + await store.deleteForRun('run-art') + expect(await store.list('run-art')).toEqual([]) + } + }) + }) + + describe('blobs', () => { + it('puts, gets, heads, lists, and deletes', async () => { + const store = persistence.stores.blobs + if (!store) return + + const put = await store.put('blobs/a.txt', 'hello world', { + contentType: 'text/plain', + customMetadata: { owner: 'alice' }, + }) + expect(put.key).toBe('blobs/a.txt') + expect(put.size).toBe('hello world'.length) + expect(put.contentType).toBe('text/plain') + + const got = await store.get('blobs/a.txt') + expect(got).not.toBeNull() + expect(await got?.text()).toBe('hello world') + expect(got?.customMetadata).toEqual({ owner: 'alice' }) + + const head = await store.head('blobs/a.txt') + expect(head?.key).toBe('blobs/a.txt') + expect(head?.size).toBe('hello world'.length) + + expect(await store.get('blobs/missing')).toBeNull() + expect(await store.head('blobs/missing')).toBeNull() + + await store.put('blobs/b.txt', 'second') + await store.put('other/c.txt', 'third') + + const page = await store.list({ prefix: 'blobs/' }) + expect(page.objects.map((o) => o.key).sort()).toEqual([ + 'blobs/a.txt', + 'blobs/b.txt', + ]) + + const limited = await store.list({ prefix: 'blobs/', limit: 1 }) + expect(limited.objects).toHaveLength(1) + expect(limited.truncated).toBe(true) + expect(limited.cursor).toBeDefined() + + const rest = await store.list({ + prefix: 'blobs/', + cursor: limited.cursor, + }) + expect(rest.objects.map((o) => o.key)).not.toContain( + limited.objects[0]?.key, + ) + + await store.delete('blobs/a.txt') + expect(await store.get('blobs/a.txt')).toBeNull() + }) + + it('matches list prefixes literally and case-sensitively', async () => { + const store = persistence.stores.blobs + if (!store) return + + // `run_` contains a SQL LIKE metacharacter (`_`), and the differing-case + // keys probe case sensitivity. Only the keys that literally start with + // the exact bytes `run_` must come back — `_` must NOT match any single + // character, and `RUN_1/...` (upper-case) must be excluded. + await store.put('run_1/x', 'a') + await store.put('run_2/y', 'b') + await store.put('runX/z', 'c') // `_` as wildcard would wrongly match this + await store.put('RUN_1/w', 'd') // case-insensitive match would include this + + const page = await store.list({ prefix: 'run_' }) + expect(page.objects.map((o) => o.key).sort()).toEqual([ + 'run_1/x', + 'run_2/y', + ]) + }) + }) + + describe('locks', () => { + it('runs the critical section and returns its value', async () => { + const store = persistence.stores.locks + if (!store) return + + const order: Array = [] + const result = await store.withLock('lock-key', async () => { + order.push('inside') + return 42 + }) + expect(result).toBe(42) + expect(order).toEqual(['inside']) + }) + }) + }) +} diff --git a/packages/ai-persistence/src/types.ts b/packages/ai-persistence/src/types.ts index f54dc81a7..7c3bb4415 100644 --- a/packages/ai-persistence/src/types.ts +++ b/packages/ai-persistence/src/types.ts @@ -1,46 +1,15 @@ -import type { - LockStore, - ModelMessage, - StreamChunk, - TokenUsage, -} from '@tanstack/ai' +import type { LockStore, ModelMessage, TokenUsage } from '@tanstack/ai' export type PersistenceMode = 'messages' | 'chat' | 'agent' export type PersistenceFeature = | 'messages' - | 'durable-replay' | 'interrupts' - | 'internal-events' | 'metadata' | 'locks' | 'artifacts' | 'blobs' -export class AppendConflictError extends Error { - constructor(message: string) { - super(message) - this.name = 'AppendConflictError' - } -} - -export interface PersistedPublicEvent { - seq: number - event: StreamChunk - cursor: string -} - -export interface PersistedInternalEvent { - seq: number - namespace: string - type: string - payload: unknown - cursor: string -} - -/** @deprecated Use PersistedPublicEvent. */ -export type PersistedEvent = PersistedPublicEvent - export interface MessageStore { loadThread: (threadId: string) => Promise> saveThread: (threadId: string, messages: Array) => Promise @@ -74,42 +43,6 @@ export interface RunStore { get: (runId: string) => Promise } -export interface PublicEventStore { - append: (input: { - runId: string - expectedSeq: number - event: StreamChunk - }) => Promise - read: ( - runId: string, - opts?: { afterSeq?: number }, - ) => AsyncIterable - hasRun: (runId: string) => Promise - latestSeq: (runId: string) => Promise -} - -export interface InternalEventStore { - append: (input: { - runId: string - expectedSeq: number - namespace: string - type: string - payload: unknown - }) => Promise - read: ( - runId: string, - opts?: { namespace?: string; afterSeq?: number }, - ) => AsyncIterable - latestSeq: (runId: string, namespace?: string) => Promise -} - -/** @deprecated Use PublicEventStore. */ -export type EventLog = PublicEventStore - -export interface DurableRunStream { - publish: (runId: string, seq: number, event: StreamChunk) => Promise -} - export interface InterruptRecord { interruptId: string runId: string @@ -132,25 +65,6 @@ export interface InterruptStore { listPendingByRun: (runId: string) => Promise> } -/** @deprecated Use InterruptRecord. */ -export interface ApprovalRecord { - approvalId: string - runId: string - threadId: string - status: 'pending' | 'granted' | 'denied' - requestedAt: number - resolvedAt?: number - payload: Record -} - -/** @deprecated Use InterruptStore. */ -export interface ApprovalStore { - create: (record: Omit) => Promise - resolve: (approvalId: string, granted: boolean) => Promise - get: (approvalId: string) => Promise - decisionsForThread: (threadId: string) => Promise> -} - export interface MetadataStore { get: (scope: string, key: string) => Promise set: (scope: string, key: string, value: unknown) => Promise @@ -234,25 +148,23 @@ export interface AIPersistence { stores: { messages?: MessageStore runs?: RunStore - publicEvents?: PublicEventStore - internalEvents?: InternalEventStore interrupts?: InterruptStore metadata?: MetadataStore locks?: LockStore artifacts?: ArtifactStore blobs?: BlobStore } - stream?: DurableRunStream } /** @deprecated Use AIPersistence. */ export type ChatPersistence = AIPersistence -const featureRequirements: Record> = { +const featureRequirements: Record< + PersistenceFeature, + Array +> = { messages: ['messages'], - 'durable-replay': ['runs', 'publicEvents'], - interrupts: ['runs', 'publicEvents', 'interrupts'], - 'internal-events': ['internalEvents'], + interrupts: ['runs', 'interrupts'], metadata: ['metadata'], locks: ['locks'], artifacts: ['artifacts'], @@ -263,10 +175,13 @@ export function validatePersistenceFeatures( persistence: AIPersistence, features: Array, ): void { - const missing = new Map>() + const missing = new Map< + PersistenceFeature, + Array + >() for (const feature of features) { const missingStores = featureRequirements[feature].filter( - (store) => !persistence.stores[store as keyof AIPersistence['stores']], + (store) => !persistence.stores[store], ) if (missingStores.length > 0) { missing.set(feature, missingStores) diff --git a/packages/ai-persistence/tests/cursor.test.ts b/packages/ai-persistence/tests/cursor.test.ts deleted file mode 100644 index 953a23129..000000000 --- a/packages/ai-persistence/tests/cursor.test.ts +++ /dev/null @@ -1,65 +0,0 @@ -import { describe, expect, it } from 'vitest' -import { - RunSequence, - decodeCursor, - encodeCursor, - isValidCursor, -} from '../src/cursor' - -describe('encodeCursor / decodeCursor', () => { - it('round-trips a runId and sequence', () => { - const cursor = encodeCursor('run-1', 42) - expect(decodeCursor(cursor)).toEqual({ runId: 'run-1', seq: 42 }) - }) - - it('is opaque (not the raw "runId:seq")', () => { - const cursor = encodeCursor('run-1', 42) - expect(cursor).not.toContain('run-1') - expect(cursor).not.toBe('run-1:42') - }) - - it('round-trips runIds that contain the delimiter', () => { - const cursor = encodeCursor('thread:run:1', 7) - expect(decodeCursor(cursor)).toEqual({ runId: 'thread:run:1', seq: 7 }) - }) - - it('orders by sequence within a run via the decoded seq', () => { - const a = decodeCursor(encodeCursor('r', 1)).seq - const b = decodeCursor(encodeCursor('r', 2)).seq - expect(b).toBeGreaterThan(a) - }) -}) - -describe('isValidCursor', () => { - it('accepts a value produced by encodeCursor', () => { - expect(isValidCursor(encodeCursor('run-1', 1))).toBe(true) - }) - - it('rejects arbitrary / malformed strings', () => { - expect(isValidCursor('not-a-cursor')).toBe(false) - expect(isValidCursor('')).toBe(false) - }) -}) - -describe('RunSequence', () => { - it('hands out monotonically increasing sequence numbers', () => { - const seq = new RunSequence('run-1') - expect(seq.next()).toBe(1) - expect(seq.next()).toBe(2) - expect(seq.next()).toBe(3) - expect(seq.current()).toBe(3) - }) - - it('resumes after an initial sequence (so resumed runs keep climbing)', () => { - const seq = new RunSequence('run-1', 10) - expect(seq.next()).toBe(11) - expect(seq.current()).toBe(11) - }) - - it('toCursor encodes the current sequence for this run', () => { - const seq = new RunSequence('run-1') - seq.next() // 1 - seq.next() // 2 - expect(decodeCursor(seq.toCursor())).toEqual({ runId: 'run-1', seq: 2 }) - }) -}) diff --git a/packages/ai-persistence/tests/generation-artifacts.test.ts b/packages/ai-persistence/tests/generation-artifacts.test.ts index a1c0ae628..ae95cff20 100644 --- a/packages/ai-persistence/tests/generation-artifacts.test.ts +++ b/packages/ai-persistence/tests/generation-artifacts.test.ts @@ -215,7 +215,6 @@ describe('withGenerationPersistence generation artifacts', () => { const persistence = defineAIPersistence({ stores: { runs: full.stores.runs, - publicEvents: full.stores.publicEvents, }, }) diff --git a/packages/ai-persistence/tests/history.test.ts b/packages/ai-persistence/tests/history.test.ts deleted file mode 100644 index 25bb1c827..000000000 --- a/packages/ai-persistence/tests/history.test.ts +++ /dev/null @@ -1,63 +0,0 @@ -import { describe, expect, it } from 'vitest' -import { EventType } from '@tanstack/ai' -import type { StreamChunk } from '@tanstack/ai' -import { memoryPersistence } from '../src/memory' -import { loadRunHistory } from '../src/history' - -const text = (delta: string): StreamChunk => ({ - type: EventType.TEXT_MESSAGE_CONTENT, - messageId: 'm1', - delta, - timestamp: 1, -}) - -describe('loadRunHistory', () => { - it('projects a persisted run into an ordered StreamChunk timeline', async () => { - const { publicEvents, internalEvents } = memoryPersistence().stores - await publicEvents!.append({ - runId: 'r1', - expectedSeq: 0, - event: text('a'), - }) - await internalEvents!.append({ - runId: 'r1', - expectedSeq: 0, - namespace: 'checkpoint', - type: 'saved', - payload: { hidden: true }, - }) - await publicEvents!.append({ - runId: 'r1', - expectedSeq: 1, - event: text('b'), - }) - - const timeline = await loadRunHistory(publicEvents!, 'r1') - expect( - timeline.map((c) => - c.type === 'TEXT_MESSAGE_CONTENT' ? c.delta : c.type, - ), - ).toEqual(['a', 'b']) - }) - - it('supports paging via afterSeq', async () => { - const { publicEvents } = memoryPersistence().stores - await publicEvents!.append({ - runId: 'r1', - expectedSeq: 0, - event: text('a'), - }) - await publicEvents!.append({ - runId: 'r1', - expectedSeq: 1, - event: text('b'), - }) - const timeline = await loadRunHistory(publicEvents!, 'r1', { afterSeq: 1 }) - expect(timeline).toHaveLength(1) - }) - - it('returns an empty timeline for an unknown run', async () => { - const { publicEvents } = memoryPersistence().stores - expect(await loadRunHistory(publicEvents!, 'nope')).toEqual([]) - }) -}) diff --git a/packages/ai-persistence/tests/interrupts.test.ts b/packages/ai-persistence/tests/interrupts.test.ts index e9fa5d5dc..9480b656c 100644 --- a/packages/ai-persistence/tests/interrupts.test.ts +++ b/packages/ai-persistence/tests/interrupts.test.ts @@ -1,7 +1,6 @@ import { describe, expect, it, vi } from 'vitest' import { EventType, chat } from '@tanstack/ai' import type { AnyTextAdapter, StreamChunk, Tool } from '@tanstack/ai' -import { createApprovalController } from '../src/approval-controller' import { memoryPersistence } from '../src/memory' import { withChatPersistence } from '../src/middleware' import { defineAIPersistence, validatePersistenceFeatures } from '../src/types' @@ -100,19 +99,18 @@ const approvalClientTool = (name: string): Tool => ({ }) describe('interrupt persistence', () => { - it('requires run, public event, and interrupt stores', () => { + it('requires run and interrupt stores', () => { const stores = memoryPersistence().stores expect(() => validatePersistenceFeatures( defineAIPersistence({ stores: { - runs: stores.runs, - publicEvents: stores.publicEvents, + messages: stores.messages, }, }), ['interrupts'], ), - ).toThrow(/interrupts.*stores\.interrupts/i) + ).toThrow(/interrupts.*stores\.runs.*stores\.interrupts/i) }) it('persists RUN_FINISHED interrupt outcomes as pending interrupt records', async () => { @@ -137,8 +135,8 @@ describe('interrupt persistence', () => { expect((await persistence.stores.runs!.get('r1'))?.status).toBe( 'interrupted', ) - expect(chunks.every((chunk) => typeof chunk.cursor === 'string')).toBe(true) - expect(await persistence.stores.publicEvents!.latestSeq('r1')).toBe(2) + // Persistence is state-only: it never stamps delivery cursors on the stream. + expect(chunks.every((chunk) => !('cursor' in chunk))).toBe(true) }) it('saves thread messages when a messages-enabled run pauses on an interrupt', async () => { @@ -230,48 +228,10 @@ describe('interrupt persistence', () => { expect(await persistence.stores.runs!.get('r2')).toBeNull() }) - it('allows cursor replay on a thread with pending interrupts', async () => { + it('treats resume entries as interrupt continuation on the same run', async () => { const persistence = memoryPersistence() const first = mockAdapter([[runStarted(), interruptFinished()]]) - const original = await collect( - chat({ - adapter: first.adapter, - messages: [{ role: 'user', content: 'hi' }], - runId: 'r1', - threadId: 't1', - middleware: [ - withChatPersistence(persistence, { features: ['interrupts'] }), - ], - }) as AsyncIterable, - ) - expect(await persistence.stores.interrupts!.listPending('t1')).toHaveLength( - 1, - ) - - const replayAdapter = mockAdapter([[text('SHOULD NOT RUN')]]) - const replay = await collect( - chat({ - adapter: replayAdapter.adapter, - messages: [], - runId: 'r1', - threadId: 't1', - cursor: original[0]!.cursor, - middleware: [ - withChatPersistence(persistence, { features: ['interrupts'] }), - ], - }) as AsyncIterable, - ) - - expect(replayAdapter.calls).toHaveLength(0) - expect(replay.map((chunk) => chunk.cursor)).toEqual( - original.slice(1).map((chunk) => chunk.cursor), - ) - }) - - it('treats cursor plus resume entries as interrupt continuation, not replay', async () => { - const persistence = memoryPersistence() - const first = mockAdapter([[runStarted(), interruptFinished()]]) - const original = await collect( + await collect( chat({ adapter: first.adapter, messages: [{ role: 'user', content: 'hi' }], @@ -295,7 +255,6 @@ describe('interrupt persistence', () => { messages: [], runId: 'r1', threadId: 't1', - cursor: original[0]!.cursor, resume: [ { interruptId: 'interrupt-1', @@ -334,7 +293,7 @@ describe('interrupt persistence', () => { } as StreamChunk, ] const first = mockAdapter([toolCallChunks()]) - const original = await collect( + await collect( chat({ adapter: first.adapter, messages: [{ role: 'user', content: 'hi' }], @@ -360,7 +319,6 @@ describe('interrupt persistence', () => { tools: [approvalClientTool('clientSearch')], runId: 'r1', threadId: 't1', - cursor: original[0]!.cursor, resume: [ { interruptId: 'approval_tool-call-1', @@ -417,7 +375,6 @@ describe('interrupt persistence', () => { tools: [clientTool('clientSearch')], runId: 'r1', threadId: 't1', - cursor: approvalChunks[0]!.cursor, resume: [ { interruptId: 'client_tool_tool-call-1', @@ -445,10 +402,10 @@ describe('interrupt persistence', () => { expect(await persistence.stores.interrupts!.listPending('t1')).toEqual([]) }) - it('rejects invalid resume entries even when the cursor is valid', async () => { + it('rejects invalid resume entries against pending interrupts', async () => { const persistence = memoryPersistence() const first = mockAdapter([[runStarted(), interruptFinished()]]) - const original = await collect( + await collect( chat({ adapter: first.adapter, messages: [{ role: 'user', content: 'hi' }], @@ -468,7 +425,6 @@ describe('interrupt persistence', () => { messages: [], runId: 'r1', threadId: 't1', - cursor: original[0]!.cursor, resume: [], middleware: [ withChatPersistence(persistence, { features: ['interrupts'] }), @@ -484,7 +440,6 @@ describe('interrupt persistence', () => { messages: [], runId: 'r1', threadId: 't1', - cursor: original[0]!.cursor, resume: [{ interruptId: 'stale-interrupt', status: 'resolved' }], middleware: [ withChatPersistence(persistence, { features: ['interrupts'] }), @@ -499,10 +454,10 @@ describe('interrupt persistence', () => { ) }) - it('rejects stale resume entries when a valid cursor has no pending interrupts', async () => { + it('rejects stale resume entries when a thread has no pending interrupts', async () => { const persistence = memoryPersistence() const first = mockAdapter([[runStarted(), text('done'), runFinished()]]) - const original = await collect( + await collect( chat({ adapter: first.adapter, messages: [{ role: 'user', content: 'hi' }], @@ -523,7 +478,6 @@ describe('interrupt persistence', () => { messages: [], runId: 'r1', threadId: 't1', - cursor: original[0]!.cursor, resume: [{ interruptId: 'stale-interrupt', status: 'resolved' }], middleware: [ withChatPersistence(persistence, { features: ['interrupts'] }), @@ -535,76 +489,6 @@ describe('interrupt persistence', () => { expect(continuation.calls).toHaveLength(0) }) - it('blocks stale or unknown cursor input while a thread has pending interrupts', async () => { - const persistence = memoryPersistence() - await persistence.stores.interrupts!.create({ - interruptId: 'interrupt-1', - runId: 'old-run', - threadId: 't1', - status: 'pending', - requestedAt: 1, - payload: {}, - }) - const { adapter } = mockAdapter([[text('SHOULD NOT RUN')]]) - - await expect( - collect( - chat({ - adapter, - messages: [{ role: 'user', content: 'new input' }], - runId: 'unknown-run', - threadId: 't1', - cursor: Buffer.from('1:unknown-run', 'utf8').toString('base64url'), - middleware: [ - withChatPersistence(persistence, { features: ['interrupts'] }), - ], - }) as AsyncIterable, - ), - ).rejects.toThrow(/unknown run unknown-run/i) - }) - - it('blocks wrong-thread cursor input while a thread has pending interrupts', async () => { - const persistence = memoryPersistence() - const other = mockAdapter([ - [runStarted(), text('other'), runFinished('r-other')], - ]) - const otherChunks = await collect( - chat({ - adapter: other.adapter, - messages: [{ role: 'user', content: 'other' }], - runId: 'r-other', - threadId: 'other-thread', - middleware: [ - withChatPersistence(persistence, { features: ['interrupts'] }), - ], - }) as AsyncIterable, - ) - await persistence.stores.interrupts!.create({ - interruptId: 'interrupt-1', - runId: 'old-run', - threadId: 't1', - status: 'pending', - requestedAt: 1, - payload: {}, - }) - const { adapter } = mockAdapter([[text('SHOULD NOT RUN')]]) - - await expect( - collect( - chat({ - adapter, - messages: [{ role: 'user', content: 'new input' }], - runId: 'r-other', - threadId: 't1', - cursor: otherChunks[0]!.cursor, - middleware: [ - withChatPersistence(persistence, { features: ['interrupts'] }), - ], - }) as AsyncIterable, - ), - ).rejects.toThrow(/belongs to thread other-thread, not request thread t1/i) - }) - it('accepts resume only when every pending interrupt has a valid matching entry', async () => { const persistence = memoryPersistence() await persistence.stores.interrupts!.create({ @@ -757,7 +641,7 @@ describe('interrupt persistence', () => { expect(nextRun.calls).toHaveLength(1) }) - it('marks durable-replay terminal interrupt outcomes as interrupted', async () => { + it('marks terminal interrupt outcomes as interrupted', async () => { const persistence = memoryPersistence() const { adapter } = mockAdapter([[runStarted(), interruptFinished()]]) @@ -768,7 +652,7 @@ describe('interrupt persistence', () => { runId: 'r1', threadId: 't1', middleware: [ - withChatPersistence(persistence, { features: ['durable-replay'] }), + withChatPersistence(persistence, { features: ['interrupts'] }), ], }) as AsyncIterable, ) @@ -776,28 +660,8 @@ describe('interrupt persistence', () => { expect((await persistence.stores.runs!.get('r1'))?.status).toBe( 'interrupted', ) - expect(await persistence.stores.interrupts!.listPending('t1')).toEqual([]) - }) - - it('keeps deprecated approval controller behavior over interrupt storage', async () => { - const persistence = memoryPersistence() - const controller = createApprovalController({ - store: persistence.stores.interrupts!, - }) - const secondController = createApprovalController({ - store: persistence.stores.interrupts!, - }) - - await controller.request({ - approvalId: 'approval-1', - runId: 'r1', - threadId: 't1', - requestedAt: 1, - payload: { tool: 'search' }, - }) - await secondController.resolve('approval-1', true) - - const decisions = await secondController.decisionsForThread('t1') - expect(decisions.get('approval-1')).toBe(true) + expect(await persistence.stores.interrupts!.listPending('t1')).toHaveLength( + 1, + ) }) }) diff --git a/packages/ai-persistence/tests/memory.conformance.test.ts b/packages/ai-persistence/tests/memory.conformance.test.ts new file mode 100644 index 000000000..0b382b6a0 --- /dev/null +++ b/packages/ai-persistence/tests/memory.conformance.test.ts @@ -0,0 +1,4 @@ +import { runPersistenceConformance } from '../src/testkit/conformance' +import { memoryPersistence } from '../src/memory' + +runPersistenceConformance('memory', () => memoryPersistence()) diff --git a/packages/ai-persistence/tests/memory.test.ts b/packages/ai-persistence/tests/memory.test.ts index ed9b8da0a..5676d0f7c 100644 --- a/packages/ai-persistence/tests/memory.test.ts +++ b/packages/ai-persistence/tests/memory.test.ts @@ -1,35 +1,29 @@ import { describe, expect, it } from 'vitest' -import { EventType } from '@tanstack/ai' import { memoryPersistence } from '../src/memory' import { defineAIPersistence, defineChatPersistence, validatePersistenceFeatures, } from '../src/types' -import { AppendConflictError } from '../src/types' -import type { StreamChunk } from '@tanstack/ai' import type { AIPersistence, ChatPersistence } from '../src/types' -const chunk = (delta: string): StreamChunk => ({ - type: EventType.TEXT_MESSAGE_CONTENT, - messageId: 'm1', - delta, - timestamp: 1, -}) - describe('memoryPersistence', () => { - it('returns a namespaced AIPersistence with every reference store present', () => { + it('returns a namespaced AIPersistence with every state store present', () => { const p = memoryPersistence() expect(p.stores.messages).toBeDefined() expect(p.stores.runs).toBeDefined() - expect(p.stores.publicEvents).toBeDefined() - expect(p.stores.internalEvents).toBeDefined() expect(p.stores.interrupts).toBeDefined() expect(p.stores.artifacts).toBeDefined() expect(p.stores.blobs).toBeDefined() expect(p.stores.locks).toBeDefined() }) + it('exposes only state stores (no delivery event stores)', () => { + const stores = memoryPersistence().stores as Record + expect('publicEvents' in stores).toBe(false) + expect('internalEvents' in stores).toBe(false) + }) + it('defineAIPersistence is an identity helper', () => { const persistence = memoryPersistence() expect(defineAIPersistence(persistence)).toBe(persistence) @@ -47,9 +41,9 @@ describe('memoryPersistence', () => { defineAIPersistence({ stores: { messages: memoryPersistence().stores.messages }, }), - ['durable-replay'], + ['interrupts'], ), - ).toThrow(/durable-replay.*stores\.runs.*stores\.publicEvents/i) + ).toThrow(/interrupts.*stores\.runs.*stores\.interrupts/i) }) it('allows message-only persistence without event stores', () => { @@ -95,110 +89,6 @@ describe('memoryPersistence', () => { }) }) - describe('publicEvents', () => { - it('appends with CAS, reports hasRun/latestSeq, and reads after a seq', async () => { - const { publicEvents } = memoryPersistence().stores - expect(await publicEvents!.hasRun('r1')).toBe(false) - await publicEvents!.append({ - runId: 'r1', - expectedSeq: 0, - event: chunk('a'), - }) - await publicEvents!.append({ - runId: 'r1', - expectedSeq: 1, - event: chunk('b'), - }) - await publicEvents!.append({ - runId: 'r1', - expectedSeq: 2, - event: chunk('c'), - }) - expect(await publicEvents!.hasRun('r1')).toBe(true) - expect(await publicEvents!.latestSeq('r1')).toBe(3) - - const seen: Array = [] - for await (const e of publicEvents!.read('r1', { afterSeq: 1 })) { - seen.push(e.seq) - } - expect(seen).toEqual([2, 3]) - }) - - it('returns an existing identical event for idempotent CAS append', async () => { - const { publicEvents } = memoryPersistence().stores - const event = chunk('a') - const first = await publicEvents!.append({ - runId: 'r1', - expectedSeq: 0, - event, - }) - const second = await publicEvents!.append({ - runId: 'r1', - expectedSeq: 0, - event, - }) - expect(second).toEqual(first) - }) - - it('throws AppendConflictError when CAS append observes a different event', async () => { - const { publicEvents } = memoryPersistence().stores - await publicEvents!.append({ - runId: 'r1', - expectedSeq: 0, - event: chunk('a'), - }) - await expect( - publicEvents!.append({ - runId: 'r1', - expectedSeq: 0, - event: chunk('b'), - }), - ).rejects.toBeInstanceOf(AppendConflictError) - }) - - it('reads all events when no afterSeq is given', async () => { - const { publicEvents } = memoryPersistence().stores - await publicEvents!.append({ - runId: 'r1', - expectedSeq: 0, - event: chunk('a'), - }) - await publicEvents!.append({ - runId: 'r1', - expectedSeq: 1, - event: chunk('b'), - }) - const seen: Array = [] - for await (const e of publicEvents!.read('r1')) { - if (e.event.type === 'TEXT_MESSAGE_CONTENT') seen.push(e.event.delta) - } - expect(seen).toEqual(['a', 'b']) - }) - }) - - describe('internalEvents', () => { - it('uses the same CAS semantics for namespaced events', async () => { - const { internalEvents } = memoryPersistence().stores - const first = await internalEvents!.append({ - runId: 'r1', - expectedSeq: 0, - namespace: 'checkpoint', - type: 'saved', - payload: { ok: true }, - }) - expect(first.seq).toBe(1) - await expect( - internalEvents!.append({ - runId: 'r1', - expectedSeq: 0, - namespace: 'checkpoint', - type: 'saved', - payload: { ok: false }, - }), - ).rejects.toBeInstanceOf(AppendConflictError) - }) - }) - describe('messages', () => { it('round-trips a thread transcript', async () => { const { messages } = memoryPersistence().stores diff --git a/packages/ai-persistence/tests/state-only.test.ts b/packages/ai-persistence/tests/state-only.test.ts new file mode 100644 index 000000000..ecefc0917 --- /dev/null +++ b/packages/ai-persistence/tests/state-only.test.ts @@ -0,0 +1,105 @@ +import { describe, expect, it } from 'vitest' +import { EventType, chat } from '@tanstack/ai' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' +import { memoryPersistence } from '../src/memory' +import { withChatPersistence } from '../src/middleware' + +function mockAdapter(iterations: Array>) { + const calls: Array = [] + let i = 0 + const adapter = { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': {}, + chatStream: (opts: unknown) => { + calls.push(opts) + const chunks = iterations[i] ?? [] + i++ + return (async function* () { + for (const c of chunks) yield c + })() + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + } as unknown as AnyTextAdapter + return { adapter, calls } +} + +const script = (): Array> => [ + [ + { type: EventType.RUN_STARTED, runId: 'r1', threadId: 't1', timestamp: 1 }, + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm1', + delta: 'hello world', + timestamp: 1, + }, + { + type: EventType.RUN_FINISHED, + runId: 'r1', + threadId: 't1', + finishReason: 'stop', + timestamp: 1, + }, + ], +] + +async function collect(stream: AsyncIterable) { + const out: Array = [] + for await (const c of stream) out.push(c) + return out +} + +describe('state-only persistence', () => { + it('exposes only state stores (no delivery event stores)', () => { + const stores = memoryPersistence().stores as Record + expect('publicEvents' in stores).toBe(false) + expect('internalEvents' in stores).toBe(false) + // The delivery stream sink is gone too. + expect('stream' in memoryPersistence()).toBe(false) + }) + + it('persists thread messages and run status', async () => { + const persistence = memoryPersistence() + const { adapter } = mockAdapter(script()) + + await collect( + chat({ + adapter, + messages: [{ role: 'user', content: 'hi' }], + runId: 'r1', + threadId: 't1', + middleware: [withChatPersistence(persistence)], + }) as AsyncIterable, + ) + + expect((await persistence.stores.runs!.get('r1'))?.status).toBe('completed') + expect( + (await persistence.stores.messages!.loadThread('t1')).length, + ).toBeGreaterThan(0) + }) + + it('produces chunks byte-identical to a non-persisted run (no cursor)', async () => { + const persisted = await collect( + chat({ + adapter: mockAdapter(script()).adapter, + messages: [{ role: 'user', content: 'hi' }], + runId: 'r1', + threadId: 't1', + middleware: [withChatPersistence(memoryPersistence())], + }) as AsyncIterable, + ) + + const plain = await collect( + chat({ + adapter: mockAdapter(script()).adapter, + messages: [{ role: 'user', content: 'hi' }], + runId: 'r1', + threadId: 't1', + }) as AsyncIterable, + ) + + expect(JSON.stringify(persisted)).toEqual(JSON.stringify(plain)) + expect(persisted.every((c) => !('cursor' in c))).toBe(true) + }) +}) diff --git a/packages/ai-persistence/tests/with-persistence.test.ts b/packages/ai-persistence/tests/with-persistence.test.ts index 591f25dcc..070368765 100644 --- a/packages/ai-persistence/tests/with-persistence.test.ts +++ b/packages/ai-persistence/tests/with-persistence.test.ts @@ -1,12 +1,9 @@ import { describe, expect, it } from 'vitest' import { EventType, chat } from '@tanstack/ai' import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' -import { encodeCursor } from '../src/cursor' import { memoryPersistence } from '../src/memory' import { withChatPersistence } from '../src/middleware' -import { createResumeSource } from '../src/resume-source' import { defineAIPersistence } from '../src/types' -import type { PersistedPublicEvent, PublicEventStore } from '../src/types' // --- minimal mock text adapter --------------------------------------------- @@ -84,31 +81,8 @@ async function expectCollectRejects( await expect(collect(stream)).rejects.toThrow(pattern) } -function sparsePublicEvents( - runId: string, - event: StreamChunk, -): PublicEventStore { - const persisted: PersistedPublicEvent = { - seq: 2, - cursor: encodeCursor(runId, 2), - event: { ...event, cursor: encodeCursor(runId, 2) }, - } - return { - append: async () => persisted, - read: (requestedRunId, opts) => - (async function* () { - if (requestedRunId !== runId) return - if ((opts?.afterSeq ?? -Infinity) < persisted.seq) { - yield persisted - } - })(), - hasRun: async (requestedRunId) => requestedRunId === runId, - latestSeq: async (requestedRunId) => (requestedRunId === runId ? 2 : 0), - } -} - -describe('withChatPersistence (no sandbox)', () => { - it('persists events with cursors, completes the run, and saves the transcript', async () => { +describe('withChatPersistence (state-only)', () => { + it('completes the run and saves the transcript', async () => { const persistence = memoryPersistence() const { adapter } = mockAdapter([ [ev.runStarted(), ev.text('hello'), ev.runFinished()], @@ -124,149 +98,24 @@ describe('withChatPersistence (no sandbox)', () => { }) as AsyncIterable, ) - // Every emitted chunk carries an in-band cursor. + // The persistence middleware never stamps delivery cursors on the stream. expect(chunks.length).toBeGreaterThan(0) - expect(chunks.every((c) => typeof c.cursor === 'string')).toBe(true) + expect(chunks.every((c) => !('cursor' in c))).toBe(true) - // Events are in the log, run is completed, transcript saved. - expect(await persistence.stores.publicEvents!.hasRun('r1')).toBe(true) + // Run is completed and the transcript is saved. expect((await persistence.stores.runs!.get('r1'))?.status).toBe('completed') expect( (await persistence.stores.messages!.loadThread('t1')).length, ).toBeGreaterThan(0) }) - it('resumes by replaying events after the cursor without running the adapter', async () => { - const persistence = memoryPersistence() - const first = mockAdapter([ - [ev.runStarted(), ev.text('hello'), ev.text(' world'), ev.runFinished()], - ]) - - const original = await collect( - chat({ - adapter: first.adapter, - messages: [{ role: 'user', content: 'hi' }], - runId: 'r1', - threadId: 't1', - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - ) - - // Resume from the first chunk's cursor — should replay the rest, no adapter. - const afterCursor = original[0]!.cursor! - const resumeAdapter = mockAdapter([[ev.text('SHOULD NOT RUN')]]) - const replay = await collect( - chat({ - adapter: resumeAdapter.adapter, - messages: [], - runId: 'r1', - threadId: 't1', - cursor: afterCursor, - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - ) - - expect(resumeAdapter.calls.length).toBe(0) - // Replayed cursors are exactly the original tail after the resume point. - expect(replay.map((c) => c.cursor)).toEqual( - original.slice(1).map((c) => c.cursor), - ) - }) - - it('fails loudly when the cursor run id does not match the request run id', async () => { + it('records an interrupt and marks the run interrupted', async () => { const persistence = memoryPersistence() - const first = mockAdapter([[ev.runStarted(), ev.text('hello')]]) - - const original = await collect( - chat({ - adapter: first.adapter, - messages: [{ role: 'user', content: 'hi' }], - runId: 'r1', - threadId: 't1', - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - ) - - const resumeAdapter = mockAdapter([[ev.text('SHOULD NOT RUN')]]) - await expectCollectRejects( - chat({ - adapter: resumeAdapter.adapter, - messages: [], - runId: 'r2', - threadId: 't1', - cursor: original[0]!.cursor!, - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - /cursor runId r1 does not match request runId r2/i, - ) - expect(resumeAdapter.calls.length).toBe(0) - }) - - it('fails loudly when the persisted run thread does not match the request thread', async () => { - const persistence = memoryPersistence() - const first = mockAdapter([[ev.runStarted(), ev.text('hello')]]) - - const original = await collect( - chat({ - adapter: first.adapter, - messages: [{ role: 'user', content: 'hi' }], - runId: 'r1', - threadId: 't1', - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - ) - - const resumeAdapter = mockAdapter([[ev.text('SHOULD NOT RUN')]]) - await expectCollectRejects( - chat({ - adapter: resumeAdapter.adapter, - messages: [], - runId: 'r1', - threadId: 't2', - cursor: original[0]!.cursor!, - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - /belongs to thread t1, not request thread t2/i, - ) - expect(resumeAdapter.calls.length).toBe(0) - }) - - it('fails loudly when the cursor points beyond persisted public events', async () => { - const persistence = memoryPersistence() - const first = mockAdapter([[ev.runStarted(), ev.text('hello')]]) + const { adapter } = mockAdapter([[ev.runStarted(), ev.interrupted()]]) await collect( chat({ - adapter: first.adapter, - messages: [{ role: 'user', content: 'hi' }], - runId: 'r1', - threadId: 't1', - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - ) - - const resumeAdapter = mockAdapter([[ev.text('SHOULD NOT RUN')]]) - await expectCollectRejects( - chat({ - adapter: resumeAdapter.adapter, - messages: [], - runId: 'r1', - threadId: 't1', - cursor: encodeCursor('r1', 999), - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - /beyond latest persisted sequence/i, - ) - expect(resumeAdapter.calls.length).toBe(0) - }) - - it('fails loudly when the cursor sequence is zero', async () => { - const persistence = memoryPersistence() - const first = mockAdapter([[ev.runStarted(), ev.text('hello')]]) - - await collect( - chat({ - adapter: first.adapter, + adapter, messages: [{ role: 'user', content: 'hi' }], runId: 'r1', threadId: 't1', @@ -274,68 +123,12 @@ describe('withChatPersistence (no sandbox)', () => { }) as AsyncIterable, ) - const resumeAdapter = mockAdapter([[ev.text('SHOULD NOT RUN')]]) - await expectCollectRejects( - chat({ - adapter: resumeAdapter.adapter, - messages: [], - runId: 'r1', - threadId: 't1', - cursor: encodeCursor('r1', 0), - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - /sequence 0 is invalid/i, - ) - expect(resumeAdapter.calls.length).toBe(0) - }) - - it('fails loudly when the cursor sequence is missing from a sparse public event log', async () => { - const full = memoryPersistence() - await full.stores.runs!.createOrResume({ - runId: 'r1', - threadId: 't1', - startedAt: 1, - }) - const persistence = defineAIPersistence({ - stores: { - runs: full.stores.runs, - publicEvents: sparsePublicEvents('r1', ev.text('tail')), - }, - }) - const resumeAdapter = mockAdapter([[ev.text('SHOULD NOT RUN')]]) - - await expectCollectRejects( - chat({ - adapter: resumeAdapter.adapter, - messages: [], - runId: 'r1', - threadId: 't1', - cursor: encodeCursor('r1', 1), - middleware: [ - withChatPersistence(persistence, { features: ['durable-replay'] }), - ], - }) as AsyncIterable, - /sequence 1 does not reference a persisted public event/i, + expect((await persistence.stores.runs!.get('r1'))?.status).toBe( + 'interrupted', ) - expect(resumeAdapter.calls.length).toBe(0) - }) - - it('fails loudly when the cursor references an unknown run', async () => { - const persistence = memoryPersistence() - const resumeAdapter = mockAdapter([[ev.text('SHOULD NOT RUN')]]) - - await expectCollectRejects( - chat({ - adapter: resumeAdapter.adapter, - messages: [], - runId: 'missing-run', - threadId: 't1', - cursor: encodeCursor('missing-run', 1), - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - /unknown run missing-run/i, + expect(await persistence.stores.interrupts!.listPending('t1')).toHaveLength( + 1, ) - expect(resumeAdapter.calls.length).toBe(0) }) it('blocks normal new input while a thread has pending interrupts', async () => { @@ -366,48 +159,6 @@ describe('withChatPersistence (no sandbox)', () => { expect(next.calls.length).toBe(0) }) - it('does not let a forged in-range cursor bypass pending interrupts', async () => { - const full = memoryPersistence() - await full.stores.runs!.createOrResume({ - runId: 'r1', - threadId: 't1', - startedAt: 1, - }) - await full.stores.interrupts!.create({ - interruptId: 'interrupt-1', - runId: 'r1', - threadId: 't1', - status: 'pending', - requestedAt: 1, - payload: {}, - }) - const persistence = defineAIPersistence({ - stores: { - runs: full.stores.runs, - publicEvents: sparsePublicEvents('r1', ev.text('tail')), - interrupts: full.stores.interrupts, - }, - }) - - const resumeAdapter = mockAdapter([[ev.text('SHOULD NOT RUN')]]) - await expectCollectRejects( - chat({ - adapter: resumeAdapter.adapter, - messages: [], - runId: 'r1', - threadId: 't1', - cursor: encodeCursor('r1', 1), - middleware: [ - withChatPersistence(persistence, { features: ['interrupts'] }), - ], - }) as AsyncIterable, - /sequence 1 does not reference a persisted public event/i, - ) - - expect(resumeAdapter.calls.length).toBe(0) - expect(await full.stores.interrupts!.listPending('t1')).toHaveLength(1) - }) - it('requires resume entries to match all pending interrupts', async () => { const persistence = memoryPersistence() const first = mockAdapter([[ev.runStarted(), ev.interrupted()]]) @@ -478,53 +229,7 @@ describe('withChatPersistence (no sandbox)', () => { ).toBe('resolved') }) - it('does not block new input for internal/background waits without public interrupts', async () => { - const persistence = memoryPersistence() - await persistence.stores.internalEvents!.append({ - runId: 'internal-run', - expectedSeq: 0, - namespace: 'background', - type: 'wait', - payload: { reason: 'background' }, - }) - - const next = mockAdapter([[ev.runStarted('r2', 't1'), ev.text('fresh')]]) - const chunks = await collect( - chat({ - adapter: next.adapter, - messages: [{ role: 'user', content: 'new input' }], - runId: 'r2', - threadId: 't1', - middleware: [withChatPersistence(persistence)], - }) as AsyncIterable, - ) - - expect(next.calls.length).toBe(1) - expect( - chunks.some((chunk) => chunk.type === EventType.TEXT_MESSAGE_CONTENT), - ).toBe(true) - }) - - it('requires a run record for persistence-backed resume source replay when a run store is provided', async () => { - const persistence = memoryPersistence() - await persistence.stores.publicEvents!.append({ - runId: 'ghost-run', - expectedSeq: 0, - event: { ...ev.text('ghost'), cursor: encodeCursor('ghost-run', 1) }, - }) - - const source = createResumeSource( - persistence.stores.publicEvents!, - persistence.stores.runs, - ) - - await expect(source.hasRun('ghost-run')).resolves.toBe(false) - await expect( - collect(source.replay('ghost-run', encodeCursor('ghost-run', 1))), - ).rejects.toThrow(/unknown run ghost-run/i) - }) - - it('persists messages without requiring event stores', async () => { + it('persists messages without requiring a run store', async () => { const full = memoryPersistence() const persistence = defineAIPersistence({ stores: { messages: full.stores.messages }, @@ -548,46 +253,17 @@ describe('withChatPersistence (no sandbox)', () => { expect(await persistence.stores.messages!.loadThread('t1')).not.toEqual([]) }) - it('treats cursor as a no-op with messages-only persistence', async () => { - const full = memoryPersistence() - const persistence = defineAIPersistence({ - stores: { messages: full.stores.messages }, - }) - const { adapter, calls } = mockAdapter([ - [ev.runStarted(), ev.text('hello'), ev.runFinished()], - ]) - - const chunks = await collect( - chat({ - adapter, - messages: [{ role: 'user', content: 'hi' }], - runId: 'r1', - threadId: 't1', - cursor: 'not-a-persistence-cursor', - middleware: [ - withChatPersistence(persistence, { features: ['messages'] }), - ], - }) as AsyncIterable, - ) - - expect(calls.length).toBe(1) - expect(chunks.every((chunk) => chunk.cursor === undefined)).toBe(true) - expect(await persistence.stores.messages!.loadThread('t1')).not.toEqual([]) - }) - - it('throws a descriptive error when durable replay is requested without event stores', async () => { + it('fails loudly when interrupts are requested without a run store', () => { const persistence = defineAIPersistence({ stores: { messages: memoryPersistence().stores.messages }, }) - const { adapter } = mockAdapter([[ev.text('unused')]]) expect(() => - withChatPersistence(persistence, { features: ['durable-replay'] }), - ).toThrow(/durable-replay.*stores\.runs.*stores\.publicEvents/i) - void adapter + withChatPersistence(persistence, { features: ['interrupts'] }), + ).toThrow(/interrupts.*stores\.runs.*stores\.interrupts/i) }) - it('is a no-op without the middleware: chunks carry no cursor', async () => { + it('is a no-op without the middleware: the stream is unchanged', async () => { const { adapter } = mockAdapter([ [ev.runStarted(), ev.text('plain'), ev.runFinished()], ]) @@ -599,6 +275,6 @@ describe('withChatPersistence (no sandbox)', () => { threadId: 't1', }) as AsyncIterable, ) - expect(chunks.every((c) => c.cursor === undefined)).toBe(true) + expect(chunks.every((c) => !('cursor' in c))).toBe(true) }) }) diff --git a/packages/ai-persistence/vite.config.ts b/packages/ai-persistence/vite.config.ts index 11f5b20b7..95f0f3d72 100644 --- a/packages/ai-persistence/vite.config.ts +++ b/packages/ai-persistence/vite.config.ts @@ -30,8 +30,12 @@ const config = defineConfig({ export default mergeConfig( config, tanstackViteConfig({ - entry: ['./src/index.ts'], + entry: ['./src/index.ts', './src/testkit/conformance.ts'], srcDir: './src', + // The conformance testkit imports Vitest; keep it external so the built + // artifact references the consumer's Vitest at test time instead of + // bundling the runner. + externalDeps: ['vitest'], cjs: false, }), ) From a2d91a709ed4c9f3d1eed99dcdbf16e9b16e6909 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Thu, 9 Jul 2026 18:51:49 +0200 Subject: [PATCH 4/9] feat(persistence)!: drizzle + prisma state backends; remove SQL-driver sprawl --- .../ai-persistence-cloudflare/package.json | 52 - .../ai-persistence-cloudflare/src/index.ts | 908 ------------ .../tests/d1.test.ts | 1301 ----------------- .../ai-persistence-cloudflare/tsconfig.json | 8 - .../ai-persistence-cloudflare/vite.config.ts | 37 - .../ai-persistence-drizzle/drizzle.config.ts | 14 + .../drizzle/0000_cute_vision.sql | 54 + .../drizzle/meta/0000_snapshot.json | 338 +++++ .../drizzle/meta/_journal.json | 13 + packages/ai-persistence-drizzle/package.json | 22 +- packages/ai-persistence-drizzle/src/bin.ts | 17 - .../ai-persistence-drizzle/src/cli-core.ts | 31 - packages/ai-persistence-drizzle/src/cli.ts | 9 - packages/ai-persistence-drizzle/src/index.ts | 214 +-- packages/ai-persistence-drizzle/src/schema.ts | 102 ++ packages/ai-persistence-drizzle/src/stores.ts | 498 +++++++ .../ai-persistence-drizzle/tests/cli.test.ts | 97 -- .../tests/drizzle.conformance.test.ts | 6 + .../tests/drizzle.test.ts | 37 - .../ai-persistence-drizzle/tsup.bin.config.ts | 11 - packages/ai-persistence-postgres/package.json | 61 - packages/ai-persistence-postgres/src/index.ts | 134 -- .../tests/driver.test.ts | 79 - .../ai-persistence-postgres/tsconfig.json | 8 - .../ai-persistence-postgres/vite.config.ts | 37 - packages/ai-persistence-prisma/.gitignore | 3 + packages/ai-persistence-prisma/package.json | 42 +- .../20260709120042_init/migration.sql | 61 + .../prisma/migrations/migration_lock.toml | 3 + .../prisma/schema.prisma | 101 ++ packages/ai-persistence-prisma/src/bin.ts | 17 - .../ai-persistence-prisma/src/cli-core.ts | 31 - packages/ai-persistence-prisma/src/cli.ts | 9 - packages/ai-persistence-prisma/src/index.ts | 100 +- packages/ai-persistence-prisma/src/stores.ts | 492 +++++++ .../ai-persistence-prisma/tests/cli.test.ts | 91 -- .../tests/prisma.conformance.test.ts | 59 + .../tests/prisma.test.ts | 55 - .../ai-persistence-prisma/tsup.bin.config.ts | 11 - packages/ai-persistence-sql/package.json | 64 - packages/ai-persistence-sql/src/bin.ts | 17 - packages/ai-persistence-sql/src/cli-core.ts | 250 ---- packages/ai-persistence-sql/src/cli.ts | 9 - packages/ai-persistence-sql/src/driver.ts | 115 -- packages/ai-persistence-sql/src/index.ts | 43 - packages/ai-persistence-sql/src/migrations.ts | 125 -- .../ai-persistence-sql/src/sql-persistence.ts | 106 -- packages/ai-persistence-sql/src/stores.ts | 700 --------- packages/ai-persistence-sql/tests/cli.test.ts | 129 -- .../tests/sql-persistence.test.ts | 765 ---------- .../ai-persistence-sql/tests/sqlite-driver.ts | 32 - packages/ai-persistence-sql/tsconfig.json | 8 - .../ai-persistence-sql/tsup.bin.config.ts | 12 - packages/ai-persistence-sql/vite.config.ts | 37 - packages/ai-persistence-sqlite/package.json | 54 - packages/ai-persistence-sqlite/src/index.ts | 95 -- .../tests/sqlite.test.ts | 78 - packages/ai-persistence-sqlite/tsconfig.json | 8 - packages/ai-persistence-sqlite/vite.config.ts | 37 - 59 files changed, 1951 insertions(+), 5896 deletions(-) delete mode 100644 packages/ai-persistence-cloudflare/package.json delete mode 100644 packages/ai-persistence-cloudflare/src/index.ts delete mode 100644 packages/ai-persistence-cloudflare/tests/d1.test.ts delete mode 100644 packages/ai-persistence-cloudflare/tsconfig.json delete mode 100644 packages/ai-persistence-cloudflare/vite.config.ts create mode 100644 packages/ai-persistence-drizzle/drizzle.config.ts create mode 100644 packages/ai-persistence-drizzle/drizzle/0000_cute_vision.sql create mode 100644 packages/ai-persistence-drizzle/drizzle/meta/0000_snapshot.json create mode 100644 packages/ai-persistence-drizzle/drizzle/meta/_journal.json delete mode 100644 packages/ai-persistence-drizzle/src/bin.ts delete mode 100644 packages/ai-persistence-drizzle/src/cli-core.ts delete mode 100644 packages/ai-persistence-drizzle/src/cli.ts create mode 100644 packages/ai-persistence-drizzle/src/schema.ts create mode 100644 packages/ai-persistence-drizzle/src/stores.ts delete mode 100644 packages/ai-persistence-drizzle/tests/cli.test.ts create mode 100644 packages/ai-persistence-drizzle/tests/drizzle.conformance.test.ts delete mode 100644 packages/ai-persistence-drizzle/tests/drizzle.test.ts delete mode 100644 packages/ai-persistence-drizzle/tsup.bin.config.ts delete mode 100644 packages/ai-persistence-postgres/package.json delete mode 100644 packages/ai-persistence-postgres/src/index.ts delete mode 100644 packages/ai-persistence-postgres/tests/driver.test.ts delete mode 100644 packages/ai-persistence-postgres/tsconfig.json delete mode 100644 packages/ai-persistence-postgres/vite.config.ts create mode 100644 packages/ai-persistence-prisma/.gitignore create mode 100644 packages/ai-persistence-prisma/prisma/migrations/20260709120042_init/migration.sql create mode 100644 packages/ai-persistence-prisma/prisma/migrations/migration_lock.toml create mode 100644 packages/ai-persistence-prisma/prisma/schema.prisma delete mode 100644 packages/ai-persistence-prisma/src/bin.ts delete mode 100644 packages/ai-persistence-prisma/src/cli-core.ts delete mode 100644 packages/ai-persistence-prisma/src/cli.ts create mode 100644 packages/ai-persistence-prisma/src/stores.ts delete mode 100644 packages/ai-persistence-prisma/tests/cli.test.ts create mode 100644 packages/ai-persistence-prisma/tests/prisma.conformance.test.ts delete mode 100644 packages/ai-persistence-prisma/tests/prisma.test.ts delete mode 100644 packages/ai-persistence-prisma/tsup.bin.config.ts delete mode 100644 packages/ai-persistence-sql/package.json delete mode 100644 packages/ai-persistence-sql/src/bin.ts delete mode 100644 packages/ai-persistence-sql/src/cli-core.ts delete mode 100644 packages/ai-persistence-sql/src/cli.ts delete mode 100644 packages/ai-persistence-sql/src/driver.ts delete mode 100644 packages/ai-persistence-sql/src/index.ts delete mode 100644 packages/ai-persistence-sql/src/migrations.ts delete mode 100644 packages/ai-persistence-sql/src/sql-persistence.ts delete mode 100644 packages/ai-persistence-sql/src/stores.ts delete mode 100644 packages/ai-persistence-sql/tests/cli.test.ts delete mode 100644 packages/ai-persistence-sql/tests/sql-persistence.test.ts delete mode 100644 packages/ai-persistence-sql/tests/sqlite-driver.ts delete mode 100644 packages/ai-persistence-sql/tsconfig.json delete mode 100644 packages/ai-persistence-sql/tsup.bin.config.ts delete mode 100644 packages/ai-persistence-sql/vite.config.ts delete mode 100644 packages/ai-persistence-sqlite/package.json delete mode 100644 packages/ai-persistence-sqlite/src/index.ts delete mode 100644 packages/ai-persistence-sqlite/tests/sqlite.test.ts delete mode 100644 packages/ai-persistence-sqlite/tsconfig.json delete mode 100644 packages/ai-persistence-sqlite/vite.config.ts diff --git a/packages/ai-persistence-cloudflare/package.json b/packages/ai-persistence-cloudflare/package.json deleted file mode 100644 index afbab2574..000000000 --- a/packages/ai-persistence-cloudflare/package.json +++ /dev/null @@ -1,52 +0,0 @@ -{ - "name": "@tanstack/ai-persistence-cloudflare", - "version": "0.1.0", - "description": "Cloudflare backend for TanStack AI persistence — D1-backed SQL stores (compile-verified), with optional Durable Object locks.", - "author": "", - "license": "MIT", - "repository": { - "type": "git", - "url": "git+https://github.com/TanStack/ai.git", - "directory": "packages/ai-persistence-cloudflare" - }, - "keywords": [ - "ai", - "tanstack", - "persistence" - ], - "type": "module", - "module": "./dist/esm/index.js", - "types": "./dist/esm/index.d.ts", - "exports": { - ".": { - "types": "./dist/esm/index.d.ts", - "import": "./dist/esm/index.js" - } - }, - "files": [ - "dist", - "src" - ], - "scripts": { - "build": "vite build", - "clean": "premove ./build ./dist", - "lint:fix": "eslint ./src --fix", - "test:build": "publint --strict", - "test:eslint": "eslint ./src", - "test:lib": "vitest", - "test:lib:dev": "pnpm test:lib --watch", - "test:types": "tsc" - }, - "peerDependencies": { - "@tanstack/ai": "workspace:^", - "@tanstack/ai-persistence": "workspace:^", - "@tanstack/ai-persistence-sql": "workspace:^" - }, - "devDependencies": { - "@cloudflare/workers-types": "^4.20260317.1", - "@tanstack/ai": "workspace:*", - "@tanstack/ai-persistence": "workspace:*", - "@tanstack/ai-persistence-sql": "workspace:*", - "@vitest/coverage-v8": "4.0.14" - } -} diff --git a/packages/ai-persistence-cloudflare/src/index.ts b/packages/ai-persistence-cloudflare/src/index.ts deleted file mode 100644 index 1a39db9df..000000000 --- a/packages/ai-persistence-cloudflare/src/index.ts +++ /dev/null @@ -1,908 +0,0 @@ -/// -/** - * Cloudflare backend. D1 (SQLite-compatible) backs the shared SQL stores; an - * optional Durable Object namespace can back the distributed lock. - * - * COMPILE-VERIFIED ONLY: this package type-checks against real `@cloudflare` - * types but is not runtime-verified here (it needs a Workers runtime). The D1 - * driver itself is unit-tested against a fake D1 so the adapter logic is - * covered. - */ -import { createSqlPersistence } from '@tanstack/ai-persistence-sql' -import type { - SqlDriver, - SqlPersistence, - SqlRow, -} from '@tanstack/ai-persistence-sql' -import type { - ArtifactRecord, - ArtifactStore, - BlobObject, - BlobRecord, - BlobStore, - PersistenceMode, -} from '@tanstack/ai-persistence' -import type { LockStore } from '@tanstack/ai' - -const DEFAULT_R2_BLOB_PREFIX = 'tanstack-ai/blobs/' -const DEFAULT_R2_ARTIFACT_PREFIX = 'tanstack-ai/artifacts/' -const DEFAULT_LOCK_LEASE_MS = 30_000 -const DEFAULT_LOCK_POLL_MS = 50 -const MIN_LOCK_TIMING_MS = 1 - -export interface R2ArtifactStoreOptions { - /** Prefix for all artifact metadata, run indexes, and blob objects. */ - prefix?: string -} - -export interface R2BlobStoreOptions { - /** Prefix for all blob objects. Logical keys passed to the store do not include it. */ - prefix?: string -} - -export interface CloudflareArtifactStoreOptions { - /** Logical BlobStore key prefix for artifact bytes. */ - blobKeyPrefix?: string - now?: () => number - /** Run artifact table migrations on first use (default false). */ - migrate?: boolean -} - -export interface DurableObjectLockStoreOptions { - leaseMs?: number - pollMs?: number -} - -type R2ArtifactMetadata = Omit & { - blobKey?: string -} - -function normalizeR2Prefix(prefix?: string): string { - const normalized = (prefix ?? DEFAULT_R2_ARTIFACT_PREFIX).replace(/^\/+/, '') - return normalized.endsWith('/') ? normalized : `${normalized}/` -} - -function normalizeBlobPrefix(prefix?: string): string { - const normalized = (prefix ?? DEFAULT_R2_BLOB_PREFIX).replace(/^\/+/, '') - return normalized.endsWith('/') ? normalized : `${normalized}/` -} - -function keyPart(value: string): string { - return encodeURIComponent(value) -} - -function orderedKeyPart(value: string): string { - return `${keyPart(value)}/${Date.now()}-${crypto.randomUUID()}` -} - -function toPhysicalKey(prefix: string, key: string): string { - return `${prefix}${key.replace(/^\/+/, '')}` -} - -function toLogicalKey(prefix: string, key: string): string { - return key.startsWith(prefix) ? key.slice(prefix.length) : key -} - -function r2RecordFromObject( - prefix: string, - object: { - key: string - size?: number - etag?: string - uploaded?: Date - httpMetadata?: { contentType?: string } - customMetadata?: Record - }, -): BlobRecord { - return { - key: toLogicalKey(prefix, object.key), - size: object.size, - etag: object.etag, - contentType: object.httpMetadata?.contentType, - customMetadata: object.customMetadata, - createdAt: object.uploaded?.getTime(), - updatedAt: object.uploaded?.getTime(), - } -} - -/** Build an R2-backed generic blob store. */ -export function createR2BlobStore( - bucket: R2Bucket, - options?: R2BlobStoreOptions, -): BlobStore { - const prefix = normalizeBlobPrefix(options?.prefix) - - return { - async put(key, body, putOptions) { - const result = await bucket.put(toPhysicalKey(prefix, key), body, { - httpMetadata: putOptions?.contentType - ? { contentType: putOptions.contentType } - : undefined, - customMetadata: putOptions?.customMetadata, - }) - return r2RecordFromObject(prefix, { - key: toPhysicalKey(prefix, key), - size: result.size, - etag: result.etag, - uploaded: result.uploaded, - httpMetadata: result.httpMetadata, - customMetadata: result.customMetadata, - }) - }, - - async get(key) { - const object = await bucket.get(toPhysicalKey(prefix, key)) - if (!object) return null - const record = r2RecordFromObject(prefix, object) - return { - ...record, - body: object.body, - arrayBuffer: () => object.arrayBuffer(), - text: () => object.text(), - } satisfies BlobObject - }, - - async head(key) { - const object = await bucket.head(toPhysicalKey(prefix, key)) - return object ? r2RecordFromObject(prefix, object) : null - }, - - async delete(key) { - await bucket.delete(toPhysicalKey(prefix, key)) - }, - - async list(listOptions) { - const result = await bucket.list({ - prefix: toPhysicalKey(prefix, listOptions?.prefix ?? ''), - cursor: listOptions?.cursor, - limit: listOptions?.limit, - }) - return { - objects: result.objects.map((object) => - r2RecordFromObject(prefix, object), - ), - cursor: 'cursor' in result ? result.cursor : undefined, - truncated: result.truncated, - } - }, - } -} - -function metadataFor( - record: ArtifactRecord, - blobKey?: string, -): R2ArtifactMetadata { - const { bytes: _bytes, ...metadata } = record - return blobKey ? { ...metadata, blobKey } : metadata -} - -async function readMetadata( - bucket: R2Bucket, - key: string, -): Promise { - const object = await bucket.get(key) - if (!object) return null - return JSON.parse(await object.text()) as R2ArtifactMetadata -} - -function recordFromMetadata(metadata: R2ArtifactMetadata): ArtifactRecord { - const { blobKey: _blobKey, ...record } = metadata - return record -} - -function withBytes( - metadata: R2ArtifactMetadata, - bytes: Uint8Array, -): ArtifactRecord { - return { ...recordFromMetadata(metadata), bytes } -} - -/** Build an R2-backed artifact store with metadata indexes separated from blob bytes. */ -export function createR2ArtifactStore( - bucket: R2Bucket, - options?: R2ArtifactStoreOptions, -): ArtifactStore { - const prefix = normalizeR2Prefix(options?.prefix) - const idKey = (artifactId: string) => - `${prefix}by-id/${keyPart(artifactId)}/metadata.json` - const runKey = (runId: string, artifactId: string) => - `${prefix}by-run/${keyPart(runId)}/${keyPart(artifactId)}.json` - const runPrefix = (runId: string) => `${prefix}by-run/${keyPart(runId)}/` - const blobKey = (artifactId: string) => - `${prefix}blobs/${keyPart(artifactId)}/${Date.now()}-${crypto.randomUUID()}` - const listRunIndexEntries = async (runId: string) => { - const entries: Array<{ key: string; metadata: R2ArtifactMetadata }> = [] - let cursor: string | undefined - do { - const result = await bucket.list({ - prefix: runPrefix(runId), - cursor, - }) - for (const object of result.objects) { - const metadata = await readMetadata(bucket, object.key) - if (metadata) entries.push({ key: object.key, metadata }) - } - cursor = result.truncated ? result.cursor : undefined - } while (cursor) - return entries - } - const listArtifacts = async (runId: string) => { - const records: Array = [] - const entries = await listRunIndexEntries(runId) - for (const entry of entries) { - const metadata = await readMetadata( - bucket, - idKey(entry.metadata.artifactId), - ) - if ( - metadata?.artifactId === entry.metadata.artifactId && - metadata.runId === runId - ) { - records.push(recordFromMetadata(metadata)) - } - } - return records - } - - return { - async save(record) { - const existing = await readMetadata(bucket, idKey(record.artifactId)) - const nextBlobKey = record.bytes ? blobKey(record.artifactId) : undefined - const nextMetadata = metadataFor(record, nextBlobKey) - const encodedMetadata = JSON.stringify(nextMetadata) - - if (record.bytes && nextBlobKey) { - await bucket.put(nextBlobKey, record.bytes) - } - - try { - await bucket.put( - runKey(record.runId, record.artifactId), - encodedMetadata, - ) - } catch (error) { - if (nextBlobKey) { - try { - await bucket.delete(nextBlobKey) - } catch { - // Preserve the commit failure; cleanup is best-effort. - } - } - throw error - } - - try { - await bucket.put(idKey(record.artifactId), encodedMetadata) - } catch (error) { - if (existing?.runId === record.runId) { - try { - await bucket.put( - runKey(existing.runId, existing.artifactId), - JSON.stringify(existing), - ) - } catch { - // Preserve the commit failure; cleanup is best-effort. - } - } else { - try { - await bucket.delete(runKey(record.runId, record.artifactId)) - } catch { - // Preserve the commit failure; cleanup is best-effort. - } - } - if (nextBlobKey) { - try { - await bucket.delete(nextBlobKey) - } catch { - // Preserve the commit failure; cleanup is best-effort. - } - } - throw error - } - - if (existing && existing.runId !== record.runId) { - try { - await bucket.delete(runKey(existing.runId, existing.artifactId)) - } catch { - // New metadata is committed; stale run index cleanup is best-effort. - } - } - if (existing?.blobKey && existing.blobKey !== nextBlobKey) { - try { - await bucket.delete(existing.blobKey) - } catch { - // New metadata is committed; stale blob cleanup is best-effort. - } - } - }, - - async get(artifactId) { - const metadata = await readMetadata(bucket, idKey(artifactId)) - if (!metadata) return null - if (!metadata.blobKey) return recordFromMetadata(metadata) - - const blob = await bucket.get(metadata.blobKey) - if (!blob) return recordFromMetadata(metadata) - return withBytes(metadata, new Uint8Array(await blob.arrayBuffer())) - }, - - async list(runId) { - return listArtifacts(runId) - }, - - async delete(artifactId) { - const metadata = await readMetadata(bucket, idKey(artifactId)) - if (!metadata) return - if (metadata.blobKey) await bucket.delete(metadata.blobKey) - await bucket.delete(runKey(metadata.runId, metadata.artifactId)) - await bucket.delete(idKey(artifactId)) - }, - - async deleteForRun(runId) { - const entries = await listRunIndexEntries(runId) - for (const entry of entries) { - const metadata = await readMetadata( - bucket, - idKey(entry.metadata.artifactId), - ) - if (!metadata || metadata.runId !== runId) { - await bucket.delete(entry.key) - continue - } - if (metadata.blobKey) await bucket.delete(metadata.blobKey) - await bucket.delete(entry.key) - await bucket.delete(idKey(entry.metadata.artifactId)) - } - }, - } -} - -/** Build a {@link SqlDriver} over a Cloudflare D1 database (SQLite dialect). */ -export function createD1Driver(d1: D1Database): SqlDriver { - const driver: SqlDriver = { - dialect: 'sqlite', - async exec(sql, params = []) { - await d1 - .prepare(sql) - .bind(...(params as Array)) - .run() - }, - async query( - sql: string, - params: ReadonlyArray = [], - ) { - const result = await d1 - .prepare(sql) - .bind(...(params as Array)) - .all() - return result.results - }, - // D1 has no interactive transaction API (only batch); stores use - // constraints plus reconciliation for CAS when this is a pass-through. - transaction(fn) { - return fn(driver) - }, - } - return driver -} - -export interface CloudflarePersistenceOptions { - d1: D1Database - /** Optional R2 bucket for artifact bytes and generic blob storage. */ - r2?: R2Bucket - /** Optional R2 prefix for artifact bytes and generic blob storage. */ - r2ArtifactPrefix?: string - /** Optional R2 prefix for generic blob storage. Defaults to r2ArtifactPrefix when provided. */ - r2BlobPrefix?: string - /** Optional Durable Object namespace for the distributed lock (see {@link createDurableObjectLockStore}). */ - durableObjects?: DurableObjectNamespace - durableObjectLocks?: DurableObjectLockStoreOptions - mode?: PersistenceMode - /** Run migrations on first use (default false). */ - migrate?: boolean -} - -export type CloudflarePersistence = SqlPersistence & { - stores: SqlPersistence['stores'] & { - locks?: LockStore - artifacts?: ArtifactStore - blobs?: BlobStore - } - /** @deprecated Use stores.locks. */ - locks?: LockStore - /** @deprecated Use stores.artifacts. */ - artifacts?: ArtifactStore -} - -/** Cloudflare-backed {@link CloudflarePersistence} (D1 SQL stores). */ -export function cloudflarePersistence( - opts: CloudflarePersistenceOptions, -): CloudflarePersistence { - const driver = createD1Driver(opts.d1) - const persistence = createSqlPersistence(driver, { - mode: opts.mode, - migrate: opts.migrate, - }) as CloudflarePersistence - if (opts.durableObjects) { - const locks = createDurableObjectLockStore( - opts.durableObjects, - opts.durableObjectLocks, - ) - persistence.stores.locks = locks - persistence.locks = locks - } - if (opts.r2) { - const blobs = createR2BlobStore(opts.r2, { - prefix: opts.r2BlobPrefix ?? opts.r2ArtifactPrefix, - }) - const artifactBlobs = createR2BlobStore(opts.r2, { - prefix: opts.r2ArtifactPrefix ?? DEFAULT_R2_ARTIFACT_PREFIX, - }) - const artifacts = createCloudflareArtifactStore(driver, artifactBlobs, { - migrate: opts.migrate, - }) - persistence.stores.blobs = blobs - persistence.stores.artifacts = artifacts - persistence.artifacts = artifacts - } - return persistence -} - -type LockHolder = { - owner: string - expiresAt: number -} - -async function readJson(request: Request): Promise> { - try { - const value = await request.json() - return value && typeof value === 'object' - ? (value as Record) - : {} - } catch { - return {} - } -} - -function jsonResponse(body: unknown, init?: ResponseInit): Response { - return new Response(JSON.stringify(body), { - ...init, - headers: { - 'content-type': 'application/json', - ...init?.headers, - }, - }) -} - -function validatePositiveFiniteNumber(name: string, value: number): number { - if (!Number.isFinite(value) || value < MIN_LOCK_TIMING_MS) { - throw new TypeError( - `${name} must be a finite positive number of milliseconds.`, - ) - } - return value -} - -/** - * Durable Object class that owns one lock holder per object instance. Bind a DO - * namespace to {@link createDurableObjectLockStore}; each lock key maps to one - * object id via `idFromName`. - */ -export class LockDurableObject { - constructor( - private readonly state: DurableObjectState, - _env?: unknown, - ) {} - - async fetch(request: Request): Promise { - const url = new URL(request.url) - const input = await readJson(request) - const owner = typeof input.owner === 'string' ? input.owner : '' - let leaseMs = DEFAULT_LOCK_LEASE_MS - if (input.leaseMs !== undefined) { - if (typeof input.leaseMs !== 'number') { - return jsonResponse( - { ok: false, error: 'leaseMs must be a finite positive number.' }, - { status: 400 }, - ) - } - try { - leaseMs = validatePositiveFiniteNumber('leaseMs', input.leaseMs) - } catch (error) { - return jsonResponse( - { - ok: false, - error: error instanceof Error ? error.message : 'invalid leaseMs', - }, - { status: 400 }, - ) - } - } - if (!owner) { - return jsonResponse( - { ok: false, error: 'missing owner' }, - { status: 400 }, - ) - } - - if (url.pathname === '/acquire') return this.acquire(owner, leaseMs) - if (url.pathname === '/renew') return this.renew(owner, leaseMs) - if (url.pathname === '/release') return this.release(owner) - return jsonResponse({ ok: false, error: 'not found' }, { status: 404 }) - } - - async alarm(): Promise { - const holder = await this.getHolder() - if (holder && holder.expiresAt <= Date.now()) { - await this.state.storage.delete('holder') - } - } - - private async acquire(owner: string, leaseMs: number): Promise { - const currentTime = Date.now() - const holder = await this.getHolder() - if (!holder || holder.expiresAt <= currentTime || holder.owner === owner) { - const next = { owner, expiresAt: currentTime + leaseMs } - await this.state.storage.put('holder', next) - await this.state.storage.setAlarm(next.expiresAt) - return jsonResponse({ ok: true, expiresAt: next.expiresAt }) - } - - return jsonResponse( - { ok: false, retryAfterMs: Math.max(0, holder.expiresAt - currentTime) }, - { status: 409 }, - ) - } - - private async renew(owner: string, leaseMs: number): Promise { - const holder = await this.getHolder() - if (!holder || holder.owner !== owner) { - return jsonResponse({ ok: false }, { status: 409 }) - } - const next = { owner, expiresAt: Date.now() + leaseMs } - await this.state.storage.put('holder', next) - await this.state.storage.setAlarm(next.expiresAt) - return jsonResponse({ ok: true, expiresAt: next.expiresAt }) - } - - private async release(owner: string): Promise { - const holder = await this.getHolder() - if (!holder || holder.owner !== owner) { - return jsonResponse({ ok: false }, { status: 409 }) - } - await this.state.storage.delete('holder') - await this.state.storage.deleteAlarm() - return jsonResponse({ ok: true }) - } - - private async getHolder(): Promise { - const holder = await this.state.storage.get('holder') - if (!holder) return null - if (holder.expiresAt <= Date.now()) { - await this.state.storage.delete('holder') - return null - } - return holder - } -} - -export function createDurableObjectLockStore( - ns: DurableObjectNamespace, - options?: DurableObjectLockStoreOptions, -): LockStore { - const leaseMs = validatePositiveFiniteNumber( - 'leaseMs', - options?.leaseMs ?? DEFAULT_LOCK_LEASE_MS, - ) - const pollMs = validatePositiveFiniteNumber( - 'pollMs', - options?.pollMs ?? DEFAULT_LOCK_POLL_MS, - ) - const sleep = (ms: number) => - new Promise((resolve) => setTimeout(resolve, Math.max(0, ms))) - const responseError = async (operation: string, response: Response) => { - let body = '' - try { - body = await response.text() - } catch { - body = '' - } - return new Error( - `Durable Object lock ${operation} failed with HTTP ${response.status}${ - body ? `: ${body}` : '' - }`, - ) - } - - return { - async withLock(key: string, fn: () => Promise): Promise { - const stub = ns.get(ns.idFromName(key)) - const owner = crypto.randomUUID() - for (;;) { - const response = await stub.fetch('https://lock/acquire', { - method: 'POST', - body: JSON.stringify({ owner, leaseMs }), - }) - if (response.ok) break - if (response.status !== 409) { - throw await responseError('acquire', response) - } - await sleep(pollMs) - } - - let released = false - let renewing = false - let failRenewal: (error: Error) => void = () => {} - const renewalFailed = new Promise((_, reject) => { - failRenewal = reject - }) - const renew = setInterval( - () => { - if (released) return - if (renewing) return - renewing = true - void (async () => { - try { - const response = await stub.fetch('https://lock/renew', { - method: 'POST', - body: JSON.stringify({ owner, leaseMs }), - }) - if (!response.ok) { - released = true - clearInterval(renew) - failRenewal(await responseError('renew', response)) - } - } catch (cause) { - released = true - clearInterval(renew) - failRenewal( - new Error('Durable Object lock renew failed', { cause }), - ) - } finally { - renewing = false - } - })() - }, - Math.max(1, Math.floor(leaseMs / 2)), - ) - - let result: T | undefined - let primaryError: unknown - try { - const criticalSection = fn() - criticalSection.catch(() => {}) - result = await Promise.race([criticalSection, renewalFailed]) - } catch (error) { - primaryError = error - } - - released = true - clearInterval(renew) - - let releaseError: unknown - try { - const response = await stub.fetch('https://lock/release', { - method: 'POST', - body: JSON.stringify({ owner }), - }) - if (!response.ok) { - releaseError = await responseError('release', response) - } - } catch (error) { - releaseError = error - } - - if (primaryError !== undefined) throw primaryError - if (releaseError !== undefined) throw releaseError - return result as T - }, - } -} - -type ArtifactRow = SqlRow & { - artifact_id: string - run_id: string - thread_id: string - name: string - mime_type: string - size: number - external_url: string | null - blob_key: string | null - created_at: number - updated_at: number -} - -function artifactFromRow(row: ArtifactRow): ArtifactRecord { - return { - artifactId: String(row.artifact_id), - runId: String(row.run_id), - threadId: String(row.thread_id), - name: String(row.name), - mimeType: String(row.mime_type), - size: Number(row.size), - externalUrl: - row.external_url === null ? undefined : String(row.external_url), - createdAt: Number(row.created_at), - } -} - -export function cloudflareArtifactDdl(): Array { - return [ - ` - CREATE TABLE IF NOT EXISTS artifacts ( - artifact_id TEXT PRIMARY KEY, - run_id TEXT NOT NULL, - thread_id TEXT NOT NULL, - name TEXT NOT NULL, - mime_type TEXT NOT NULL, - size INTEGER NOT NULL, - external_url TEXT, - blob_key TEXT, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL - ) - `, - 'CREATE INDEX IF NOT EXISTS idx_artifacts_run_id ON artifacts (run_id)', - ] -} - -/** Build a D1/SQL-indexed artifact store whose bytes live in a BlobStore. */ -export function createCloudflareArtifactStore( - driver: SqlDriver, - blobStore: BlobStore, - options?: CloudflareArtifactStoreOptions, -): ArtifactStore { - const blobKeyPrefix = normalizeBlobPrefix( - options?.blobKeyPrefix ?? 'artifacts/', - ) - const now = options?.now ?? Date.now - const migrate = options?.migrate ?? false - let migrated: Promise | undefined - const blobKey = (artifactId: string) => - `${blobKeyPrefix}${orderedKeyPart(artifactId)}` - const ensureArtifactsTable = () => { - if (!migrate) return Promise.resolve() - migrated ??= (async () => { - for (const statement of cloudflareArtifactDdl()) { - await driver.exec(statement) - } - })() - return migrated - } - const findRow = async (artifactId: string) => { - await ensureArtifactsTable() - const rows = await driver.query( - 'SELECT * FROM artifacts WHERE artifact_id = ?', - [artifactId], - ) - return rows[0] ?? null - } - - return { - async save(record) { - await ensureArtifactsTable() - const existing = await findRow(record.artifactId) - const nextBlobKey = record.bytes ? blobKey(record.artifactId) : undefined - let wroteNextBlob = false - - if (record.bytes && nextBlobKey) { - await blobStore.put(nextBlobKey, record.bytes, { - contentType: record.mimeType, - customMetadata: { - artifactId: record.artifactId, - runId: record.runId, - threadId: record.threadId, - }, - }) - wroteNextBlob = true - } - - try { - await driver.exec( - ` - INSERT INTO artifacts ( - artifact_id, - run_id, - thread_id, - name, - mime_type, - size, - external_url, - blob_key, - created_at, - updated_at - ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - ON CONFLICT (artifact_id) DO UPDATE SET - run_id = excluded.run_id, - thread_id = excluded.thread_id, - name = excluded.name, - mime_type = excluded.mime_type, - size = excluded.size, - external_url = excluded.external_url, - blob_key = excluded.blob_key, - created_at = excluded.created_at, - updated_at = excluded.updated_at - `, - [ - record.artifactId, - record.runId, - record.threadId, - record.name, - record.mimeType, - record.size, - record.externalUrl ?? null, - nextBlobKey ?? null, - record.createdAt, - now(), - ], - ) - } catch (error) { - if (wroteNextBlob && nextBlobKey) { - try { - await blobStore.delete(nextBlobKey) - } catch { - // Preserve the SQL failure; leaked staged blobs are non-indexed. - } - } - throw error - } - - const oldBlobKey = existing?.blob_key - if (oldBlobKey && oldBlobKey !== nextBlobKey) { - try { - await blobStore.delete(String(oldBlobKey)) - } catch { - // The row already points at the new blob/external URL, so stale blob - // cleanup is best-effort and cannot be rediscovered by retrying save. - } - } - }, - - async get(artifactId) { - const row = await findRow(artifactId) - if (!row) return null - const record = artifactFromRow(row) - if (!row.blob_key) return record - const object = await blobStore.get(String(row.blob_key)) - if (!object) return record - return { - ...record, - bytes: new Uint8Array(await object.arrayBuffer()), - } - }, - - async list(runId) { - await ensureArtifactsTable() - const rows = await driver.query( - 'SELECT * FROM artifacts WHERE run_id = ? ORDER BY created_at, artifact_id', - [runId], - ) - return rows.map(artifactFromRow) - }, - - async delete(artifactId) { - const row = await findRow(artifactId) - if (!row) return - if (row.blob_key) await blobStore.delete(String(row.blob_key)) - await driver.exec('DELETE FROM artifacts WHERE artifact_id = ?', [ - artifactId, - ]) - }, - - async deleteForRun(runId) { - await ensureArtifactsTable() - const rows = await driver.query( - 'SELECT * FROM artifacts WHERE run_id = ?', - [runId], - ) - await Promise.all( - rows - .map((row) => row.blob_key) - .filter((key): key is string => typeof key === 'string') - .map((key) => blobStore.delete(key)), - ) - await driver.exec('DELETE FROM artifacts WHERE run_id = ?', [runId]) - }, - } -} diff --git a/packages/ai-persistence-cloudflare/tests/d1.test.ts b/packages/ai-persistence-cloudflare/tests/d1.test.ts deleted file mode 100644 index c8c8ce4ba..000000000 --- a/packages/ai-persistence-cloudflare/tests/d1.test.ts +++ /dev/null @@ -1,1301 +0,0 @@ -/// -import { describe, expect, it } from 'vitest' -import { DatabaseSync } from 'node:sqlite' -import { EventType } from '@tanstack/ai' -import type { StreamChunk } from '@tanstack/ai' -import { - cloudflarePersistence, - createD1Driver, - createR2BlobStore, - createR2ArtifactStore, - createDurableObjectLockStore, - createCloudflareArtifactStore, - cloudflareArtifactDdl, - LockDurableObject, -} from '../src/index' -import type { SqlDriver } from '@tanstack/ai-persistence-sql' -import type { ArtifactRecord } from '@tanstack/ai-persistence' - -/** - * A fake D1Database backed by node:sqlite. D1 is SQLite-compatible, so this - * exercises the real D1 driver + SQL stores end-to-end without a Workers runtime. - */ -function fakeD1(): D1Database { - const db = new DatabaseSync(':memory:') - const prepare = (sql: string) => { - let bound: Array = [] - const api = { - bind: (...values: Array) => { - bound = values - return api - }, - run: () => { - db.prepare(sql).run(...(bound as Array)) - return Promise.resolve({ success: true }) - }, - all: () => - Promise.resolve({ - results: db.prepare(sql).all(...(bound as Array)), - }), - first: () => Promise.resolve(null), - raw: () => Promise.resolve([]), - } - return api - } - return { prepare } as unknown as D1Database -} - -const text = (delta: string): StreamChunk => ({ - type: EventType.TEXT_MESSAGE_CONTENT, - messageId: 'm1', - delta, - timestamp: 1, -}) - -class FakeR2Bucket { - readonly objects = new Map< - string, - { - body: Uint8Array - httpMetadata?: { contentType?: string } - customMetadata?: Record - uploaded: Date - etag: string - } - >() - readonly deleted: Array = [] - readonly failDeletes = new Set() - readonly failPuts = new Set() - getCount = 0 - - async put( - key: string, - value: string | ArrayBuffer | ArrayBufferView, - options?: { - httpMetadata?: { contentType?: string } - customMetadata?: Record - }, - ) { - if (this.failPuts.has(key)) { - throw new Error(`put failed for ${key}`) - } - let body: Uint8Array - if (typeof value === 'string') { - body = new TextEncoder().encode(value) - } else if (value instanceof ArrayBuffer) { - body = new Uint8Array(value) - } else { - body = new Uint8Array(value.buffer, value.byteOffset, value.byteLength) - } - this.objects.set(key, { - body, - httpMetadata: options?.httpMetadata, - customMetadata: options?.customMetadata, - uploaded: new Date(), - etag: `"${key}:${body.byteLength}"`, - }) - return { - key, - size: body.byteLength, - etag: `"${key}:${body.byteLength}"`, - uploaded: new Date(), - httpMetadata: options?.httpMetadata, - customMetadata: options?.customMetadata, - } - } - - async get(key: string) { - this.getCount++ - return this.object(key) - } - - async head(key: string) { - return this.object(key) - } - - private object(key: string) { - const value = this.objects.get(key) - if (!value) return Promise.resolve(null) - const body = value.body - return { - key, - size: body.byteLength, - etag: value.etag, - uploaded: value.uploaded, - httpMetadata: value.httpMetadata, - customMetadata: value.customMetadata, - body: undefined, - async arrayBuffer() { - return body.buffer.slice( - body.byteOffset, - body.byteOffset + body.byteLength, - ) - }, - async text() { - return new TextDecoder().decode(body) - }, - async json() { - return JSON.parse(new TextDecoder().decode(body)) - }, - } - } - - async delete(key: string) { - if (this.failDeletes.has(key)) { - throw new Error(`delete failed for ${key}`) - } - this.deleted.push(key) - this.objects.delete(key) - } - - async list(options?: { prefix?: string; cursor?: string }) { - return { - objects: [...this.objects.keys()] - .filter((key) => key.startsWith(options?.prefix ?? '')) - .sort() - .map((key) => { - const value = this.objects.get(key)! - return { - key, - size: value.body.byteLength, - etag: value.etag, - uploaded: value.uploaded, - httpMetadata: value.httpMetadata, - customMetadata: value.customMetadata, - } - }), - truncated: false, - } - } - - bucket(): R2Bucket { - return this as unknown as R2Bucket - } -} - -async function r2MetadataBlobKey( - bucket: FakeR2Bucket, - artifactId = 'art1', -): Promise { - const object = await bucket.get( - `test-artifacts/by-id/${artifactId}/metadata.json`, - ) - const metadata = (await object?.json()) as { blobKey?: string } | undefined - if (!metadata?.blobKey) throw new Error(`missing blob key for ${artifactId}`) - return metadata.blobKey -} - -function failNextArtifactUpsert(driver: SqlDriver): SqlDriver { - let shouldFail = true - return { - ...driver, - async exec(sql: string, params?: ReadonlyArray) { - if (shouldFail && /INSERT INTO artifacts/.test(sql)) { - shouldFail = false - throw new Error('artifact upsert failed') - } - await driver.exec(sql, params) - }, - } -} - -class FakeDurableObjectState { - readonly storage = { - values: new Map(), - get: async (key: string) => this.storage.values.get(key), - put: async (key: string, value: unknown) => { - this.storage.values.set(key, value) - }, - delete: async (key: string) => { - this.storage.values.delete(key) - }, - setAlarm: async (_timestamp: number) => {}, - deleteAlarm: async () => {}, - } -} - -class FakeDurableObjectNamespace { - readonly instances = new Map< - string, - { object: LockDurableObject; chain: Promise } - >() - - idFromName(name: string) { - return { name } - } - - get(id: { name: string }) { - const current = this.instances.get(id.name) ?? { - object: new LockDurableObject( - new FakeDurableObjectState() as unknown as DurableObjectState, - ), - chain: Promise.resolve(), - } - this.instances.set(id.name, current) - - return { - fetch: (input: RequestInfo | URL, init?: RequestInit) => { - const request = - input instanceof Request ? input : new Request(input, init) - const run = current.chain.then(() => current.object.fetch(request)) - current.chain = run.then( - () => undefined, - () => undefined, - ) - return run - }, - } - } - - namespace(): DurableObjectNamespace { - return this as unknown as DurableObjectNamespace - } -} - -class FetchDurableObjectNamespace { - constructor( - private readonly fetchImpl: ( - input: RequestInfo | URL, - init?: RequestInit, - ) => Promise, - ) {} - - idFromName(name: string) { - return { name } - } - - get(_id: { name: string }) { - return { - fetch: this.fetchImpl, - } - } - - namespace(): DurableObjectNamespace { - return this as unknown as DurableObjectNamespace - } -} - -const artifact = (overrides: Partial = {}): ArtifactRecord => ({ - artifactId: 'art1', - runId: 'run1', - threadId: 'thread1', - name: 'out.txt', - mimeType: 'text/plain', - size: 5, - createdAt: 10, - ...overrides, -}) - -describe('cloudflarePersistence (D1 via fake backed by node:sqlite)', () => { - it('does not create core SQL tables by default', async () => { - const p = cloudflarePersistence({ d1: fakeD1() }) - - await expect( - p.runs!.createOrResume({ runId: 'r1', threadId: 't1', startedAt: 1 }), - ).rejects.toThrow(/no such table: runs/) - }) - - it('round-trips runs and events through the D1 driver', async () => { - const p = cloudflarePersistence({ d1: fakeD1(), migrate: true }) - await p.runs!.createOrResume({ runId: 'r1', threadId: 't1', startedAt: 1 }) - await p.events!.append('r1', 1, text('a')) - await p.events!.append('r1', 2, text('b')) - - expect((await p.runs!.get('r1'))?.status).toBe('running') - expect(await p.events!.latestSeq('r1')).toBe(2) - - const deltas: Array = [] - for await (const e of p.events!.read('r1', { afterSeq: 0 })) { - if (e.event.type === 'TEXT_MESSAGE_CONTENT') deltas.push(e.event.delta) - } - expect(deltas).toEqual(['a', 'b']) - }) - - it('createD1Driver reports the sqlite dialect', () => { - expect(createD1Driver(fakeD1()).dialect).toBe('sqlite') - }) -}) - -describe('createR2ArtifactStore', () => { - it('stores artifact bytes separately from metadata and hydrates get results', async () => { - const bucket = new FakeR2Bucket() - const store = createR2ArtifactStore(bucket.bucket(), { - prefix: 'test-artifacts/', - }) - const bytes = new TextEncoder().encode('hello') - - await store.save(artifact({ bytes })) - - const metadataObject = await bucket.get( - 'test-artifacts/by-id/art1/metadata.json', - ) - expect(await metadataObject?.json()).not.toHaveProperty('bytes') - expect( - [...bucket.objects.keys()].some((key) => - key.startsWith('test-artifacts/blobs/art1/'), - ), - ).toBe(true) - - await expect(store.get('art1')).resolves.toMatchObject({ - artifactId: 'art1', - runId: 'run1', - bytes, - }) - - const list = await store.list('run1') - expect(list).toHaveLength(1) - expect(list[0]).toMatchObject({ - artifactId: 'art1', - runId: 'run1', - name: 'out.txt', - }) - expect(list[0]?.bytes).toBeUndefined() - }) - - it('stores externalUrl-only artifacts without requiring bytes', async () => { - const store = createR2ArtifactStore(new FakeR2Bucket().bucket(), { - prefix: 'test-artifacts/', - }) - - await store.save( - artifact({ - artifactId: 'url-artifact', - externalUrl: 'https://example.com/out.txt', - size: 0, - }), - ) - - const stored = await store.get('url-artifact') - expect(stored).toMatchObject({ - artifactId: 'url-artifact', - externalUrl: 'https://example.com/out.txt', - }) - expect(stored?.bytes).toBeUndefined() - - const list = await store.list('run1') - expect(list).toMatchObject([ - { - artifactId: 'url-artifact', - externalUrl: 'https://example.com/out.txt', - }, - ]) - expect(list[0]?.bytes).toBeUndefined() - }) - - it('wires R2 artifacts into cloudflarePersistence without replacing D1 core stores', () => { - const p = cloudflarePersistence({ - d1: fakeD1(), - r2: new FakeR2Bucket().bucket(), - }) - - expect(p.stores.artifacts).toBeDefined() - expect(p.stores.runs).toBeDefined() - expect(p.stores.publicEvents).toBeDefined() - expect(p.runs).toBe(p.stores.runs) - }) - - it('cleans up stale run indexes and blob bytes on overwrite', async () => { - const bucket = new FakeR2Bucket() - const store = createR2ArtifactStore(bucket.bucket(), { - prefix: 'test-artifacts/', - }) - await store.save( - artifact({ - runId: 'old-run', - bytes: new TextEncoder().encode('hello'), - }), - ) - - await store.save( - artifact({ - runId: 'new-run', - externalUrl: 'https://example.com/out.txt', - }), - ) - - await expect(store.list('old-run')).resolves.toEqual([]) - await expect(store.list('new-run')).resolves.toHaveLength(1) - expect(bucket.deleted).toContain('test-artifacts/by-run/old-run/art1.json') - expect( - bucket.deleted.some((key) => - key.startsWith('test-artifacts/blobs/art1/'), - ), - ).toBe(true) - }) - - it('keeps committed metadata when stale cleanup fails after overwrite', async () => { - const bucket = new FakeR2Bucket() - const store = createR2ArtifactStore(bucket.bucket(), { - prefix: 'test-artifacts/', - }) - await store.save( - artifact({ - runId: 'old-run', - bytes: new TextEncoder().encode('hello'), - }), - ) - bucket.failDeletes.add('test-artifacts/by-run/old-run/art1.json') - - await expect( - store.save( - artifact({ - runId: 'new-run', - externalUrl: 'https://example.com/out.txt', - size: 0, - }), - ), - ).resolves.toBeUndefined() - - await expect(store.get('art1')).resolves.toMatchObject({ - artifactId: 'art1', - runId: 'new-run', - externalUrl: 'https://example.com/out.txt', - size: 0, - }) - expect((await store.get('art1'))?.bytes).toBeUndefined() - await expect(store.list('new-run')).resolves.toMatchObject([ - { - artifactId: 'art1', - runId: 'new-run', - externalUrl: 'https://example.com/out.txt', - }, - ]) - await expect(store.list('old-run')).resolves.toEqual([]) - }) - - it('keeps old bytes when overwrite metadata commit fails after new blob write', async () => { - const bucket = new FakeR2Bucket() - const store = createR2ArtifactStore(bucket.bucket(), { - prefix: 'test-artifacts/', - }) - await store.save( - artifact({ - bytes: new TextEncoder().encode('old'), - size: 3, - }), - ) - const oldBlobKeys = [...bucket.objects.keys()].filter((key) => - key.startsWith('test-artifacts/blobs/art1'), - ) - - bucket.failPuts.add('test-artifacts/by-run/run1/art1.json') - - await expect( - store.save( - artifact({ - bytes: new TextEncoder().encode('new'), - size: 3, - }), - ), - ).rejects.toThrow(/put failed/) - - await expect(store.get('art1')).resolves.toMatchObject({ - artifactId: 'art1', - bytes: new TextEncoder().encode('old'), - }) - expect( - [...bucket.objects.keys()].filter((key) => - key.startsWith('test-artifacts/blobs/art1'), - ), - ).toEqual(oldBlobKeys) - }) - - it('keeps by-id as source of truth when by-id commit fails after by-run commit', async () => { - const bucket = new FakeR2Bucket() - const store = createR2ArtifactStore(bucket.bucket(), { - prefix: 'test-artifacts/', - }) - await store.save( - artifact({ - runId: 'old-run', - bytes: new TextEncoder().encode('old'), - size: 3, - }), - ) - - bucket.failPuts.add('test-artifacts/by-id/art1/metadata.json') - - await expect( - store.save( - artifact({ - runId: 'new-run', - bytes: new TextEncoder().encode('new'), - size: 3, - }), - ), - ).rejects.toThrow(/put failed/) - - await expect(store.get('art1')).resolves.toMatchObject({ - artifactId: 'art1', - runId: 'old-run', - bytes: new TextEncoder().encode('old'), - }) - await expect(store.list('new-run')).resolves.toEqual([]) - expect( - [...bucket.objects.keys()].filter((key) => - key.startsWith('test-artifacts/blobs/art1'), - ), - ).toHaveLength(1) - }) - - it('returns committed by-id metadata when same-run by-run metadata is stale', async () => { - const bucket = new FakeR2Bucket() - const store = createR2ArtifactStore(bucket.bucket(), { - prefix: 'test-artifacts/', - }) - await store.save( - artifact({ - runId: 'run1', - name: 'committed.txt', - bytes: new TextEncoder().encode('old'), - size: 3, - }), - ) - - await bucket.put( - 'test-artifacts/by-run/run1/art1.json', - JSON.stringify( - artifact({ - runId: 'run1', - name: 'uncommitted.txt', - bytes: new TextEncoder().encode('new'), - size: 3, - }), - ), - ) - - await expect(store.list('run1')).resolves.toMatchObject([ - { - artifactId: 'art1', - runId: 'run1', - name: 'committed.txt', - }, - ]) - }) - - it('restores the old run index when same-run by-id commit fails after by-run commit', async () => { - const bucket = new FakeR2Bucket() - const store = createR2ArtifactStore(bucket.bucket(), { - prefix: 'test-artifacts/', - }) - await store.save( - artifact({ - runId: 'run1', - bytes: new TextEncoder().encode('old'), - size: 3, - }), - ) - - bucket.failPuts.add('test-artifacts/by-id/art1/metadata.json') - - await expect( - store.save( - artifact({ - runId: 'run1', - bytes: new TextEncoder().encode('new'), - size: 3, - }), - ), - ).rejects.toThrow(/put failed/) - - await expect(store.list('run1')).resolves.toMatchObject([ - { - artifactId: 'art1', - runId: 'run1', - size: 3, - }, - ]) - await expect(store.get('art1')).resolves.toMatchObject({ - artifactId: 'art1', - runId: 'run1', - bytes: new TextEncoder().encode('old'), - }) - }) - - it('ignores stale run index entries when deleting artifacts for an old run', async () => { - const bucket = new FakeR2Bucket() - const store = createR2ArtifactStore(bucket.bucket(), { - prefix: 'test-artifacts/', - }) - await store.save( - artifact({ - runId: 'old-run', - bytes: new TextEncoder().encode('old'), - }), - ) - bucket.failDeletes.add('test-artifacts/by-run/old-run/art1.json') - await store.save( - artifact({ - runId: 'new-run', - bytes: new TextEncoder().encode('new'), - }), - ) - bucket.failDeletes.delete('test-artifacts/by-run/old-run/art1.json') - - await store.deleteForRun?.('old-run') - - await expect(store.get('art1')).resolves.toMatchObject({ - artifactId: 'art1', - runId: 'new-run', - bytes: new TextEncoder().encode('new'), - }) - await expect(store.list('old-run')).resolves.toEqual([]) - await expect(store.list('new-run')).resolves.toHaveLength(1) - }) - - it('deleteForRun works when destructured from the store', async () => { - const bucket = new FakeR2Bucket() - const store = createR2ArtifactStore(bucket.bucket(), { - prefix: 'test-artifacts/', - }) - await store.save( - artifact({ - bytes: new TextEncoder().encode('hello'), - }), - ) - - const deleteForRun = store.deleteForRun! - await deleteForRun('run1') - - await expect(store.list('run1')).resolves.toEqual([]) - expect(bucket.deleted).toContain('test-artifacts/by-id/art1/metadata.json') - expect( - bucket.deleted.some((key) => - key.startsWith('test-artifacts/blobs/art1/'), - ), - ).toBe(true) - }) - - it('keeps metadata indexes when blob delete by artifact id fails so delete can retry', async () => { - const bucket = new FakeR2Bucket() - const store = createR2ArtifactStore(bucket.bucket(), { - prefix: 'test-artifacts/', - }) - await store.save(artifact({ bytes: new TextEncoder().encode('hello') })) - const blobKey = await r2MetadataBlobKey(bucket) - bucket.failDeletes.add(blobKey) - - await expect(store.delete?.('art1')).rejects.toThrow(/delete failed/) - - expect(bucket.objects.has('test-artifacts/by-id/art1/metadata.json')).toBe( - true, - ) - expect(bucket.objects.has('test-artifacts/by-run/run1/art1.json')).toBe( - true, - ) - expect(bucket.objects.has(blobKey)).toBe(true) - - bucket.failDeletes.delete(blobKey) - await store.delete?.('art1') - await expect(store.get('art1')).resolves.toBeNull() - expect(bucket.objects.has(blobKey)).toBe(false) - }) - - it('keeps by-id metadata when by-run delete by artifact id fails so delete can retry', async () => { - const bucket = new FakeR2Bucket() - const store = createR2ArtifactStore(bucket.bucket(), { - prefix: 'test-artifacts/', - }) - await store.save(artifact({ bytes: new TextEncoder().encode('hello') })) - bucket.failDeletes.add('test-artifacts/by-run/run1/art1.json') - - await expect(store.delete?.('art1')).rejects.toThrow(/delete failed/) - - expect(bucket.objects.has('test-artifacts/by-id/art1/metadata.json')).toBe( - true, - ) - expect(bucket.objects.has('test-artifacts/by-run/run1/art1.json')).toBe( - true, - ) - expect( - [...bucket.objects.keys()].some((key) => - key.startsWith('test-artifacts/blobs/art1/'), - ), - ).toBe(false) - - bucket.failDeletes.delete('test-artifacts/by-run/run1/art1.json') - await store.delete?.('art1') - await expect(store.get('art1')).resolves.toBeNull() - await expect(store.list('run1')).resolves.toEqual([]) - }) - - it('keeps metadata indexes when blob delete by run fails so deleteForRun can retry', async () => { - const bucket = new FakeR2Bucket() - const store = createR2ArtifactStore(bucket.bucket(), { - prefix: 'test-artifacts/', - }) - await store.save( - artifact({ - artifactId: 'art1', - bytes: new TextEncoder().encode('one'), - }), - ) - await store.save( - artifact({ - artifactId: 'art2', - bytes: new TextEncoder().encode('two'), - }), - ) - const blobKey = await r2MetadataBlobKey(bucket, 'art1') - bucket.failDeletes.add(blobKey) - - await expect(store.deleteForRun?.('run1')).rejects.toThrow(/delete failed/) - - expect(bucket.objects.has('test-artifacts/by-id/art1/metadata.json')).toBe( - true, - ) - expect(bucket.objects.has('test-artifacts/by-run/run1/art1.json')).toBe( - true, - ) - expect(bucket.objects.has(blobKey)).toBe(true) - - bucket.failDeletes.delete(blobKey) - await store.deleteForRun?.('run1') - await expect(store.list('run1')).resolves.toEqual([]) - expect(bucket.objects.has(blobKey)).toBe(false) - }) -}) - -describe('createR2BlobStore', () => { - it('stores, reads, heads, lists, and deletes logical keys with internal R2 prefixes', async () => { - const bucket = new FakeR2Bucket() - const store = createR2BlobStore(bucket.bucket(), { - prefix: 'internal/blobs/', - }) - - await expect( - store.put('runs/run1/out.txt', 'hello', { - contentType: 'text/plain', - customMetadata: { runId: 'run1' }, - }), - ).resolves.toMatchObject({ - key: 'runs/run1/out.txt', - size: 5, - contentType: 'text/plain', - customMetadata: { runId: 'run1' }, - }) - - expect(bucket.objects.has('internal/blobs/runs/run1/out.txt')).toBe(true) - - await expect(store.head('runs/run1/out.txt')).resolves.toMatchObject({ - key: 'runs/run1/out.txt', - size: 5, - contentType: 'text/plain', - customMetadata: { runId: 'run1' }, - }) - await expect((await store.get('runs/run1/out.txt'))?.text()).resolves.toBe( - 'hello', - ) - - const page = await store.list({ prefix: 'runs/run1/' }) - expect(page.objects).toMatchObject([ - { - key: 'runs/run1/out.txt', - contentType: 'text/plain', - customMetadata: { runId: 'run1' }, - }, - ]) - - await store.delete('runs/run1/out.txt') - expect(bucket.objects.has('internal/blobs/runs/run1/out.txt')).toBe(false) - }) -}) - -describe('D1-indexed R2 artifacts', () => { - it('uses separate R2 prefixes for D1 artifacts and generic blobs', async () => { - const bucket = new FakeR2Bucket() - const p = cloudflarePersistence({ - d1: fakeD1(), - r2: bucket.bucket(), - r2ArtifactPrefix: 'artifact-bytes/', - r2BlobPrefix: 'generic-blobs/', - migrate: true, - }) - - await p.stores.artifacts!.save( - artifact({ bytes: new TextEncoder().encode('hello') }), - ) - await p.stores.blobs!.put('uploads/raw.txt', 'raw') - - expect( - [...bucket.objects.keys()].some((key) => - key.startsWith('artifact-bytes/artifacts/art1/'), - ), - ).toBe(true) - expect( - [...bucket.objects.keys()].some((key) => - key.startsWith('generic-blobs/uploads/raw.txt'), - ), - ).toBe(true) - }) - - it('wires BlobStore and D1-indexed ArtifactStore through cloudflarePersistence', async () => { - const bucket = new FakeR2Bucket() - const p = cloudflarePersistence({ - d1: fakeD1(), - r2: bucket.bucket(), - r2ArtifactPrefix: 'artifact-bytes/', - migrate: true, - }) - - expect(p.stores.blobs).toBeDefined() - expect(p.stores.artifacts).toBeDefined() - - await p.stores.artifacts!.save( - artifact({ bytes: new TextEncoder().encode('hello') }), - ) - - expect( - [...bucket.objects.keys()].some((key) => - key.startsWith('artifact-bytes/artifacts/art1/'), - ), - ).toBe(true) - await expect(p.stores.artifacts!.get('art1')).resolves.toMatchObject({ - artifactId: 'art1', - bytes: new TextEncoder().encode('hello'), - }) - - bucket.getCount = 0 - await expect(p.stores.artifacts!.list('run1')).resolves.toMatchObject([ - { artifactId: 'art1', runId: 'run1', name: 'out.txt' }, - ]) - expect(bucket.getCount).toBe(0) - }) - - it('deletes artifact metadata and R2 bytes by artifact id and run id', async () => { - const bucket = new FakeR2Bucket() - const p = cloudflarePersistence({ - d1: fakeD1(), - r2: bucket.bucket(), - r2ArtifactPrefix: 'artifact-bytes/', - migrate: true, - }) - const store = p.stores.artifacts! - - await store.save( - artifact({ - artifactId: 'art1', - bytes: new TextEncoder().encode('one'), - }), - ) - await store.save( - artifact({ - artifactId: 'art2', - bytes: new TextEncoder().encode('two'), - }), - ) - - await store.delete?.('art1') - await expect(store.get('art1')).resolves.toBeNull() - expect( - bucket.deleted.some((key) => - key.startsWith('artifact-bytes/artifacts/art1/'), - ), - ).toBe(true) - - await store.deleteForRun?.('run1') - await expect(store.list('run1')).resolves.toEqual([]) - expect( - bucket.deleted.some((key) => - key.startsWith('artifact-bytes/artifacts/art2/'), - ), - ).toBe(true) - }) - - it('deletes stale bytes when an artifact is overwritten with an external URL only', async () => { - const bucket = new FakeR2Bucket() - const p = cloudflarePersistence({ - d1: fakeD1(), - r2: bucket.bucket(), - r2ArtifactPrefix: 'artifact-bytes/', - migrate: true, - }) - const store = p.stores.artifacts! - - await store.save( - artifact({ - bytes: new TextEncoder().encode('hello'), - }), - ) - await store.save( - artifact({ - externalUrl: 'https://example.com/out.txt', - size: 0, - }), - ) - - expect( - bucket.deleted.some((key) => - key.startsWith('artifact-bytes/artifacts/art1/'), - ), - ).toBe(true) - await expect(store.get('art1')).resolves.toMatchObject({ - artifactId: 'art1', - externalUrl: 'https://example.com/out.txt', - }) - expect((await store.get('art1'))?.bytes).toBeUndefined() - }) - - it('keeps an external URL overwrite when post-commit stale blob cleanup fails', async () => { - const bucket = new FakeR2Bucket() - const p = cloudflarePersistence({ - d1: fakeD1(), - r2: bucket.bucket(), - r2ArtifactPrefix: 'artifact-bytes/', - migrate: true, - }) - const store = p.stores.artifacts! - - await store.save( - artifact({ - bytes: new TextEncoder().encode('hello'), - }), - ) - const staleBlobKey = [...bucket.objects.keys()].find((key) => - key.startsWith('artifact-bytes/artifacts/art1/'), - )! - bucket.failDeletes.add(staleBlobKey) - - await expect( - store.save( - artifact({ - externalUrl: 'https://example.com/out.txt', - size: 0, - }), - ), - ).resolves.toBeUndefined() - - await expect(store.get('art1')).resolves.toMatchObject({ - artifactId: 'art1', - externalUrl: 'https://example.com/out.txt', - size: 0, - }) - expect((await store.get('art1'))?.bytes).toBeUndefined() - expect(bucket.objects.has(staleBlobKey)).toBe(true) - }) - - it('does not auto-create artifact tables by default', async () => { - const p = cloudflarePersistence({ - d1: fakeD1(), - r2: new FakeR2Bucket().bucket(), - }) - - await expect(p.stores.artifacts!.list('run1')).rejects.toThrow( - /no such table: artifacts/, - ) - }) - - it('keeps existing metadata and bytes consistent when overwrite upsert fails after blob write', async () => { - const bucket = new FakeR2Bucket() - const driver = createD1Driver(fakeD1()) - const blobs = createR2BlobStore(bucket.bucket(), { - prefix: 'artifact-bytes/', - }) - const store = createCloudflareArtifactStore(driver, blobs, { - migrate: true, - }) - await store.save( - artifact({ - bytes: new TextEncoder().encode('old'), - size: 3, - }), - ) - - const failingStore = createCloudflareArtifactStore( - failNextArtifactUpsert(driver), - blobs, - { migrate: true }, - ) - - await expect( - failingStore.save( - artifact({ - bytes: new TextEncoder().encode('new'), - size: 3, - }), - ), - ).rejects.toThrow('artifact upsert failed') - - await expect(store.get('art1')).resolves.toMatchObject({ - bytes: new TextEncoder().encode('old'), - }) - expect( - [...bucket.objects.keys()].filter((key) => - key.includes('artifact-bytes/artifacts/art1'), - ), - ).toHaveLength(1) - }) - - it('keeps a D1 artifact row when blob delete by artifact id fails so delete can retry', async () => { - const bucket = new FakeR2Bucket() - const p = cloudflarePersistence({ - d1: fakeD1(), - r2: bucket.bucket(), - r2ArtifactPrefix: 'artifact-bytes/', - migrate: true, - }) - const store = p.stores.artifacts! - await store.save(artifact({ bytes: new TextEncoder().encode('hello') })) - const stored = await store.get('art1') - expect(stored?.bytes).toEqual(new TextEncoder().encode('hello')) - const physicalKey = [...bucket.objects.keys()].find((key) => - key.includes('artifact-bytes/artifacts/art1'), - )! - bucket.failDeletes.add(physicalKey) - - await expect(store.delete?.('art1')).rejects.toThrow(/delete failed/) - - await expect(store.get('art1')).resolves.toMatchObject({ - artifactId: 'art1', - bytes: new TextEncoder().encode('hello'), - }) - - bucket.failDeletes.delete(physicalKey) - await store.delete?.('art1') - await expect(store.get('art1')).resolves.toBeNull() - }) - - it('keeps D1 artifact rows when blob delete by run fails so deleteForRun can retry', async () => { - const bucket = new FakeR2Bucket() - const p = cloudflarePersistence({ - d1: fakeD1(), - r2: bucket.bucket(), - r2ArtifactPrefix: 'artifact-bytes/', - migrate: true, - }) - const store = p.stores.artifacts! - await store.save( - artifact({ - artifactId: 'art1', - bytes: new TextEncoder().encode('one'), - }), - ) - await store.save( - artifact({ - artifactId: 'art2', - bytes: new TextEncoder().encode('two'), - }), - ) - const physicalKey = [...bucket.objects.keys()].find((key) => - key.includes('artifact-bytes/artifacts/art1'), - )! - bucket.failDeletes.add(physicalKey) - - await expect(store.deleteForRun?.('run1')).rejects.toThrow(/delete failed/) - - await expect(store.list('run1')).resolves.toHaveLength(2) - - bucket.failDeletes.delete(physicalKey) - await store.deleteForRun?.('run1') - await expect(store.list('run1')).resolves.toEqual([]) - }) - - it('exports artifact table DDL for self-managed Cloudflare migrations', () => { - const statements = cloudflareArtifactDdl() - expect(statements.join('\n')).toMatch( - /CREATE TABLE IF NOT EXISTS artifacts/, - ) - expect(statements.join('\n')).toMatch(/idx_artifacts_run_id/) - }) -}) - -describe('Durable Object locks', () => { - it('serializes same-key work and releases after thrown callbacks', async () => { - const ns = new FakeDurableObjectNamespace() - const locks = createDurableObjectLockStore(ns.namespace(), { - leaseMs: 250, - pollMs: 1, - }) - const events: Array = [] - - const first = locks.withLock('thread1', async () => { - events.push('first:start') - await new Promise((resolve) => setTimeout(resolve, 20)) - events.push('first:end') - }) - const second = locks.withLock('thread1', async () => { - events.push('second') - }) - - await Promise.all([first, second]) - expect(events).toEqual(['first:start', 'first:end', 'second']) - - await expect( - locks.withLock('thread1', async () => { - throw new Error('boom') - }), - ).rejects.toThrow('boom') - - await expect( - locks.withLock('thread1', async () => 'after-error'), - ).resolves.toBe('after-error') - }) - - it('does not release a holder when the owner token does not match', async () => { - const object = new LockDurableObject( - new FakeDurableObjectState() as unknown as DurableObjectState, - ) - - const acquire = await object.fetch( - new Request('https://lock/acquire', { - method: 'POST', - body: JSON.stringify({ owner: 'owner-a', leaseMs: 60_000 }), - }), - ) - expect(acquire.status).toBe(200) - - const wrongRelease = await object.fetch( - new Request('https://lock/release', { - method: 'POST', - body: JSON.stringify({ owner: 'owner-b' }), - }), - ) - expect(wrongRelease.status).toBe(409) - - const blocked = await object.fetch( - new Request('https://lock/acquire', { - method: 'POST', - body: JSON.stringify({ owner: 'owner-c', leaseMs: 60_000 }), - }), - ) - expect(blocked.status).toBe(409) - }) - - it('surfaces unexpected acquire failures instead of retrying forever', async () => { - let attempts = 0 - const ns = new FetchDurableObjectNamespace(async () => { - attempts++ - if (attempts > 1) throw new Error('retried unexpected acquire failure') - return Response.json({ error: 'storage unavailable' }, { status: 503 }) - }) - const locks = createDurableObjectLockStore(ns.namespace(), { - leaseMs: 250, - pollMs: 1, - }) - - await expect( - locks.withLock('thread1', async () => 'never'), - ).rejects.toThrow(/acquire.*503.*storage unavailable/) - expect(attempts).toBe(1) - }) - - it('rejects the critical section when lock renewal fails', async () => { - let renewCalls = 0 - const ns = new FetchDurableObjectNamespace(async (input) => { - const path = new URL(String(input)).pathname - if (path === '/acquire') return Response.json({ ok: true }) - if (path === '/renew') { - renewCalls++ - return Response.json({ error: 'lease lost' }, { status: 409 }) - } - return Response.json({ ok: true }) - }) - const locks = createDurableObjectLockStore(ns.namespace(), { - leaseMs: 10, - pollMs: 1, - }) - - await expect( - locks.withLock( - 'thread1', - () => new Promise((resolve) => setTimeout(() => resolve('late'), 50)), - ), - ).rejects.toThrow(/renew.*409.*lease lost/) - expect(renewCalls).toBe(1) - }) - - it('preserves callback errors when release also fails', async () => { - const ns = new FetchDurableObjectNamespace(async (input) => { - const path = new URL(String(input)).pathname - if (path === '/release') { - return Response.json({ error: 'release failed' }, { status: 503 }) - } - return Response.json({ ok: true }) - }) - const locks = createDurableObjectLockStore(ns.namespace(), { - leaseMs: 250, - pollMs: 1, - }) - - await expect( - locks.withLock('thread1', async () => { - throw new Error('callback failed') - }), - ).rejects.toThrow('callback failed') - }) - - it('preserves renewal errors when release also fails', async () => { - const ns = new FetchDurableObjectNamespace(async (input) => { - const path = new URL(String(input)).pathname - if (path === '/acquire') return Response.json({ ok: true }) - if (path === '/renew') { - return Response.json({ error: 'lease lost' }, { status: 409 }) - } - return Response.json({ error: 'release failed' }, { status: 503 }) - }) - const locks = createDurableObjectLockStore(ns.namespace(), { - leaseMs: 10, - pollMs: 1, - }) - - await expect( - locks.withLock( - 'thread1', - () => new Promise((resolve) => setTimeout(() => resolve('late'), 50)), - ), - ).rejects.toThrow(/renew.*lease lost/) - }) - - it('surfaces release failures when the critical section succeeds', async () => { - const ns = new FetchDurableObjectNamespace(async (input) => { - const path = new URL(String(input)).pathname - if (path === '/release') { - return Response.json({ error: 'release failed' }, { status: 503 }) - } - return Response.json({ ok: true }) - }) - const locks = createDurableObjectLockStore(ns.namespace(), { - leaseMs: 250, - pollMs: 1, - }) - - await expect(locks.withLock('thread1', async () => 'ok')).rejects.toThrow( - /release.*503.*release failed/, - ) - }) - - it('rejects invalid lock timing options before creating a store', () => { - expect(() => - createDurableObjectLockStore( - new FakeDurableObjectNamespace().namespace(), - { - leaseMs: 0, - }, - ), - ).toThrow(/leaseMs/) - expect(() => - createDurableObjectLockStore( - new FakeDurableObjectNamespace().namespace(), - { - pollMs: Number.POSITIVE_INFINITY, - }, - ), - ).toThrow(/pollMs/) - }) - - it('rejects invalid lock timing request payloads', async () => { - const object = new LockDurableObject( - new FakeDurableObjectState() as unknown as DurableObjectState, - ) - - const response = await object.fetch( - new Request('https://lock/acquire', { - method: 'POST', - body: JSON.stringify({ owner: 'owner-a', leaseMs: 0 }), - }), - ) - - expect(response.status).toBe(400) - await expect(response.text()).resolves.toMatch(/leaseMs/) - }) -}) diff --git a/packages/ai-persistence-cloudflare/tsconfig.json b/packages/ai-persistence-cloudflare/tsconfig.json deleted file mode 100644 index c38689f4e..000000000 --- a/packages/ai-persistence-cloudflare/tsconfig.json +++ /dev/null @@ -1,8 +0,0 @@ -{ - "extends": "../../tsconfig.base.json", - "compilerOptions": { - "outDir": "dist" - }, - "include": ["src", "tests"], - "exclude": ["node_modules", "dist"] -} diff --git a/packages/ai-persistence-cloudflare/vite.config.ts b/packages/ai-persistence-cloudflare/vite.config.ts deleted file mode 100644 index 11f5b20b7..000000000 --- a/packages/ai-persistence-cloudflare/vite.config.ts +++ /dev/null @@ -1,37 +0,0 @@ -import { defineConfig, mergeConfig } from 'vitest/config' -import { tanstackViteConfig } from '@tanstack/vite-config' -import packageJson from './package.json' - -const config = defineConfig({ - test: { - name: packageJson.name, - dir: './', - watch: false, - - globals: true, - environment: 'node', - include: ['tests/**/*.test.ts'], - coverage: { - provider: 'v8', - reporter: ['text', 'json', 'html', 'lcov'], - exclude: [ - 'node_modules/', - 'dist/', - 'tests/', - '**/*.test.ts', - '**/*.config.ts', - '**/types.ts', - ], - include: ['src/**/*.ts'], - }, - }, -}) - -export default mergeConfig( - config, - tanstackViteConfig({ - entry: ['./src/index.ts'], - srcDir: './src', - cjs: false, - }), -) diff --git a/packages/ai-persistence-drizzle/drizzle.config.ts b/packages/ai-persistence-drizzle/drizzle.config.ts new file mode 100644 index 000000000..f9d243c25 --- /dev/null +++ b/packages/ai-persistence-drizzle/drizzle.config.ts @@ -0,0 +1,14 @@ +import { defineConfig } from 'drizzle-kit' + +/** + * drizzle-kit config for the batteries-included sqlite backend. Migrations are + * generated from `src/schema.ts` into `drizzle/` and shipped with the package + * (see the `files` field); `sqlPersistence({ migrate: true })` applies them. + * + * Regenerate with `pnpm db:generate` after any schema change. + */ +export default defineConfig({ + dialect: 'sqlite', + schema: './src/schema.ts', + out: './drizzle', +}) diff --git a/packages/ai-persistence-drizzle/drizzle/0000_cute_vision.sql b/packages/ai-persistence-drizzle/drizzle/0000_cute_vision.sql new file mode 100644 index 000000000..76e81e8fd --- /dev/null +++ b/packages/ai-persistence-drizzle/drizzle/0000_cute_vision.sql @@ -0,0 +1,54 @@ +CREATE TABLE `artifacts` ( + `artifact_id` text PRIMARY KEY NOT NULL, + `run_id` text NOT NULL, + `thread_id` text NOT NULL, + `name` text NOT NULL, + `mime_type` text NOT NULL, + `size` integer NOT NULL, + `external_url` text, + `created_at` integer NOT NULL +); +--> statement-breakpoint +CREATE TABLE `blobs` ( + `key` text PRIMARY KEY NOT NULL, + `content_type` text, + `size` integer, + `etag` text, + `custom_metadata_json` text, + `created_at` integer, + `updated_at` integer, + `body` blob +); +--> statement-breakpoint +CREATE TABLE `interrupts` ( + `interrupt_id` text PRIMARY KEY NOT NULL, + `run_id` text NOT NULL, + `thread_id` text NOT NULL, + `status` text NOT NULL, + `requested_at` integer NOT NULL, + `resolved_at` integer, + `payload_json` text NOT NULL, + `response_json` text +); +--> statement-breakpoint +CREATE TABLE `messages` ( + `thread_id` text PRIMARY KEY NOT NULL, + `messages_json` text NOT NULL +); +--> statement-breakpoint +CREATE TABLE `metadata` ( + `scope` text NOT NULL, + `key` text NOT NULL, + `value_json` text NOT NULL, + PRIMARY KEY(`scope`, `key`) +); +--> statement-breakpoint +CREATE TABLE `runs` ( + `run_id` text PRIMARY KEY NOT NULL, + `thread_id` text NOT NULL, + `status` text NOT NULL, + `started_at` integer NOT NULL, + `finished_at` integer, + `error` text, + `usage_json` text +); diff --git a/packages/ai-persistence-drizzle/drizzle/meta/0000_snapshot.json b/packages/ai-persistence-drizzle/drizzle/meta/0000_snapshot.json new file mode 100644 index 000000000..589e34589 --- /dev/null +++ b/packages/ai-persistence-drizzle/drizzle/meta/0000_snapshot.json @@ -0,0 +1,338 @@ +{ + "version": "6", + "dialect": "sqlite", + "id": "5e224922-d0ca-4cac-bf7b-5813ad6f6d5d", + "prevId": "00000000-0000-0000-0000-000000000000", + "tables": { + "artifacts": { + "name": "artifacts", + "columns": { + "artifact_id": { + "name": "artifact_id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "thread_id": { + "name": "thread_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "name": { + "name": "name", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "mime_type": { + "name": "mime_type", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "size": { + "name": "size", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "external_url": { + "name": "external_url", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "created_at": { + "name": "created_at", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "blobs": { + "name": "blobs", + "columns": { + "key": { + "name": "key", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "content_type": { + "name": "content_type", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "size": { + "name": "size", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "etag": { + "name": "etag", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "custom_metadata_json": { + "name": "custom_metadata_json", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "created_at": { + "name": "created_at", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "updated_at": { + "name": "updated_at", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "body": { + "name": "body", + "type": "blob", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "interrupts": { + "name": "interrupts", + "columns": { + "interrupt_id": { + "name": "interrupt_id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "thread_id": { + "name": "thread_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "requested_at": { + "name": "requested_at", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "resolved_at": { + "name": "resolved_at", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "payload_json": { + "name": "payload_json", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "response_json": { + "name": "response_json", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "messages": { + "name": "messages", + "columns": { + "thread_id": { + "name": "thread_id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "messages_json": { + "name": "messages_json", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "metadata": { + "name": "metadata", + "columns": { + "scope": { + "name": "scope", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "key": { + "name": "key", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "value_json": { + "name": "value_json", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "metadata_scope_key_pk": { + "columns": [ + "scope", + "key" + ], + "name": "metadata_scope_key_pk" + } + }, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "runs": { + "name": "runs", + "columns": { + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "thread_id": { + "name": "thread_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "started_at": { + "name": "started_at", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "finished_at": { + "name": "finished_at", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "error": { + "name": "error", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "usage_json": { + "name": "usage_json", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + } + }, + "views": {}, + "enums": {}, + "_meta": { + "schemas": {}, + "tables": {}, + "columns": {} + }, + "internal": { + "indexes": {} + } +} \ No newline at end of file diff --git a/packages/ai-persistence-drizzle/drizzle/meta/_journal.json b/packages/ai-persistence-drizzle/drizzle/meta/_journal.json new file mode 100644 index 000000000..ce52b8084 --- /dev/null +++ b/packages/ai-persistence-drizzle/drizzle/meta/_journal.json @@ -0,0 +1,13 @@ +{ + "version": "7", + "dialect": "sqlite", + "entries": [ + { + "idx": 0, + "version": "6", + "when": 1783594836281, + "tag": "0000_cute_vision", + "breakpoints": true + } + ] +} \ No newline at end of file diff --git a/packages/ai-persistence-drizzle/package.json b/packages/ai-persistence-drizzle/package.json index b44e5c61e..70a5f8b49 100644 --- a/packages/ai-persistence-drizzle/package.json +++ b/packages/ai-persistence-drizzle/package.json @@ -1,7 +1,7 @@ { "name": "@tanstack/ai-persistence-drizzle", "version": "0.1.0", - "description": "Drizzle ORM backend for TanStack AI persistence — bring your own Drizzle db (sqlite or postgres dialect); persists over the underlying client.", + "description": "Drizzle ORM backend for TanStack AI state persistence — the batteries-included sqlPersistence (bundled migrations) and bring-your-own drizzlePersistence(db), with an exported schema you can drive with drizzle-kit.", "author": "", "license": "MIT", "repository": { @@ -12,14 +12,13 @@ "keywords": [ "ai", "tanstack", - "persistence" + "persistence", + "drizzle", + "sqlite" ], "type": "module", "module": "./dist/esm/index.js", "types": "./dist/esm/index.d.ts", - "bin": { - "tanstack-ai-persistence-drizzle": "./dist/bin/bin.js" - }, "exports": { ".": { "types": "./dist/esm/index.d.ts", @@ -28,11 +27,14 @@ }, "files": [ "dist", - "src" + "src", + "drizzle", + "drizzle.config.ts" ], "scripts": { - "build": "vite build && tsup --config tsup.bin.config.ts", + "build": "vite build", "clean": "premove ./build ./dist", + "db:generate": "drizzle-kit generate", "lint:fix": "eslint ./src --fix", "test:build": "publint --strict", "test:eslint": "eslint ./src", @@ -43,13 +45,13 @@ "peerDependencies": { "@tanstack/ai": "workspace:^", "@tanstack/ai-persistence": "workspace:^", - "@tanstack/ai-persistence-sql": "workspace:^" + "drizzle-orm": ">=0.44.0" }, "devDependencies": { "@tanstack/ai": "workspace:*", "@tanstack/ai-persistence": "workspace:*", - "@tanstack/ai-persistence-sql": "workspace:*", "@vitest/coverage-v8": "4.0.14", - "tsup": "^8.5.1" + "drizzle-kit": "^0.31.0", + "drizzle-orm": "^0.45.0" } } diff --git a/packages/ai-persistence-drizzle/src/bin.ts b/packages/ai-persistence-drizzle/src/bin.ts deleted file mode 100644 index 447ca521a..000000000 --- a/packages/ai-persistence-drizzle/src/bin.ts +++ /dev/null @@ -1,17 +0,0 @@ -import { runCli } from './cli' - -try { - const result = await runCli(process.argv.slice(2), { - writeStdout: (value) => process.stdout.write(value), - }) - if (result.kind === 'help') { - process.stdout.write(result.text) - } else if (result.kind === 'file') { - process.stdout.write(`Wrote ${result.path}\n`) - } -} catch (error) { - process.stderr.write( - `${error instanceof Error ? error.message : String(error)}\n`, - ) - process.exitCode = 1 -} diff --git a/packages/ai-persistence-drizzle/src/cli-core.ts b/packages/ai-persistence-drizzle/src/cli-core.ts deleted file mode 100644 index c5fdb5992..000000000 --- a/packages/ai-persistence-drizzle/src/cli-core.ts +++ /dev/null @@ -1,31 +0,0 @@ -import { join } from 'node:path' -import { - formatMigrationSql, - runMigrationCli as runSharedMigrationCli, -} from '@tanstack/ai-persistence-sql/cli-core' -import type { - CliOptions, - CliResult, - MigrationCliConfig, -} from '@tanstack/ai-persistence-sql/cli-core' - -export type CliDialect = 'sqlite' | 'postgres' -export type { CliOptions, CliResult } - -const drizzleMigrationCliConfig: MigrationCliConfig = { - commandName: 'tanstack-ai-persistence-drizzle', - defaultDialect: 'sqlite', - allowedDialects: ['sqlite', 'postgres'], - defaultName: 'tanstack_ai_persistence', - defaultOut: ({ cwd, timestamp, name }) => - join(cwd, 'drizzle', `${timestamp}_${name}.sql`), -} - -export { formatMigrationSql } - -export function runMigrationCli( - argv: Array, - options?: CliOptions, -): Promise { - return runSharedMigrationCli(argv, drizzleMigrationCliConfig, options) -} diff --git a/packages/ai-persistence-drizzle/src/cli.ts b/packages/ai-persistence-drizzle/src/cli.ts deleted file mode 100644 index 6c4518c11..000000000 --- a/packages/ai-persistence-drizzle/src/cli.ts +++ /dev/null @@ -1,9 +0,0 @@ -import { runMigrationCli } from './cli-core' -import type { CliOptions, CliResult } from './cli-core' - -export function runCli( - argv: Array, - options?: CliOptions, -): Promise { - return runMigrationCli(argv, options) -} diff --git a/packages/ai-persistence-drizzle/src/index.ts b/packages/ai-persistence-drizzle/src/index.ts index 6081f805d..7a7dd8f6f 100644 --- a/packages/ai-persistence-drizzle/src/index.ts +++ b/packages/ai-persistence-drizzle/src/index.ts @@ -1,105 +1,145 @@ /** - * Drizzle backend (bring-your-own). A Drizzle `db` exposes its underlying client - * at `db.$client`: a `better-sqlite3`/`node:sqlite`-shaped handle for the sqlite - * dialect, or a node-postgres `Pool` for the postgres dialect. This adapter - * unwraps `$client` and persists over it through the shared SQL stores — so it - * reuses the exact, tested driver logic rather than re-deriving SQL through - * Drizzle's query builder. + * Drizzle-backed **state** persistence for TanStack AI. * - * Schema: the tables are the same as the raw SQL backend. Re-export `ddl` so you - * can apply them with your own Drizzle migration workflow, or opt in to lazy - * backend migrations with `migrate: true`. + * Two entry points, one `AIPersistence` contract: + * + * - {@link sqlPersistence} — batteries-included. Give it a `dialect` + `url` and + * it builds the database, applies the migrations bundled in `drizzle/` + * (generated by drizzle-kit), and returns a ready `AIPersistence`. + * - {@link drizzlePersistence} — bring-your-own. Pass a Drizzle sqlite `db` you + * already constructed and migrated (via `drizzle-kit` against the exported + * {@link schema}); this just wires the stores over it. + * + * Locks are not part of the SQL schema — an in-memory lock is provided as a dev + * default so the returned persistence is complete. Swap in a distributed lock + * for multi-process deployments. */ -import { createSqlPersistence, ddl } from '@tanstack/ai-persistence-sql' -import type { - Dialect, - SqlDriver, - SqlPersistence, - SqlRow, -} from '@tanstack/ai-persistence-sql' -import type { PersistenceMode } from '@tanstack/ai-persistence' +import { existsSync, readFileSync } from 'node:fs' +import { dirname, join } from 'node:path' +import { fileURLToPath } from 'node:url' +import { DatabaseSync } from 'node:sqlite' +import { drizzle } from 'drizzle-orm/sqlite-proxy' +import { InMemoryLockStore } from '@tanstack/ai' +import { + createArtifactStore, + createBlobStore, + createInterruptStore, + createMessageStore, + createMetadataStore, + createRunStore, +} from './stores' +import { schema } from './schema' +import type { DrizzleDb } from './stores' +import type { AIPersistence } from '@tanstack/ai-persistence' -export { ddl } +export { schema } +export type { DrizzleDb } from './stores' +export * from './schema' -/** SQLite-shaped client (node:sqlite DatabaseSync / better-sqlite3 Database). */ -interface SqliteClient { - prepare: (sql: string) => { - run: (...params: Array) => unknown - all: (...params: Array) => Array +/** Wire the AIPersistence stores over an existing Drizzle sqlite database. */ +export function drizzlePersistence(db: DrizzleDb): AIPersistence { + return { + stores: { + messages: createMessageStore(db), + runs: createRunStore(db), + interrupts: createInterruptStore(db), + metadata: createMetadataStore(db), + artifacts: createArtifactStore(db), + blobs: createBlobStore(db), + locks: new InMemoryLockStore(), + }, } } -/** node-postgres Pool-shaped client. */ -interface PgClient { - query: ( - sql: string, - params?: ReadonlyArray, - ) => Promise<{ rows: Array> }> -} -/** A Drizzle db exposing its underlying driver client. */ -export interface DrizzleDb { - $client: unknown +export interface SqlPersistenceOptions { + /** + * SQL dialect. Only `'sqlite'` is bundled with batteries (migrations shipped + * in `drizzle/`), so it is the sole accepted value and the default. For + * `postgres`/`mysql`, construct your own Drizzle db and use + * {@link drizzlePersistence} with migrations generated from the exported + * {@link schema} — `sqlPersistence` cannot build those for you. + */ + dialect?: 'sqlite' + /** Database URL/path. For sqlite: `:memory:` or a file path (`file:` prefix optional). */ + url: string + /** Apply the bundled migrations on construction (default false). */ + migrate?: boolean } -function sqliteDriver(client: SqliteClient): SqlDriver { - const driver: SqlDriver = { - dialect: 'sqlite', - exec(sql, params = []) { - client.prepare(sql).run(...params) - return Promise.resolve() - }, - query( - sql: string, - params: ReadonlyArray = [], - ) { - return Promise.resolve(client.prepare(sql).all(...params) as Array) - }, - transaction(fn) { - return fn(driver) - }, +/** Build an {@link AIPersistence} from a dialect + url, applying bundled migrations. */ +export function sqlPersistence(opts: SqlPersistenceOptions): AIPersistence { + // `dialect` is typed `'sqlite'` only, but guard at runtime so an untyped JS + // caller passing another dialect fails loud instead of silently getting a + // sqlite db. + const dialect = opts.dialect ?? 'sqlite' + // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition + if (dialect !== 'sqlite') { + throw new Error( + `sqlPersistence bundles the sqlite dialect only; for '${dialect}' construct a Drizzle db yourself and use drizzlePersistence(db) with migrations generated from the exported schema.`, + ) } - return driver + const sqlite = new DatabaseSync(normalizeSqliteUrl(opts.url)) + if (opts.migrate) applyBundledMigrations(sqlite) + const db = drizzle((sql, params, method) => { + const stmt = sqlite.prepare(sql) + if (method === 'run') { + stmt.run(...params) + return Promise.resolve({ rows: [] }) + } + if (method === 'get') { + const row = stmt.get(...params) + // The store impl only ever selects via `.all()`, so `get` is unused; a + // missing row maps to an empty positional array either way. + return Promise.resolve({ rows: row ? Object.values(row) : [] }) + } + const rows = stmt.all(...params) + return Promise.resolve({ rows: rows.map((row) => Object.values(row)) }) + }) + return drizzlePersistence(db) } -function pgDriver(client: PgClient): SqlDriver { - const driver: SqlDriver = { - dialect: 'postgres', - async exec(sql, params = []) { - await client.query(sql, params) - }, - async query( - sql: string, - params: ReadonlyArray = [], - ) { - return (await client.query(sql, params)).rows as Array - }, - // Without a dedicated connection we can't BEGIN/COMMIT; stores use - // constraints plus reconciliation for CAS when this is a pass-through. - transaction(fn) { - return fn(driver) - }, - } - return driver +function normalizeSqliteUrl(url: string): string { + if (url === ':memory:') return url + return url.startsWith('file:') ? url.slice('file:'.length) : url } -export interface DrizzlePersistenceOptions { - db: DrizzleDb - dialect: Dialect - mode?: PersistenceMode - /** Run migrations on first use (default false). */ - migrate?: boolean +interface MigrationJournal { + entries?: Array<{ tag: string }> } -/** Drizzle-backed {@link SqlPersistence}. */ -export function drizzlePersistence( - opts: DrizzlePersistenceOptions, -): SqlPersistence { - const driver = - opts.dialect === 'postgres' - ? pgDriver(opts.db.$client as PgClient) - : sqliteDriver(opts.db.$client as SqliteClient) - return createSqlPersistence(driver, { - mode: opts.mode, - migrate: opts.migrate, - }) +/** Locate the bundled `drizzle/` migrations dir relative to this module (src or dist). */ +function findMigrationsDir(): string { + let dir = dirname(fileURLToPath(import.meta.url)) + for (let i = 0; i < 8; i++) { + const candidate = join(dir, 'drizzle') + if (existsSync(join(candidate, 'meta', '_journal.json'))) return candidate + const parent = dirname(dir) + if (parent === dir) break + dir = parent + } + throw new Error('Could not locate the bundled drizzle migrations directory.') +} + +function applyBundledMigrations(sqlite: DatabaseSync): void { + const migrationsDir = findMigrationsDir() + const journal = JSON.parse( + readFileSync(join(migrationsDir, 'meta', '_journal.json'), 'utf8'), + ) as MigrationJournal + sqlite.exec( + 'CREATE TABLE IF NOT EXISTS __drizzle_migrations (tag TEXT PRIMARY KEY, applied_at INTEGER NOT NULL)', + ) + const appliedRows = sqlite + .prepare('SELECT tag FROM __drizzle_migrations') + .all() as Array<{ tag: string }> + const applied = new Set(appliedRows.map((row) => row.tag)) + for (const entry of journal.entries ?? []) { + if (applied.has(entry.tag)) continue + const sql = readFileSync(join(migrationsDir, `${entry.tag}.sql`), 'utf8') + sqlite.exec(sql) + sqlite + .prepare( + 'INSERT INTO __drizzle_migrations (tag, applied_at) VALUES (?, ?)', + ) + .run(entry.tag, Date.now()) + } } diff --git a/packages/ai-persistence-drizzle/src/schema.ts b/packages/ai-persistence-drizzle/src/schema.ts new file mode 100644 index 000000000..5a9d1c9b2 --- /dev/null +++ b/packages/ai-persistence-drizzle/src/schema.ts @@ -0,0 +1,102 @@ +/** + * Drizzle schema for the TanStack AI **state** persistence contract. + * + * Each table mirrors the corresponding record in + * `@tanstack/ai-persistence`'s `types.ts` column-for-column. JSON-valued fields + * are stored in `*_json` text columns (portable across sqlite/pg/mysql); epoch + * millisecond timestamps are stored as integers. + * + * This schema is the single source of truth for migrations: run + * `pnpm db:generate` (drizzle-kit) after any change here. It is also exported + * from the package so bring-your-own-drizzle users can drive their own + * migration workflow against it. + * + * NOTE: keep this in sync with the sibling Prisma schema fragment + * (`@tanstack/ai-persistence-prisma`). See coupling `persistence-schema-dual-source`. + */ +import { blob, integer, primaryKey, sqliteTable, text } from 'drizzle-orm/sqlite-core' +import type { + InterruptRecord, + RunStatus, +} from '@tanstack/ai-persistence' +import type { ModelMessage, TokenUsage } from '@tanstack/ai' + +/** Thread message history (`MessageStore`). */ +export const messages = sqliteTable('messages', { + threadId: text('thread_id').primaryKey(), + messagesJson: text('messages_json', { mode: 'json' }) + .$type>() + .notNull(), +}) + +/** Run lifecycle records (`RunStore`). */ +export const runs = sqliteTable('runs', { + runId: text('run_id').primaryKey(), + threadId: text('thread_id').notNull(), + status: text('status').$type().notNull(), + startedAt: integer('started_at').notNull(), + finishedAt: integer('finished_at'), + error: text('error'), + usageJson: text('usage_json', { mode: 'json' }).$type(), +}) + +/** Interrupt / approval records (`InterruptStore`). */ +export const interrupts = sqliteTable('interrupts', { + interruptId: text('interrupt_id').primaryKey(), + runId: text('run_id').notNull(), + threadId: text('thread_id').notNull(), + status: text('status').$type().notNull(), + requestedAt: integer('requested_at').notNull(), + resolvedAt: integer('resolved_at'), + payloadJson: text('payload_json', { mode: 'json' }) + .$type>() + .notNull(), + responseJson: text('response_json', { mode: 'json' }).$type(), +}) + +/** Scoped key/value metadata (`MetadataStore`). */ +export const metadata = sqliteTable( + 'metadata', + { + scope: text('scope').notNull(), + key: text('key').notNull(), + valueJson: text('value_json', { mode: 'json' }).$type().notNull(), + }, + (table) => [primaryKey({ columns: [table.scope, table.key] })], +) + +/** Generation artifact references (`ArtifactStore`). Bytes live in the blob store. */ +export const artifacts = sqliteTable('artifacts', { + artifactId: text('artifact_id').primaryKey(), + runId: text('run_id').notNull(), + threadId: text('thread_id').notNull(), + name: text('name').notNull(), + mimeType: text('mime_type').notNull(), + size: integer('size').notNull(), + externalUrl: text('external_url'), + createdAt: integer('created_at').notNull(), +}) + +/** Generic blob objects (`BlobStore`). */ +export const blobs = sqliteTable('blobs', { + key: text('key').primaryKey(), + contentType: text('content_type'), + size: integer('size'), + etag: text('etag'), + customMetadataJson: text('custom_metadata_json', { mode: 'json' }).$type< + Record + >(), + createdAt: integer('created_at'), + updatedAt: integer('updated_at'), + body: blob('body', { mode: 'buffer' }), +}) + +/** The full state schema, for `drizzlePersistence(db)` and drizzle-kit. */ +export const schema = { + messages, + runs, + interrupts, + metadata, + artifacts, + blobs, +} diff --git a/packages/ai-persistence-drizzle/src/stores.ts b/packages/ai-persistence-drizzle/src/stores.ts new file mode 100644 index 000000000..08af6bb84 --- /dev/null +++ b/packages/ai-persistence-drizzle/src/stores.ts @@ -0,0 +1,498 @@ +/** + * AIPersistence store implementations over a Drizzle sqlite database. + * + * Each method mirrors the semantics of the reference in-memory backend + * (`@tanstack/ai-persistence`'s `memory.ts`). JSON columns are handled by + * Drizzle's `text({ mode: 'json' })`; blob bytes by `blob({ mode: 'buffer' })`. + */ +import { and, asc, eq, gt, gte, lt } from 'drizzle-orm' +import { artifacts, blobs, interrupts, messages, metadata, runs } from './schema' +import type { SQL } from 'drizzle-orm' +import type { BaseSQLiteDatabase } from 'drizzle-orm/sqlite-core' +import type { ModelMessage } from '@tanstack/ai' +import type { + ArtifactRecord, + ArtifactStore, + BlobBody, + BlobListOptions, + BlobListPage, + BlobObject, + BlobRecord, + BlobStore, + InterruptRecord, + InterruptStore, + MessageStore, + MetadataStore, + RunRecord, + RunStore, +} from '@tanstack/ai-persistence' + +/** + * Any Drizzle sqlite database (better-sqlite3, libsql, node:sqlite proxy, D1, …). + * + * Typed as the schema-agnostic slice of the query builder we actually use, so a + * BYO `db` constructed with any `{ schema }` is assignable regardless of its + * `TFullSchema` (which is invariant on the full `BaseSQLiteDatabase`). + */ +export type DrizzleDb = Pick< + BaseSQLiteDatabase<'sync' | 'async', unknown>, + 'select' | 'insert' | 'update' | 'delete' +> + +const textEncoder = new TextEncoder() +const textDecoder = new TextDecoder() + +/** + * Smallest string strictly greater than every string that starts with `prefix`, + * used as the exclusive upper bound of a literal prefix range scan. Increments + * the last non-`U+FFFF` UTF-16 code unit (carrying over trailing `U+FFFF`s). + * Returns `undefined` when no finite bound exists — an empty prefix, or a prefix + * consisting solely of `U+FFFF` — in which case the caller omits the upper bound. + */ +function prefixUpperBound(prefix: string): string | undefined { + let i = prefix.length - 1 + while (i >= 0 && prefix.charCodeAt(i) === 0xffff) i-- + if (i < 0) return undefined + return prefix.slice(0, i) + String.fromCharCode(prefix.charCodeAt(i) + 1) +} + +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.') +} + +export function createMessageStore(db: DrizzleDb): MessageStore { + return { + async loadThread(threadId) { + const rows = await db + .select({ messagesJson: messages.messagesJson }) + .from(messages) + .where(eq(messages.threadId, threadId)) + return rows[0]?.messagesJson ?? [] + }, + async saveThread(threadId, msgs: Array) { + await db + .insert(messages) + .values({ threadId, messagesJson: msgs }) + .onConflictDoUpdate({ + target: messages.threadId, + set: { messagesJson: msgs }, + }) + }, + } +} + +function mapRun(row: typeof runs.$inferSelect): RunRecord { + return { + runId: row.runId, + threadId: row.threadId, + status: row.status, + startedAt: row.startedAt, + ...(row.finishedAt != null ? { finishedAt: row.finishedAt } : {}), + ...(row.error != null ? { error: row.error } : {}), + ...(row.usageJson != null ? { usage: row.usageJson } : {}), + } +} + +export function createRunStore(db: DrizzleDb): RunStore { + const store: RunStore = { + async createOrResume(input) { + const existing = await store.get(input.runId) + if (existing) return existing + const record: RunRecord = { + runId: input.runId, + threadId: input.threadId, + status: input.status ?? 'running', + startedAt: input.startedAt, + } + await db + .insert(runs) + .values({ + runId: record.runId, + threadId: record.threadId, + status: record.status, + startedAt: record.startedAt, + }) + .onConflictDoNothing({ target: runs.runId }) + return (await store.get(input.runId)) ?? record + }, + async update(runId, patch) { + const set: Partial = {} + if (patch.status !== undefined) set.status = patch.status + if (patch.finishedAt !== undefined) set.finishedAt = patch.finishedAt + if (patch.error !== undefined) set.error = patch.error + if (patch.usage !== undefined) set.usageJson = patch.usage + if (Object.keys(set).length === 0) return + await db.update(runs).set(set).where(eq(runs.runId, runId)) + }, + async get(runId) { + const rows = await db.select().from(runs).where(eq(runs.runId, runId)) + const row = rows[0] + return row ? mapRun(row) : null + }, + } + return store +} + +function mapInterrupt(row: typeof interrupts.$inferSelect): InterruptRecord { + return { + interruptId: row.interruptId, + runId: row.runId, + threadId: row.threadId, + status: row.status, + requestedAt: row.requestedAt, + ...(row.resolvedAt != null ? { resolvedAt: row.resolvedAt } : {}), + payload: row.payloadJson, + ...(row.responseJson != null ? { response: row.responseJson } : {}), + } +} + +export function createInterruptStore(db: DrizzleDb): InterruptStore { + return { + async create(record) { + await db + .insert(interrupts) + .values({ + interruptId: record.interruptId, + runId: record.runId, + threadId: record.threadId, + status: record.status, + requestedAt: record.requestedAt, + payloadJson: record.payload, + responseJson: record.response ?? null, + }) + .onConflictDoNothing({ target: interrupts.interruptId }) + }, + async resolve(interruptId, response) { + await db + .update(interrupts) + .set({ + status: 'resolved', + resolvedAt: Date.now(), + responseJson: response ?? null, + }) + .where(eq(interrupts.interruptId, interruptId)) + }, + async cancel(interruptId) { + await db + .update(interrupts) + .set({ status: 'cancelled', resolvedAt: Date.now() }) + .where(eq(interrupts.interruptId, interruptId)) + }, + async get(interruptId) { + const rows = await db + .select() + .from(interrupts) + .where(eq(interrupts.interruptId, interruptId)) + const row = rows[0] + return row ? mapInterrupt(row) : null + }, + async list(threadId) { + const rows = await db + .select() + .from(interrupts) + .where(eq(interrupts.threadId, threadId)) + .orderBy(asc(interrupts.requestedAt)) + return rows.map(mapInterrupt) + }, + async listPending(threadId) { + const rows = await db + .select() + .from(interrupts) + .where( + and( + eq(interrupts.threadId, threadId), + eq(interrupts.status, 'pending'), + ), + ) + .orderBy(asc(interrupts.requestedAt)) + return rows.map(mapInterrupt) + }, + async listByRun(runId) { + const rows = await db + .select() + .from(interrupts) + .where(eq(interrupts.runId, runId)) + .orderBy(asc(interrupts.requestedAt)) + return rows.map(mapInterrupt) + }, + async listPendingByRun(runId) { + const rows = await db + .select() + .from(interrupts) + .where( + and(eq(interrupts.runId, runId), eq(interrupts.status, 'pending')), + ) + .orderBy(asc(interrupts.requestedAt)) + return rows.map(mapInterrupt) + }, + } +} + +export function createMetadataStore(db: DrizzleDb): MetadataStore { + return { + async get(scope, key) { + const rows = await db + .select({ valueJson: metadata.valueJson }) + .from(metadata) + .where(and(eq(metadata.scope, scope), eq(metadata.key, key))) + const row = rows[0] + return row ? row.valueJson : null + }, + async set(scope, key, value) { + await db + .insert(metadata) + .values({ scope, key, valueJson: value }) + .onConflictDoUpdate({ + target: [metadata.scope, metadata.key], + set: { valueJson: value }, + }) + }, + async delete(scope, key) { + await db + .delete(metadata) + .where(and(eq(metadata.scope, scope), eq(metadata.key, key))) + }, + } +} + +function mapArtifact(row: typeof artifacts.$inferSelect): ArtifactRecord { + return { + artifactId: row.artifactId, + runId: row.runId, + threadId: row.threadId, + name: row.name, + mimeType: row.mimeType, + size: row.size, + ...(row.externalUrl != null ? { externalUrl: row.externalUrl } : {}), + createdAt: row.createdAt, + } +} + +export function createArtifactStore(db: DrizzleDb): ArtifactStore { + return { + async save(record) { + const values = { + artifactId: record.artifactId, + runId: record.runId, + threadId: record.threadId, + name: record.name, + mimeType: record.mimeType, + size: record.size, + externalUrl: record.externalUrl ?? null, + createdAt: record.createdAt, + } + await db + .insert(artifacts) + .values(values) + .onConflictDoUpdate({ + target: artifacts.artifactId, + set: { + runId: values.runId, + threadId: values.threadId, + name: values.name, + mimeType: values.mimeType, + size: values.size, + externalUrl: values.externalUrl, + createdAt: values.createdAt, + }, + }) + }, + async get(artifactId) { + const rows = await db + .select() + .from(artifacts) + .where(eq(artifacts.artifactId, artifactId)) + const row = rows[0] + return row ? mapArtifact(row) : null + }, + async list(runId) { + const rows = await db + .select() + .from(artifacts) + .where(eq(artifacts.runId, runId)) + .orderBy(asc(artifacts.createdAt), asc(artifacts.artifactId)) + return rows.map(mapArtifact) + }, + async delete(artifactId) { + await db.delete(artifacts).where(eq(artifacts.artifactId, artifactId)) + }, + async deleteForRun(runId) { + await db.delete(artifacts).where(eq(artifacts.runId, runId)) + }, + } +} + +function blobRecordSnapshot(row: typeof blobs.$inferSelect): BlobRecord { + return { + key: row.key, + ...(row.size != null ? { size: row.size } : {}), + ...(row.etag != null ? { etag: row.etag } : {}), + ...(row.contentType != null ? { contentType: row.contentType } : {}), + ...(row.customMetadataJson != null + ? { customMetadata: { ...row.customMetadataJson } } + : {}), + ...(row.createdAt != null ? { createdAt: row.createdAt } : {}), + ...(row.updatedAt != null ? { updatedAt: row.updatedAt } : {}), + } +} + +function blobObject(row: typeof blobs.$inferSelect): BlobObject { + const bytes = row.body ? copyBytes(new Uint8Array(row.body)) : new Uint8Array() + return { + ...blobRecordSnapshot(row), + body: new ReadableStream({ + start(controller) { + controller.enqueue(copyBytes(bytes)) + controller.close() + }, + }), + arrayBuffer: () => Promise.resolve(bytesToArrayBuffer(bytes)), + text: () => Promise.resolve(textDecoder.decode(bytes)), + } +} + +export function createBlobStore(db: DrizzleDb): BlobStore { + const readRow = async (key: string) => { + const rows = await db.select().from(blobs).where(eq(blobs.key, key)) + return rows[0] ?? null + } + return { + async put(key, body, options) { + const bytes = await bytesFromBlobBody(body) + const existing = await readRow(key) + const now = Date.now() + const contentType = + options?.contentType ?? + (typeof Blob !== 'undefined' && body instanceof Blob + ? body.type || undefined + : undefined) + const customMetadata = options?.customMetadata + ? { ...options.customMetadata } + : undefined + const record = { + key, + contentType: contentType ?? null, + size: bytes.byteLength, + etag: crypto.randomUUID(), + customMetadataJson: customMetadata ?? null, + createdAt: existing?.createdAt ?? now, + updatedAt: now, + body: Buffer.from(bytes), + } + await db + .insert(blobs) + .values(record) + .onConflictDoUpdate({ + target: blobs.key, + set: { + contentType: record.contentType, + size: record.size, + etag: record.etag, + customMetadataJson: record.customMetadataJson, + createdAt: record.createdAt, + updatedAt: record.updatedAt, + body: record.body, + }, + }) + const stored = await readRow(key) + return stored ? blobRecordSnapshot(stored) : blobRecordSnapshot(record) + }, + async get(key) { + const row = await readRow(key) + return row ? blobObject(row) : null + }, + async head(key) { + const row = await readRow(key) + return row ? blobRecordSnapshot(row) : null + }, + async delete(key) { + await db.delete(blobs).where(eq(blobs.key, key)) + }, + async list(options?: BlobListOptions): Promise { + const limit = options?.limit + if (limit === 0) return { objects: [], truncated: false } + const prefix = options?.prefix ?? '' + // Match the prefix LITERALLY and case-sensitively to mirror the reference + // in-memory backend's `key.startsWith(prefix)`. A SQL `LIKE '${prefix}%'` + // would treat `_`/`%` in the caller's prefix as wildcards and, on SQLite, + // match case-insensitively for ASCII. Instead use a half-open range on the + // (BINARY-collated) key column: `key >= prefix AND key < upperBound`, where + // upperBound is the smallest string strictly greater than every key that + // starts with `prefix`. This contains no LIKE metacharacters and relies on + // the default binary/byte ordering, giving literal, case-sensitive matching. + const conditions: Array = [] + if (prefix !== '') { + conditions.push(gte(blobs.key, prefix)) + const upperBound = prefixUpperBound(prefix) + if (upperBound !== undefined) { + conditions.push(lt(blobs.key, upperBound)) + } + } + if (options?.cursor !== undefined) { + conditions.push(gt(blobs.key, options.cursor)) + } + const base = db + .select() + .from(blobs) + .where(and(...conditions)) + .orderBy(asc(blobs.key)) + const rows = await (limit === undefined ? base : base.limit(limit + 1)) + const pageRows = limit === undefined ? rows : rows.slice(0, limit) + const objects = pageRows.map(blobRecordSnapshot) + const truncated = limit !== undefined && rows.length > limit + return { + objects, + ...(truncated ? { cursor: pageRows.at(-1)?.key, truncated } : {}), + } + }, + } +} diff --git a/packages/ai-persistence-drizzle/tests/cli.test.ts b/packages/ai-persistence-drizzle/tests/cli.test.ts deleted file mode 100644 index f619f902b..000000000 --- a/packages/ai-persistence-drizzle/tests/cli.test.ts +++ /dev/null @@ -1,97 +0,0 @@ -import { mkdtemp, readFile, writeFile } from 'node:fs/promises' -import { join } from 'node:path' -import { tmpdir } from 'node:os' -import { describe, expect, it } from 'vitest' -import { runCli } from '../src/cli' - -describe('Drizzle persistence CLI', () => { - it('writes the default Drizzle migration path', async () => { - const cwd = await mkdtemp(join(tmpdir(), 'tanstack-ai-drizzle-')) - - const result = await runCli( - ['--dialect', 'postgres', '--timestamp', '20260706123456'], - { cwd }, - ) - - expect(result).toEqual({ - kind: 'file', - path: join(cwd, 'drizzle', '20260706123456_tanstack_ai_persistence.sql'), - }) - if (result.kind === 'file') { - await expect(readFile(result.path, 'utf8')).resolves.toContain( - 'TanStack AI persistence migration for postgres', - ) - } - }) - - it('writes a custom output path', async () => { - const cwd = await mkdtemp(join(tmpdir(), 'tanstack-ai-drizzle-')) - const out = join(cwd, 'db', 'migrations', 'custom.sql') - - const result = await runCli(['--dialect', 'sqlite', '--out', out], { cwd }) - - expect(result).toEqual({ kind: 'file', path: out }) - await expect(readFile(out, 'utf8')).resolves.toContain( - 'TanStack AI persistence migration for sqlite', - ) - }) - - it('prints SQL to stdout without writing a file', async () => { - const cwd = await mkdtemp(join(tmpdir(), 'tanstack-ai-drizzle-')) - - const result = await runCli(['--dialect', 'sqlite', '--stdout'], { cwd }) - - expect(result.kind).toBe('stdout') - if (result.kind === 'stdout') { - expect(result.sql).toContain( - 'TanStack AI persistence migration for sqlite', - ) - expect(result.sql).toContain('CREATE TABLE IF NOT EXISTS runs') - } - }) - - it('rejects invalid dialects', async () => { - await expect(runCli(['--dialect', 'mysql'])).rejects.toThrow( - 'Invalid --dialect', - ) - }) - - it('rejects path traversal migration names', async () => { - await expect(runCli(['--name', '../escape'])).rejects.toThrow( - 'Invalid --name', - ) - }) - - it('refuses to overwrite an existing file without --force', async () => { - const cwd = await mkdtemp(join(tmpdir(), 'tanstack-ai-drizzle-')) - const out = join(cwd, 'migration.sql') - await writeFile(out, 'existing') - - await expect(runCli(['--dialect', 'sqlite', '--out', out])).rejects.toThrow( - 'already exists', - ) - }) - - it('overwrites an existing file with --force', async () => { - const cwd = await mkdtemp(join(tmpdir(), 'tanstack-ai-drizzle-')) - const out = join(cwd, 'migration.sql') - await writeFile(out, 'existing') - - const result = await runCli([ - '--dialect', - 'sqlite', - '--out', - out, - '--force', - ]) - - expect(result).toEqual({ kind: 'file', path: out }) - await expect(readFile(out, 'utf8')).resolves.toContain( - 'TanStack AI persistence migration for sqlite', - ) - }) - - it('rejects missing option values', async () => { - await expect(runCli(['--out'])).rejects.toThrow('Missing value for --out') - }) -}) diff --git a/packages/ai-persistence-drizzle/tests/drizzle.conformance.test.ts b/packages/ai-persistence-drizzle/tests/drizzle.conformance.test.ts new file mode 100644 index 000000000..a624bc933 --- /dev/null +++ b/packages/ai-persistence-drizzle/tests/drizzle.conformance.test.ts @@ -0,0 +1,6 @@ +import { runPersistenceConformance } from '@tanstack/ai-persistence/testkit' +import { sqlPersistence } from '../src/index' + +runPersistenceConformance('drizzle-sqlite', () => + sqlPersistence({ dialect: 'sqlite', url: ':memory:', migrate: true }), +) diff --git a/packages/ai-persistence-drizzle/tests/drizzle.test.ts b/packages/ai-persistence-drizzle/tests/drizzle.test.ts deleted file mode 100644 index d22745253..000000000 --- a/packages/ai-persistence-drizzle/tests/drizzle.test.ts +++ /dev/null @@ -1,37 +0,0 @@ -import { describe, expect, it } from 'vitest' -import { DatabaseSync } from 'node:sqlite' -import { EventType } from '@tanstack/ai' -import type { StreamChunk } from '@tanstack/ai' -import { drizzlePersistence } from '../src/index' - -const text = (delta: string): StreamChunk => ({ - type: EventType.TEXT_MESSAGE_CONTENT, - messageId: 'm1', - delta, - timestamp: 1, -}) - -describe('drizzlePersistence (sqlite dialect, $client = node:sqlite)', () => { - it('does not create schema by default', async () => { - const db = { $client: new DatabaseSync(':memory:') } - const p = drizzlePersistence({ db, dialect: 'sqlite' }) - - await expect( - p.runs!.createOrResume({ runId: 'r1', threadId: 't1', startedAt: 1 }), - ).rejects.toThrow(/no such table: runs/) - }) - - it('persists via the unwrapped Drizzle client', async () => { - // A Drizzle sqlite db exposes its driver client at `$client`; node:sqlite's - // DatabaseSync is prepare().run/all-shaped like better-sqlite3. - const db = { $client: new DatabaseSync(':memory:') } - const p = drizzlePersistence({ db, dialect: 'sqlite', migrate: true }) - - await p.runs!.createOrResume({ runId: 'r1', threadId: 't1', startedAt: 1 }) - await p.events!.append('r1', 1, text('a')) - await p.events!.append('r1', 2, text('b')) - - expect((await p.runs!.get('r1'))?.threadId).toBe('t1') - expect(await p.events!.latestSeq('r1')).toBe(2) - }) -}) diff --git a/packages/ai-persistence-drizzle/tsup.bin.config.ts b/packages/ai-persistence-drizzle/tsup.bin.config.ts deleted file mode 100644 index d40ae7297..000000000 --- a/packages/ai-persistence-drizzle/tsup.bin.config.ts +++ /dev/null @@ -1,11 +0,0 @@ -import { defineConfig } from 'tsup' - -export default defineConfig({ - entry: { bin: 'src/bin.ts' }, - outDir: 'dist/bin', - format: ['esm'], - platform: 'node', - target: 'node18', - external: ['@tanstack/ai-persistence-sql'], - banner: { js: '#!/usr/bin/env node' }, -}) diff --git a/packages/ai-persistence-postgres/package.json b/packages/ai-persistence-postgres/package.json deleted file mode 100644 index 14dc22acc..000000000 --- a/packages/ai-persistence-postgres/package.json +++ /dev/null @@ -1,61 +0,0 @@ -{ - "name": "@tanstack/ai-persistence-postgres", - "version": "0.1.0", - "description": "Postgres backend for TanStack AI persistence — durable runs, messages, events, interrupts, and metadata over node-postgres (pg) or a bring-your-own pool.", - "author": "", - "license": "MIT", - "repository": { - "type": "git", - "url": "git+https://github.com/TanStack/ai.git", - "directory": "packages/ai-persistence-postgres" - }, - "keywords": [ - "ai", - "tanstack", - "persistence", - "postgres", - "pg" - ], - "type": "module", - "module": "./dist/esm/index.js", - "types": "./dist/esm/index.d.ts", - "exports": { - ".": { - "types": "./dist/esm/index.d.ts", - "import": "./dist/esm/index.js" - } - }, - "files": [ - "dist", - "src" - ], - "scripts": { - "build": "vite build", - "clean": "premove ./build ./dist", - "lint:fix": "eslint ./src --fix", - "test:build": "publint --strict", - "test:eslint": "eslint ./src", - "test:lib": "vitest", - "test:lib:dev": "pnpm test:lib --watch", - "test:types": "tsc" - }, - "peerDependencies": { - "@tanstack/ai": "workspace:^", - "@tanstack/ai-persistence": "workspace:^", - "@tanstack/ai-persistence-sql": "workspace:^", - "pg": "^8.13.0" - }, - "peerDependenciesMeta": { - "pg": { - "optional": true - } - }, - "devDependencies": { - "@tanstack/ai": "workspace:*", - "@tanstack/ai-persistence": "workspace:*", - "@tanstack/ai-persistence-sql": "workspace:*", - "@types/pg": "^8.11.10", - "@vitest/coverage-v8": "4.0.14", - "pg": "^8.13.0" - } -} diff --git a/packages/ai-persistence-postgres/src/index.ts b/packages/ai-persistence-postgres/src/index.ts deleted file mode 100644 index dba4cd313..000000000 --- a/packages/ai-persistence-postgres/src/index.ts +++ /dev/null @@ -1,134 +0,0 @@ -/** - * Postgres backend. Wraps a `pg`-style pool in the shared `SqlDriver` and - * assembles a `SqlPersistence` via `@tanstack/ai-persistence-sql`. - * - * - convenience: `postgresPersistence({ connectionString })` lazily creates a - * `pg.Pool` (the `pg` package is an optional peer, imported on first use). - * - BYO: pass `{ client }` — any pool exposing `query(sql, params) => { rows }` - * and `connect() => client` (node-postgres `Pool`). - */ -import { createSqlPersistence } from '@tanstack/ai-persistence-sql' -import type { - SqlDriver, - SqlPersistence, - SqlRow, -} from '@tanstack/ai-persistence-sql' -import type { PersistenceMode } from '@tanstack/ai-persistence' - -/** Minimal node-postgres surface the driver relies on. */ -export interface PgQueryable { - query: ( - sql: string, - params?: ReadonlyArray, - ) => Promise<{ rows: Array> }> -} -export interface PgPool extends PgQueryable { - connect: () => Promise -} -export interface PgPoolClient extends PgQueryable { - release: () => void -} - -export interface PostgresDriverOptions { - connectionString?: string - /** Bring-your-own pool (node-postgres `Pool`). */ - client?: PgPool -} - -/** - * Build a driver over a queryable. When `pool` is provided, `transaction` - * acquires a dedicated connection for BEGIN/COMMIT; otherwise (inside a - * transaction's bound client) it runs the callback directly. - */ -function makeDriver(queryable: PgQueryable, pool: PgPool | null): SqlDriver { - const driver: SqlDriver = { - dialect: 'postgres', - async exec(sql, params = []) { - await queryable.query(sql, params) - }, - async query( - sql: string, - params: ReadonlyArray = [], - ) { - const result = await queryable.query(sql, params) - return result.rows as Array - }, - async transaction(fn) { - if (!pool) return fn(driver) - const client = await pool.connect() - try { - await client.query('BEGIN') - const result = await fn(makeDriver(client, null)) - await client.query('COMMIT') - return result - } catch (err) { - await client.query('ROLLBACK') - throw err - } finally { - client.release() - } - }, - } - return driver -} - -/** Build a {@link SqlDriver} backed by Postgres (pool resolved lazily). */ -export function createPostgresDriver(opts: PostgresDriverOptions): SqlDriver { - let poolPromise: Promise | undefined - const getPool = (): Promise => { - if (opts.client) return Promise.resolve(opts.client) - poolPromise ??= (async () => { - const { Pool } = await import('pg') - const pool = new Pool({ connectionString: opts.connectionString }) - // Adapt the real pg Pool to our minimal interface (no cast needed): - // both `query(sql, params) => { rows }` and `connect()` already match. - return { - query: (sql, params) => pool.query(sql, params as Array), - connect: async () => { - const client = await pool.connect() - return { - query: (sql, params) => client.query(sql, params as Array), - release: () => client.release(), - } - }, - } - })() - return poolPromise - } - - return { - dialect: 'postgres', - async exec(sql, params) { - const pool = await getPool() - await makeDriver(pool, pool).exec(sql, params) - }, - async query(sql, params) { - const pool = await getPool() - return makeDriver(pool, pool).query(sql, params) - }, - async transaction(fn) { - const pool = await getPool() - return makeDriver(pool, pool).transaction(fn) - }, - } -} - -export interface PostgresPersistenceOptions extends PostgresDriverOptions { - mode?: PersistenceMode - /** Run migrations on first use (default false). */ - migrate?: boolean -} - -/** Postgres-backed {@link SqlPersistence}. */ -export function postgresPersistence( - opts: PostgresPersistenceOptions, -): SqlPersistence { - const driver = createPostgresDriver({ - connectionString: opts.connectionString, - client: opts.client, - }) - return createSqlPersistence(driver, { - mode: opts.mode, - migrate: opts.migrate, - }) -} diff --git a/packages/ai-persistence-postgres/tests/driver.test.ts b/packages/ai-persistence-postgres/tests/driver.test.ts deleted file mode 100644 index 2ab9d803b..000000000 --- a/packages/ai-persistence-postgres/tests/driver.test.ts +++ /dev/null @@ -1,79 +0,0 @@ -import { describe, expect, it } from 'vitest' -import { createPostgresDriver } from '../src/index' -import type { PgPool, PgPoolClient } from '../src/index' - -/** A fake pg pool recording SQL, returning canned rows, tracking tx control. */ -function fakePool() { - const calls: Array<{ sql: string; params?: ReadonlyArray }> = [] - let nextRows: Array> = [] - let released = false - const client: PgPoolClient = { - query: (sql, params) => { - calls.push({ sql, params }) - return Promise.resolve({ rows: sql.startsWith('SELECT') ? nextRows : [] }) - }, - release: () => { - released = true - }, - } - const pool: PgPool = { - query: (sql, params) => { - calls.push({ sql, params }) - return Promise.resolve({ rows: sql.startsWith('SELECT') ? nextRows : [] }) - }, - connect: () => Promise.resolve(client), - } - return { - pool, - calls, - setRows: (rows: Array>) => { - nextRows = rows - }, - wasReleased: () => released, - } -} - -describe('createPostgresDriver (BYO pool)', () => { - it('forwards sql + params and returns rows', async () => { - const f = fakePool() - f.setRows([{ run_id: 'r1' }]) - const driver = createPostgresDriver({ client: f.pool }) - - await driver.exec('INSERT INTO runs (run_id) VALUES ($1)', ['r1']) - const rows = await driver.query('SELECT * FROM runs WHERE run_id = $1', [ - 'r1', - ]) - - expect(rows).toEqual([{ run_id: 'r1' }]) - expect(f.calls[0]).toEqual({ - sql: 'INSERT INTO runs (run_id) VALUES ($1)', - params: ['r1'], - }) - expect(driver.dialect).toBe('postgres') - }) - - it('wraps a transaction in BEGIN/COMMIT on a dedicated connection', async () => { - const f = fakePool() - const driver = createPostgresDriver({ client: f.pool }) - await driver.transaction(async (tx) => { - await tx.exec('UPDATE runs SET status = $1', ['done']) - }) - const sqls = f.calls.map((c) => c.sql) - expect(sqls[0]).toBe('BEGIN') - expect(sqls).toContain('UPDATE runs SET status = $1') - expect(sqls[sqls.length - 1]).toBe('COMMIT') - expect(f.wasReleased()).toBe(true) - }) - - it('rolls back and rethrows on error', async () => { - const f = fakePool() - const driver = createPostgresDriver({ client: f.pool }) - await expect( - driver.transaction(async () => { - throw new Error('boom') - }), - ).rejects.toThrow('boom') - expect(f.calls.map((c) => c.sql)).toContain('ROLLBACK') - expect(f.wasReleased()).toBe(true) - }) -}) diff --git a/packages/ai-persistence-postgres/tsconfig.json b/packages/ai-persistence-postgres/tsconfig.json deleted file mode 100644 index c38689f4e..000000000 --- a/packages/ai-persistence-postgres/tsconfig.json +++ /dev/null @@ -1,8 +0,0 @@ -{ - "extends": "../../tsconfig.base.json", - "compilerOptions": { - "outDir": "dist" - }, - "include": ["src", "tests"], - "exclude": ["node_modules", "dist"] -} diff --git a/packages/ai-persistence-postgres/vite.config.ts b/packages/ai-persistence-postgres/vite.config.ts deleted file mode 100644 index 11f5b20b7..000000000 --- a/packages/ai-persistence-postgres/vite.config.ts +++ /dev/null @@ -1,37 +0,0 @@ -import { defineConfig, mergeConfig } from 'vitest/config' -import { tanstackViteConfig } from '@tanstack/vite-config' -import packageJson from './package.json' - -const config = defineConfig({ - test: { - name: packageJson.name, - dir: './', - watch: false, - - globals: true, - environment: 'node', - include: ['tests/**/*.test.ts'], - coverage: { - provider: 'v8', - reporter: ['text', 'json', 'html', 'lcov'], - exclude: [ - 'node_modules/', - 'dist/', - 'tests/', - '**/*.test.ts', - '**/*.config.ts', - '**/types.ts', - ], - include: ['src/**/*.ts'], - }, - }, -}) - -export default mergeConfig( - config, - tanstackViteConfig({ - entry: ['./src/index.ts'], - srcDir: './src', - cjs: false, - }), -) diff --git a/packages/ai-persistence-prisma/.gitignore b/packages/ai-persistence-prisma/.gitignore new file mode 100644 index 000000000..e1c30d225 --- /dev/null +++ b/packages/ai-persistence-prisma/.gitignore @@ -0,0 +1,3 @@ +# Scratch sqlite databases used to generate migrations / run tests. +prisma/*.db +prisma/*.db-journal diff --git a/packages/ai-persistence-prisma/package.json b/packages/ai-persistence-prisma/package.json index 1a5e8fd5c..b0c3e1035 100644 --- a/packages/ai-persistence-prisma/package.json +++ b/packages/ai-persistence-prisma/package.json @@ -1,7 +1,7 @@ { "name": "@tanstack/ai-persistence-prisma", "version": "0.1.0", - "description": "Prisma backend for TanStack AI persistence — bring your own PrismaClient; persists via $queryRawUnsafe/$executeRawUnsafe (sqlite or postgres dialect).", + "description": "Prisma backend for TanStack AI state persistence — bring your own PrismaClient. Ships a schema.prisma fragment mirroring @tanstack/ai-persistence-drizzle column-for-column, plus prismaPersistence(prisma) wiring the AIPersistence stores over a PrismaClient.", "author": "", "license": "MIT", "repository": { @@ -12,14 +12,13 @@ "keywords": [ "ai", "tanstack", - "persistence" + "persistence", + "prisma", + "sqlite" ], "type": "module", "module": "./dist/esm/index.js", "types": "./dist/esm/index.d.ts", - "bin": { - "tanstack-ai-persistence-prisma": "./dist/bin/bin.js" - }, "exports": { ".": { "types": "./dist/esm/index.d.ts", @@ -28,11 +27,13 @@ }, "files": [ "dist", - "src" + "src", + "prisma" ], "scripts": { - "build": "vite build && tsup --config tsup.bin.config.ts", + "build": "vite build", "clean": "premove ./build ./dist", + "db:generate": "prisma generate", "lint:fix": "eslint ./src --fix", "test:build": "publint --strict", "test:eslint": "eslint ./src", @@ -40,16 +41,35 @@ "test:lib:dev": "pnpm test:lib --watch", "test:types": "tsc" }, + "nx": { + "targets": { + "db:generate": { + "cache": false + }, + "build": { + "dependsOn": ["db:generate", "^build"] + }, + "test:lib": { + "dependsOn": ["db:generate", "^build"] + }, + "test:types": { + "dependsOn": ["db:generate", "^build"] + }, + "test:eslint": { + "dependsOn": ["db:generate", "^build"] + } + } + }, "peerDependencies": { + "@prisma/client": ">=6", "@tanstack/ai": "workspace:^", - "@tanstack/ai-persistence": "workspace:^", - "@tanstack/ai-persistence-sql": "workspace:^" + "@tanstack/ai-persistence": "workspace:^" }, "devDependencies": { + "@prisma/client": "^6.19.3", "@tanstack/ai": "workspace:*", "@tanstack/ai-persistence": "workspace:*", - "@tanstack/ai-persistence-sql": "workspace:*", "@vitest/coverage-v8": "4.0.14", - "tsup": "^8.5.1" + "prisma": "^6.19.3" } } diff --git a/packages/ai-persistence-prisma/prisma/migrations/20260709120042_init/migration.sql b/packages/ai-persistence-prisma/prisma/migrations/20260709120042_init/migration.sql new file mode 100644 index 000000000..9140976fd --- /dev/null +++ b/packages/ai-persistence-prisma/prisma/migrations/20260709120042_init/migration.sql @@ -0,0 +1,61 @@ +-- CreateTable +CREATE TABLE "messages" ( + "thread_id" TEXT NOT NULL PRIMARY KEY, + "messages_json" TEXT NOT NULL +); + +-- CreateTable +CREATE TABLE "runs" ( + "run_id" TEXT NOT NULL PRIMARY KEY, + "thread_id" TEXT NOT NULL, + "status" TEXT NOT NULL, + "started_at" BIGINT NOT NULL, + "finished_at" BIGINT, + "error" TEXT, + "usage_json" TEXT +); + +-- CreateTable +CREATE TABLE "interrupts" ( + "interrupt_id" TEXT NOT NULL PRIMARY KEY, + "run_id" TEXT NOT NULL, + "thread_id" TEXT NOT NULL, + "status" TEXT NOT NULL, + "requested_at" BIGINT NOT NULL, + "resolved_at" BIGINT, + "payload_json" TEXT NOT NULL, + "response_json" TEXT +); + +-- CreateTable +CREATE TABLE "metadata" ( + "scope" TEXT NOT NULL, + "key" TEXT NOT NULL, + "value_json" TEXT NOT NULL, + + PRIMARY KEY ("scope", "key") +); + +-- CreateTable +CREATE TABLE "artifacts" ( + "artifact_id" TEXT NOT NULL PRIMARY KEY, + "run_id" TEXT NOT NULL, + "thread_id" TEXT NOT NULL, + "name" TEXT NOT NULL, + "mime_type" TEXT NOT NULL, + "size" BIGINT NOT NULL, + "external_url" TEXT, + "created_at" BIGINT NOT NULL +); + +-- CreateTable +CREATE TABLE "blobs" ( + "key" TEXT NOT NULL PRIMARY KEY, + "content_type" TEXT, + "size" BIGINT, + "etag" TEXT, + "custom_metadata_json" TEXT, + "created_at" BIGINT, + "updated_at" BIGINT, + "body" BLOB +); diff --git a/packages/ai-persistence-prisma/prisma/migrations/migration_lock.toml b/packages/ai-persistence-prisma/prisma/migrations/migration_lock.toml new file mode 100644 index 000000000..2a5a44419 --- /dev/null +++ b/packages/ai-persistence-prisma/prisma/migrations/migration_lock.toml @@ -0,0 +1,3 @@ +# Please do not edit this file manually +# It should be added in your version-control system (e.g., Git) +provider = "sqlite" diff --git a/packages/ai-persistence-prisma/prisma/schema.prisma b/packages/ai-persistence-prisma/prisma/schema.prisma new file mode 100644 index 000000000..ae2f3031a --- /dev/null +++ b/packages/ai-persistence-prisma/prisma/schema.prisma @@ -0,0 +1,101 @@ +// Prisma schema for the TanStack AI **state** persistence contract. +// +// Kept in sync with packages/ai-persistence-drizzle/src/schema.ts — see coupling +// persistence-schema-dual-source. The two ORMs describe the SAME state tables and +// have NO auto-converter: any change here MUST be mirrored in the drizzle schema, +// with regenerated migrations for BOTH ORMs and the shared conformance suite +// re-run against memory + drizzle + prisma. +// +// JSON-valued fields are stored in `*_json` TEXT columns (portable, and matching +// drizzle's `text({ mode: 'json' })` on-disk shape); epoch-millisecond timestamps +// and sizes are stored as integers. Prisma's `Json` type is unavailable on +// sqlite, so the store layer serializes/deserializes those columns with +// JSON.stringify/parse. +// +// Integer columns use Prisma `BigInt` (not `Int`): both map to a sqlite +// `INTEGER` column — so the DDL matches drizzle's `integer()` column-for-column +// — but `BigInt` covers the full 64-bit range that epoch-millisecond timestamps +// require (`Int` caps at 32 bits). The store layer converts bigint <-> number at +// the boundary so records stay plain `number`s, matching the drizzle backend. + +generator client { + provider = "prisma-client-js" +} + +datasource db { + provider = "sqlite" + url = env("DATABASE_URL") +} + +/// Thread message history (`MessageStore`). +model Message { + threadId String @id @map("thread_id") + messagesJson String @map("messages_json") + + @@map("messages") +} + +/// Run lifecycle records (`RunStore`). +model Run { + runId String @id @map("run_id") + threadId String @map("thread_id") + status String + startedAt BigInt @map("started_at") + finishedAt BigInt? @map("finished_at") + error String? + usageJson String? @map("usage_json") + + @@map("runs") +} + +/// Interrupt / approval records (`InterruptStore`). +model Interrupt { + interruptId String @id @map("interrupt_id") + runId String @map("run_id") + threadId String @map("thread_id") + status String + requestedAt BigInt @map("requested_at") + resolvedAt BigInt? @map("resolved_at") + payloadJson String @map("payload_json") + responseJson String? @map("response_json") + + @@map("interrupts") +} + +/// Scoped key/value metadata (`MetadataStore`). +model Metadata { + scope String + key String + valueJson String @map("value_json") + + @@id([scope, key]) + @@map("metadata") +} + +/// Generation artifact references (`ArtifactStore`). Bytes live in the blob store. +model Artifact { + artifactId String @id @map("artifact_id") + runId String @map("run_id") + threadId String @map("thread_id") + name String + mimeType String @map("mime_type") + size BigInt + externalUrl String? @map("external_url") + createdAt BigInt @map("created_at") + + @@map("artifacts") +} + +/// Generic blob objects (`BlobStore`). +model Blob { + key String @id + contentType String? @map("content_type") + size BigInt? + etag String? + customMetadataJson String? @map("custom_metadata_json") + createdAt BigInt? @map("created_at") + updatedAt BigInt? @map("updated_at") + body Bytes? + + @@map("blobs") +} diff --git a/packages/ai-persistence-prisma/src/bin.ts b/packages/ai-persistence-prisma/src/bin.ts deleted file mode 100644 index 447ca521a..000000000 --- a/packages/ai-persistence-prisma/src/bin.ts +++ /dev/null @@ -1,17 +0,0 @@ -import { runCli } from './cli' - -try { - const result = await runCli(process.argv.slice(2), { - writeStdout: (value) => process.stdout.write(value), - }) - if (result.kind === 'help') { - process.stdout.write(result.text) - } else if (result.kind === 'file') { - process.stdout.write(`Wrote ${result.path}\n`) - } -} catch (error) { - process.stderr.write( - `${error instanceof Error ? error.message : String(error)}\n`, - ) - process.exitCode = 1 -} diff --git a/packages/ai-persistence-prisma/src/cli-core.ts b/packages/ai-persistence-prisma/src/cli-core.ts deleted file mode 100644 index 47d48e6bf..000000000 --- a/packages/ai-persistence-prisma/src/cli-core.ts +++ /dev/null @@ -1,31 +0,0 @@ -import { join } from 'node:path' -import { - formatMigrationSql, - runMigrationCli as runSharedMigrationCli, -} from '@tanstack/ai-persistence-sql/cli-core' -import type { - CliOptions, - CliResult, - MigrationCliConfig, -} from '@tanstack/ai-persistence-sql/cli-core' - -export type CliDialect = 'sqlite' | 'postgres' -export type { CliOptions, CliResult } - -const prismaMigrationCliConfig: MigrationCliConfig = { - commandName: 'tanstack-ai-persistence-prisma', - defaultDialect: 'sqlite', - allowedDialects: ['sqlite', 'postgres'], - defaultName: 'tanstack_ai_persistence', - defaultOut: ({ cwd, timestamp, name }) => - join(cwd, 'prisma', 'migrations', `${timestamp}_${name}`, 'migration.sql'), -} - -export { formatMigrationSql } - -export function runMigrationCli( - argv: Array, - options?: CliOptions, -): Promise { - return runSharedMigrationCli(argv, prismaMigrationCliConfig, options) -} diff --git a/packages/ai-persistence-prisma/src/cli.ts b/packages/ai-persistence-prisma/src/cli.ts deleted file mode 100644 index 6c4518c11..000000000 --- a/packages/ai-persistence-prisma/src/cli.ts +++ /dev/null @@ -1,9 +0,0 @@ -import { runMigrationCli } from './cli-core' -import type { CliOptions, CliResult } from './cli-core' - -export function runCli( - argv: Array, - options?: CliOptions, -): Promise { - return runMigrationCli(argv, options) -} diff --git a/packages/ai-persistence-prisma/src/index.ts b/packages/ai-persistence-prisma/src/index.ts index 8eb634840..09ab56573 100644 --- a/packages/ai-persistence-prisma/src/index.ts +++ b/packages/ai-persistence-prisma/src/index.ts @@ -1,70 +1,44 @@ /** - * Prisma backend (bring-your-own). Persists over a `PrismaClient`'s raw SQL - * escape hatches — `$queryRawUnsafe(sql, ...params)` for reads and - * `$executeRawUnsafe(sql, ...params)` for writes — which take a positional- - * parameter SQL string, exactly matching the shared `SqlDriver` contract. + * Prisma-backed **state** persistence for TanStack AI. * - * The tables are the same as the raw SQL backend; add the documented Prisma - * models to your `schema.prisma` and run `prisma migrate`, or opt in to lazy - * backend migrations with `migrate: true`. + * Bring your own `PrismaClient`: add the model fragment shipped in this + * package's `prisma/schema.prisma` to your own Prisma schema, run + * `prisma migrate` through your normal workflow, then pass the generated + * client to {@link prismaPersistence}. This wires the `AIPersistence` stores + * (messages, runs, interrupts, metadata, artifacts, blobs) over it. + * + * The Prisma schema mirrors `@tanstack/ai-persistence-drizzle`'s exported + * schema column-for-column — see coupling `persistence-schema-dual-source`. + * Any change to one schema MUST be mirrored in the other, with regenerated + * migrations for both ORMs and the shared conformance suite re-run. + * + * Locks are not part of the SQL schema — an in-memory lock is provided as a dev + * default so the returned persistence is complete. Swap in a distributed lock + * for multi-process deployments. */ -import { createSqlPersistence } from '@tanstack/ai-persistence-sql' -import type { - Dialect, - SqlDriver, - SqlPersistence, - SqlRow, -} from '@tanstack/ai-persistence-sql' -import type { PersistenceMode } from '@tanstack/ai-persistence' +import { InMemoryLockStore } from '@tanstack/ai' +import { + createArtifactStore, + createBlobStore, + createInterruptStore, + createMessageStore, + createMetadataStore, + createRunStore, +} from './stores' +import type { PrismaClient } from '@prisma/client' +import type { AIPersistence } from '@tanstack/ai-persistence' -/** The PrismaClient surface this adapter relies on. */ -export interface PrismaRawClient { - $queryRawUnsafe: ( - sql: string, - ...params: Array - ) => Promise - $executeRawUnsafe: (sql: string, ...params: Array) => Promise - $transaction: (fn: (tx: PrismaRawClient) => Promise) => Promise -} - -/** Build a {@link SqlDriver} over a PrismaClient. */ -export function createPrismaDriver( - prisma: PrismaRawClient, - dialect: Dialect, -): SqlDriver { - const driver: SqlDriver = { - dialect, - async exec(sql, params = []) { - await prisma.$executeRawUnsafe(sql, ...params) - }, - async query( - sql: string, - params: ReadonlyArray = [], - ) { - return prisma.$queryRawUnsafe>(sql, ...params) - }, - async transaction(fn) { - return prisma.$transaction((tx) => fn(createPrismaDriver(tx, dialect))) +/** Wire the AIPersistence stores over an existing Prisma `PrismaClient`. */ +export function prismaPersistence(prisma: PrismaClient): AIPersistence { + return { + stores: { + messages: createMessageStore(prisma), + runs: createRunStore(prisma), + interrupts: createInterruptStore(prisma), + metadata: createMetadataStore(prisma), + artifacts: createArtifactStore(prisma), + blobs: createBlobStore(prisma), + locks: new InMemoryLockStore(), }, } - return driver -} - -export interface PrismaPersistenceOptions { - prisma: PrismaRawClient - dialect: Dialect - mode?: PersistenceMode - /** Run migrations on first use (default false). */ - migrate?: boolean -} - -/** Prisma-backed {@link SqlPersistence}. */ -export function prismaPersistence( - opts: PrismaPersistenceOptions, -): SqlPersistence { - const driver = createPrismaDriver(opts.prisma, opts.dialect) - return createSqlPersistence(driver, { - mode: opts.mode, - migrate: opts.migrate, - }) } diff --git a/packages/ai-persistence-prisma/src/stores.ts b/packages/ai-persistence-prisma/src/stores.ts new file mode 100644 index 000000000..de597f0a1 --- /dev/null +++ b/packages/ai-persistence-prisma/src/stores.ts @@ -0,0 +1,492 @@ +/** + * AIPersistence store implementations over a Prisma `PrismaClient`. + * + * Each method mirrors the semantics of the reference in-memory backend + * (`@tanstack/ai-persistence`'s `memory.ts`) and the sibling Drizzle backend + * (`@tanstack/ai-persistence-drizzle`'s `stores.ts`). JSON-valued columns are + * TEXT (Prisma's `Json` type is unavailable on sqlite), so they are serialized + * with `JSON.stringify`/`JSON.parse` here; blob bytes use Prisma's `Bytes`. + */ +import type { PrismaClient } from '@prisma/client' +import type { ModelMessage } from '@tanstack/ai' +import type { + ArtifactRecord, + ArtifactStore, + BlobBody, + BlobListOptions, + BlobListPage, + BlobObject, + BlobRecord, + BlobStore, + InterruptRecord, + InterruptStore, + MessageStore, + MetadataStore, + RunRecord, + RunStatus, + RunStore, +} from '@tanstack/ai-persistence' + +const textEncoder = new TextEncoder() +const textDecoder = new TextDecoder() + +/** + * Smallest string strictly greater than every string that starts with `prefix`, + * used as the exclusive upper bound of a literal prefix range scan. Increments + * the last non-`U+FFFF` UTF-16 code unit (carrying over trailing `U+FFFF`s). + * Returns `undefined` when no finite bound exists — an empty prefix, or a prefix + * consisting solely of `U+FFFF` — in which case the caller omits the upper bound. + */ +function prefixUpperBound(prefix: string): string | undefined { + let i = prefix.length - 1 + while (i >= 0 && prefix.charCodeAt(i) === 0xffff) i-- + if (i < 0) return undefined + return prefix.slice(0, i) + String.fromCharCode(prefix.charCodeAt(i) + 1) +} + +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.') +} + +export function createMessageStore(prisma: PrismaClient): MessageStore { + return { + async loadThread(threadId) { + const row = await prisma.message.findUnique({ where: { threadId } }) + if (!row) return [] + return JSON.parse(row.messagesJson) as Array + }, + async saveThread(threadId, msgs: Array) { + const messagesJson = JSON.stringify(msgs) + await prisma.message.upsert({ + where: { threadId }, + create: { threadId, messagesJson }, + update: { messagesJson }, + }) + }, + } +} + +interface RunRow { + runId: string + threadId: string + status: string + startedAt: bigint + finishedAt: bigint | null + error: string | null + usageJson: string | null +} + +function mapRun(row: RunRow): RunRecord { + return { + runId: row.runId, + threadId: row.threadId, + status: row.status as RunStatus, + startedAt: Number(row.startedAt), + ...(row.finishedAt != null ? { finishedAt: Number(row.finishedAt) } : {}), + ...(row.error != null ? { error: row.error } : {}), + ...(row.usageJson != null + ? { usage: JSON.parse(row.usageJson) as RunRecord['usage'] } + : {}), + } +} + +export function createRunStore(prisma: PrismaClient): RunStore { + const store: RunStore = { + async createOrResume(input) { + const existing = await store.get(input.runId) + if (existing) return existing + const record: RunRecord = { + runId: input.runId, + threadId: input.threadId, + status: input.status ?? 'running', + startedAt: input.startedAt, + } + // Upsert with an empty update = insert-if-absent (never overwrites an + // existing run), matching the reference createOrResume semantics. + await prisma.run.upsert({ + where: { runId: record.runId }, + create: { + runId: record.runId, + threadId: record.threadId, + status: record.status, + startedAt: BigInt(record.startedAt), + }, + update: {}, + }) + return (await store.get(input.runId)) ?? record + }, + async update(runId, patch) { + const data: { + status?: RunStatus + finishedAt?: bigint + error?: string + usageJson?: string + } = {} + if (patch.status !== undefined) data.status = patch.status + if (patch.finishedAt !== undefined) + data.finishedAt = BigInt(patch.finishedAt) + if (patch.error !== undefined) data.error = patch.error + if (patch.usage !== undefined) data.usageJson = JSON.stringify(patch.usage) + if (Object.keys(data).length === 0) return + // updateMany no-ops (does not throw) when the run does not exist. + await prisma.run.updateMany({ where: { runId }, data }) + }, + async get(runId) { + const row = await prisma.run.findUnique({ where: { runId } }) + return row ? mapRun(row) : null + }, + } + return store +} + +interface InterruptRow { + interruptId: string + runId: string + threadId: string + status: string + requestedAt: bigint + resolvedAt: bigint | null + payloadJson: string + responseJson: string | null +} + +function mapInterrupt(row: InterruptRow): InterruptRecord { + return { + interruptId: row.interruptId, + runId: row.runId, + threadId: row.threadId, + status: row.status as InterruptRecord['status'], + requestedAt: Number(row.requestedAt), + ...(row.resolvedAt != null ? { resolvedAt: Number(row.resolvedAt) } : {}), + payload: JSON.parse(row.payloadJson) as Record, + ...(row.responseJson != null + ? { response: JSON.parse(row.responseJson) as unknown } + : {}), + } +} + +export function createInterruptStore(prisma: PrismaClient): InterruptStore { + return { + async create(record) { + await prisma.interrupt.upsert({ + where: { interruptId: record.interruptId }, + create: { + interruptId: record.interruptId, + runId: record.runId, + threadId: record.threadId, + status: record.status, + requestedAt: BigInt(record.requestedAt), + payloadJson: JSON.stringify(record.payload), + responseJson: + record.response == null ? null : JSON.stringify(record.response), + }, + update: {}, + }) + }, + async resolve(interruptId, response) { + await prisma.interrupt.updateMany({ + where: { interruptId }, + data: { + status: 'resolved', + resolvedAt: BigInt(Date.now()), + responseJson: response == null ? null : JSON.stringify(response), + }, + }) + }, + async cancel(interruptId) { + await prisma.interrupt.updateMany({ + where: { interruptId }, + data: { status: 'cancelled', resolvedAt: BigInt(Date.now()) }, + }) + }, + async get(interruptId) { + const row = await prisma.interrupt.findUnique({ where: { interruptId } }) + return row ? mapInterrupt(row) : null + }, + async list(threadId) { + const rows = await prisma.interrupt.findMany({ + where: { threadId }, + orderBy: { requestedAt: 'asc' }, + }) + return rows.map(mapInterrupt) + }, + async listPending(threadId) { + const rows = await prisma.interrupt.findMany({ + where: { threadId, status: 'pending' }, + orderBy: { requestedAt: 'asc' }, + }) + return rows.map(mapInterrupt) + }, + async listByRun(runId) { + const rows = await prisma.interrupt.findMany({ + where: { runId }, + orderBy: { requestedAt: 'asc' }, + }) + return rows.map(mapInterrupt) + }, + async listPendingByRun(runId) { + const rows = await prisma.interrupt.findMany({ + where: { runId, status: 'pending' }, + orderBy: { requestedAt: 'asc' }, + }) + return rows.map(mapInterrupt) + }, + } +} + +export function createMetadataStore(prisma: PrismaClient): MetadataStore { + return { + async get(scope, key) { + const row = await prisma.metadata.findUnique({ + where: { scope_key: { scope, key } }, + }) + return row ? (JSON.parse(row.valueJson) as unknown) : null + }, + async set(scope, key, value) { + const valueJson = JSON.stringify(value) + await prisma.metadata.upsert({ + where: { scope_key: { scope, key } }, + create: { scope, key, valueJson }, + update: { valueJson }, + }) + }, + async delete(scope, key) { + await prisma.metadata.deleteMany({ where: { scope, key } }) + }, + } +} + +interface ArtifactRow { + artifactId: string + runId: string + threadId: string + name: string + mimeType: string + size: bigint + externalUrl: string | null + createdAt: bigint +} + +function mapArtifact(row: ArtifactRow): ArtifactRecord { + return { + artifactId: row.artifactId, + runId: row.runId, + threadId: row.threadId, + name: row.name, + mimeType: row.mimeType, + size: Number(row.size), + ...(row.externalUrl != null ? { externalUrl: row.externalUrl } : {}), + createdAt: Number(row.createdAt), + } +} + +export function createArtifactStore(prisma: PrismaClient): ArtifactStore { + return { + async save(record) { + const data = { + runId: record.runId, + threadId: record.threadId, + name: record.name, + mimeType: record.mimeType, + size: BigInt(record.size), + externalUrl: record.externalUrl ?? null, + createdAt: BigInt(record.createdAt), + } + await prisma.artifact.upsert({ + where: { artifactId: record.artifactId }, + create: { artifactId: record.artifactId, ...data }, + update: data, + }) + }, + async get(artifactId) { + const row = await prisma.artifact.findUnique({ where: { artifactId } }) + return row ? mapArtifact(row) : null + }, + async list(runId) { + const rows = await prisma.artifact.findMany({ + where: { runId }, + orderBy: [{ createdAt: 'asc' }, { artifactId: 'asc' }], + }) + return rows.map(mapArtifact) + }, + async delete(artifactId) { + await prisma.artifact.deleteMany({ where: { artifactId } }) + }, + async deleteForRun(runId) { + await prisma.artifact.deleteMany({ where: { runId } }) + }, + } +} + +interface BlobRow { + key: string + contentType: string | null + size: bigint | null + etag: string | null + customMetadataJson: string | null + createdAt: bigint | null + updatedAt: bigint | null + body: Uint8Array | null +} + +function blobRecordSnapshot(row: BlobRow): BlobRecord { + return { + key: row.key, + ...(row.size != null ? { size: Number(row.size) } : {}), + ...(row.etag != null ? { etag: row.etag } : {}), + ...(row.contentType != null ? { contentType: row.contentType } : {}), + ...(row.customMetadataJson != null + ? { + customMetadata: JSON.parse(row.customMetadataJson) as Record< + string, + string + >, + } + : {}), + ...(row.createdAt != null ? { createdAt: Number(row.createdAt) } : {}), + ...(row.updatedAt != null ? { updatedAt: Number(row.updatedAt) } : {}), + } +} + +function blobObject(row: BlobRow): BlobObject { + const bytes = row.body ? copyBytes(new Uint8Array(row.body)) : new Uint8Array() + return { + ...blobRecordSnapshot(row), + body: new ReadableStream({ + start(controller) { + controller.enqueue(copyBytes(bytes)) + controller.close() + }, + }), + arrayBuffer: () => Promise.resolve(bytesToArrayBuffer(bytes)), + text: () => Promise.resolve(textDecoder.decode(bytes)), + } +} + +export function createBlobStore(prisma: PrismaClient): BlobStore { + const readRow = (key: string): Promise => + prisma.blob.findUnique({ where: { key } }) + return { + async put(key, body, options) { + const bytes = await bytesFromBlobBody(body) + const existing = await readRow(key) + const now = Date.now() + const contentType = + options?.contentType ?? + (typeof Blob !== 'undefined' && body instanceof Blob + ? body.type || undefined + : undefined) + const customMetadata = options?.customMetadata + ? { ...options.customMetadata } + : undefined + const record = { + contentType: contentType ?? null, + size: BigInt(bytes.byteLength), + etag: crypto.randomUUID(), + customMetadataJson: customMetadata + ? JSON.stringify(customMetadata) + : null, + createdAt: existing?.createdAt ?? BigInt(now), + updatedAt: BigInt(now), + body: Buffer.from(bytes), + } + const stored = await prisma.blob.upsert({ + where: { key }, + create: { key, ...record }, + update: record, + }) + return blobRecordSnapshot(stored) + }, + async get(key) { + const row = await readRow(key) + return row ? blobObject(row) : null + }, + async head(key) { + const row = await readRow(key) + return row ? blobRecordSnapshot(row) : null + }, + async delete(key) { + await prisma.blob.deleteMany({ where: { key } }) + }, + async list(options?: BlobListOptions): Promise { + const limit = options?.limit + if (limit === 0) return { objects: [], truncated: false } + const prefix = options?.prefix ?? '' + // Match the prefix LITERALLY and case-sensitively to mirror the reference + // in-memory backend's `key.startsWith(prefix)`. Prisma's `startsWith` + // compiles to a SQL `LIKE` on sqlite, which treats `_`/`%` in the prefix + // as wildcards and matches case-insensitively for ASCII. Instead use a + // half-open range on the (BINARY-collated) key column — `key >= prefix AND + // key < upperBound` — which contains no LIKE metacharacters and relies on + // binary/byte ordering, giving literal, case-sensitive matching. + const upperBound = prefixUpperBound(prefix) + const rows = await prisma.blob.findMany({ + where: { + key: { + gte: prefix, + ...(upperBound !== undefined ? { lt: upperBound } : {}), + ...(options?.cursor !== undefined ? { gt: options.cursor } : {}), + }, + }, + orderBy: { key: 'asc' }, + ...(limit === undefined ? {} : { take: limit + 1 }), + }) + const pageRows = limit === undefined ? rows : rows.slice(0, limit) + const objects = pageRows.map(blobRecordSnapshot) + const truncated = limit !== undefined && rows.length > limit + return { + objects, + ...(truncated ? { cursor: pageRows.at(-1)?.key, truncated } : {}), + } + }, + } +} diff --git a/packages/ai-persistence-prisma/tests/cli.test.ts b/packages/ai-persistence-prisma/tests/cli.test.ts deleted file mode 100644 index 573ff9e96..000000000 --- a/packages/ai-persistence-prisma/tests/cli.test.ts +++ /dev/null @@ -1,91 +0,0 @@ -import { mkdtemp, readFile, writeFile } from 'node:fs/promises' -import { join } from 'node:path' -import { tmpdir } from 'node:os' -import { describe, expect, it } from 'vitest' -import { runCli } from '../src/cli' - -describe('Prisma persistence CLI', () => { - it('writes the default Prisma migration path', async () => { - const cwd = await mkdtemp(join(tmpdir(), 'tanstack-ai-prisma-')) - - const result = await runCli( - ['--dialect', 'postgres', '--timestamp', '20260706123456'], - { cwd }, - ) - - expect(result).toEqual({ - kind: 'file', - path: join( - cwd, - 'prisma', - 'migrations', - '20260706123456_tanstack_ai_persistence', - 'migration.sql', - ), - }) - if (result.kind === 'file') { - await expect(readFile(result.path, 'utf8')).resolves.toContain( - 'TanStack AI persistence migration for postgres', - ) - } - }) - - it('prints SQL to stdout without writing a file', async () => { - const cwd = await mkdtemp(join(tmpdir(), 'tanstack-ai-prisma-')) - - const result = await runCli(['--dialect', 'sqlite', '--stdout'], { cwd }) - - expect(result.kind).toBe('stdout') - if (result.kind === 'stdout') { - expect(result.sql).toContain( - 'TanStack AI persistence migration for sqlite', - ) - expect(result.sql).toContain('CREATE TABLE IF NOT EXISTS runs') - } - }) - - it('rejects invalid dialects', async () => { - await expect(runCli(['--dialect', 'mysql'])).rejects.toThrow( - 'Invalid --dialect', - ) - }) - - it('rejects path traversal migration names', async () => { - await expect(runCli(['--name', '../escape'])).rejects.toThrow( - 'Invalid --name', - ) - }) - - it('refuses to overwrite an existing file without --force', async () => { - const cwd = await mkdtemp(join(tmpdir(), 'tanstack-ai-prisma-')) - const out = join(cwd, 'migration.sql') - await writeFile(out, 'existing') - - await expect(runCli(['--dialect', 'sqlite', '--out', out])).rejects.toThrow( - 'already exists', - ) - }) - - it('overwrites an existing file with --force', async () => { - const cwd = await mkdtemp(join(tmpdir(), 'tanstack-ai-prisma-')) - const out = join(cwd, 'migration.sql') - await writeFile(out, 'existing') - - const result = await runCli([ - '--dialect', - 'sqlite', - '--out', - out, - '--force', - ]) - - expect(result).toEqual({ kind: 'file', path: out }) - await expect(readFile(out, 'utf8')).resolves.toContain( - 'TanStack AI persistence migration for sqlite', - ) - }) - - it('rejects missing option values', async () => { - await expect(runCli(['--out'])).rejects.toThrow('Missing value for --out') - }) -}) diff --git a/packages/ai-persistence-prisma/tests/prisma.conformance.test.ts b/packages/ai-persistence-prisma/tests/prisma.conformance.test.ts new file mode 100644 index 000000000..a06cb8848 --- /dev/null +++ b/packages/ai-persistence-prisma/tests/prisma.conformance.test.ts @@ -0,0 +1,59 @@ +import { readFile, readdir } from 'node:fs/promises' +import { mkdtempSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { fileURLToPath } from 'node:url' +import { afterAll } from 'vitest' +import { PrismaClient } from '@prisma/client' +import { runPersistenceConformance } from '@tanstack/ai-persistence/testkit' +import { prismaPersistence } from '../src/index' + +const migrationsDir = fileURLToPath( + new URL('../prisma/migrations', import.meta.url), +) + +const clients: Array = [] + +/** Read every generated migration's DDL, in journal order, as raw statements. */ +async function migrationStatements(): Promise> { + const entries = await readdir(migrationsDir, { withFileTypes: true }) + const dirs = entries + .filter((entry) => entry.isDirectory()) + .map((entry) => entry.name) + .sort() + const statements: Array = [] + for (const dir of dirs) { + const sql = await readFile(join(migrationsDir, dir, 'migration.sql'), 'utf8') + for (const raw of sql.split(';')) { + const statement = raw + .split('\n') + .filter((line) => !line.trimStart().startsWith('--')) + .join('\n') + .trim() + if (statement) statements.push(statement) + } + } + return statements +} + +/** Create a PrismaClient over a fresh temp sqlite file with migrations applied. */ +async function makeTestClient(): Promise { + const dir = mkdtempSync(join(tmpdir(), 'tanstack-ai-prisma-')) + const dbPath = join(dir, 'state.db').replace(/\\/g, '/') + const prisma = new PrismaClient({ + datasources: { db: { url: `file:${dbPath}` } }, + }) + clients.push(prisma) + for (const statement of await migrationStatements()) { + await prisma.$executeRawUnsafe(statement) + } + return prisma +} + +afterAll(async () => { + await Promise.all(clients.map((client) => client.$disconnect())) +}) + +runPersistenceConformance('prisma', async () => + prismaPersistence(await makeTestClient()), +) diff --git a/packages/ai-persistence-prisma/tests/prisma.test.ts b/packages/ai-persistence-prisma/tests/prisma.test.ts deleted file mode 100644 index cea5c6767..000000000 --- a/packages/ai-persistence-prisma/tests/prisma.test.ts +++ /dev/null @@ -1,55 +0,0 @@ -import { describe, expect, it } from 'vitest' -import { DatabaseSync } from 'node:sqlite' -import { EventType } from '@tanstack/ai' -import type { StreamChunk } from '@tanstack/ai' -import { prismaPersistence } from '../src/index' -import type { PrismaRawClient } from '../src/index' - -/** A fake PrismaClient backed by node:sqlite (raw escape hatches only). */ -function fakePrisma(): PrismaRawClient { - const db = new DatabaseSync(':memory:') - const client: PrismaRawClient = { - $queryRawUnsafe: (sql: string, ...params: Array) => - Promise.resolve(db.prepare(sql).all(...(params as Array)) as T), - $executeRawUnsafe: (sql: string, ...params: Array) => { - db.prepare(sql).run(...(params as Array)) - return Promise.resolve(0) - }, - $transaction: (fn) => fn(client), - } - return client -} - -const text = (delta: string): StreamChunk => ({ - type: EventType.TEXT_MESSAGE_CONTENT, - messageId: 'm1', - delta, - timestamp: 1, -}) - -describe('prismaPersistence (sqlite dialect, raw escape hatches)', () => { - it('does not create schema by default', async () => { - const p = prismaPersistence({ prisma: fakePrisma(), dialect: 'sqlite' }) - - await expect( - p.runs!.createOrResume({ runId: 'r1', threadId: 't1', startedAt: 1 }), - ).rejects.toThrow(/no such table: runs/) - }) - - it('persists via $queryRawUnsafe / $executeRawUnsafe', async () => { - const p = prismaPersistence({ - prisma: fakePrisma(), - dialect: 'sqlite', - migrate: true, - }) - await p.runs!.createOrResume({ runId: 'r1', threadId: 't1', startedAt: 1 }) - await p.events!.append('r1', 1, text('a')) - - expect((await p.runs!.get('r1'))?.status).toBe('running') - const deltas: Array = [] - for await (const e of p.events!.read('r1')) { - if (e.event.type === 'TEXT_MESSAGE_CONTENT') deltas.push(e.event.delta) - } - expect(deltas).toEqual(['a']) - }) -}) diff --git a/packages/ai-persistence-prisma/tsup.bin.config.ts b/packages/ai-persistence-prisma/tsup.bin.config.ts deleted file mode 100644 index d40ae7297..000000000 --- a/packages/ai-persistence-prisma/tsup.bin.config.ts +++ /dev/null @@ -1,11 +0,0 @@ -import { defineConfig } from 'tsup' - -export default defineConfig({ - entry: { bin: 'src/bin.ts' }, - outDir: 'dist/bin', - format: ['esm'], - platform: 'node', - target: 'node18', - external: ['@tanstack/ai-persistence-sql'], - banner: { js: '#!/usr/bin/env node' }, -}) diff --git a/packages/ai-persistence-sql/package.json b/packages/ai-persistence-sql/package.json deleted file mode 100644 index 07a74434b..000000000 --- a/packages/ai-persistence-sql/package.json +++ /dev/null @@ -1,64 +0,0 @@ -{ - "name": "@tanstack/ai-persistence-sql", - "version": "0.1.0", - "description": "Shared SQL store core for TanStack AI persistence — durable runs, messages, public/internal events, interrupts, and metadata behind a minimal SqlDriver (sqlite | postgres | mysql dialect), with versioned migrations. Consumed by the sqlite/postgres/cloudflare/drizzle/prisma backends.", - "author": "", - "license": "MIT", - "repository": { - "type": "git", - "url": "git+https://github.com/TanStack/ai.git", - "directory": "packages/ai-persistence-sql" - }, - "keywords": [ - "ai", - "ai-sdk", - "typescript", - "tanstack", - "persistence", - "sql", - "sqlite", - "postgres", - "mysql", - "migrations" - ], - "type": "module", - "module": "./dist/esm/index.js", - "types": "./dist/esm/index.d.ts", - "bin": { - "tanstack-ai-persistence-sql": "./dist/bin/bin.js" - }, - "exports": { - ".": { - "types": "./dist/esm/index.d.ts", - "import": "./dist/esm/index.js" - }, - "./cli-core": { - "types": "./dist/esm/cli-core.d.ts", - "import": "./dist/esm/cli-core.js" - } - }, - "files": [ - "dist", - "src" - ], - "scripts": { - "build": "vite build && tsup --config tsup.bin.config.ts", - "clean": "premove ./build ./dist", - "lint:fix": "eslint ./src --fix", - "test:build": "publint --strict", - "test:eslint": "eslint ./src", - "test:lib": "vitest", - "test:lib:dev": "pnpm test:lib --watch", - "test:types": "tsc" - }, - "peerDependencies": { - "@tanstack/ai": "workspace:^", - "@tanstack/ai-persistence": "workspace:^" - }, - "devDependencies": { - "@tanstack/ai": "workspace:*", - "@tanstack/ai-persistence": "workspace:*", - "@vitest/coverage-v8": "4.0.14", - "tsup": "^8.5.1" - } -} diff --git a/packages/ai-persistence-sql/src/bin.ts b/packages/ai-persistence-sql/src/bin.ts deleted file mode 100644 index 447ca521a..000000000 --- a/packages/ai-persistence-sql/src/bin.ts +++ /dev/null @@ -1,17 +0,0 @@ -import { runCli } from './cli' - -try { - const result = await runCli(process.argv.slice(2), { - writeStdout: (value) => process.stdout.write(value), - }) - if (result.kind === 'help') { - process.stdout.write(result.text) - } else if (result.kind === 'file') { - process.stdout.write(`Wrote ${result.path}\n`) - } -} catch (error) { - process.stderr.write( - `${error instanceof Error ? error.message : String(error)}\n`, - ) - process.exitCode = 1 -} diff --git a/packages/ai-persistence-sql/src/cli-core.ts b/packages/ai-persistence-sql/src/cli-core.ts deleted file mode 100644 index 1a7e2ffd1..000000000 --- a/packages/ai-persistence-sql/src/cli-core.ts +++ /dev/null @@ -1,250 +0,0 @@ -import { mkdir, writeFile } from 'node:fs/promises' -import { dirname, isAbsolute, join } from 'node:path' -import { ddl } from './migrations' -import type { Dialect } from './driver' - -export type CliDialect = Dialect - -export interface CliOptions { - cwd?: string - now?: Date - writeStdout?: (value: string) => void -} - -export type CliResult = - | { kind: 'file'; path: string } - | { kind: 'stdout'; sql: string } - | { kind: 'help'; text: string } - -export interface MigrationCliConfig { - commandName: string - defaultDialect: TDialect - allowedDialects: ReadonlyArray - defaultName?: string - defaultOut: (params: { - cwd: string - timestamp: string - name: string - }) => string -} - -interface ParsedArgs { - dialect: TDialect - out?: string - stdout: boolean - force: boolean - help: boolean - name: string - timestamp?: string -} - -export const sqlMigrationCliConfig: MigrationCliConfig = { - commandName: 'tanstack-ai-persistence-sql', - defaultDialect: 'sqlite', - allowedDialects: ['sqlite', 'postgres', 'mysql'], - defaultName: 'tanstack_ai_persistence', - defaultOut: ({ cwd, timestamp, name }) => - join(cwd, 'migrations', `${timestamp}_${name}.sql`), -} - -export function formatMigrationSql(dialect: Dialect): string { - return ( - [ - `-- TanStack AI persistence migration for ${dialect}`, - ...ddl(dialect).map((statement) => `${statement};`), - ].join('\n\n') + '\n' - ) -} - -export async function runMigrationCli( - argv: Array, - config: MigrationCliConfig, - options: CliOptions = {}, -): Promise { - const parsed = parseArgs(argv, config) - const help = helpText(config) - if (parsed.help) return { kind: 'help', text: help } - - const cwd = options.cwd ?? process.cwd() - const timestamp = - parsed.timestamp ?? defaultTimestamp(options.now ?? new Date()) - const sql = formatMigrationSql(parsed.dialect) - - if (parsed.stdout) { - options.writeStdout?.(sql) - return { kind: 'stdout', sql } - } - - const out = parsed.out - ? resolveOut(cwd, parsed.out) - : config.defaultOut({ cwd, timestamp, name: parsed.name }) - - await writeMigrationFile(out, sql, parsed.force) - return { kind: 'file', path: out } -} - -function parseArgs( - argv: Array, - config: MigrationCliConfig, -): ParsedArgs { - const parsed: ParsedArgs = { - dialect: config.defaultDialect, - stdout: false, - force: false, - help: false, - name: config.defaultName ?? 'tanstack_ai_persistence', - } - - for (let i = 0; i < argv.length; i++) { - const arg = argv[i] - if (arg === undefined) { - throw new Error(`Missing option at index ${i}.`) - } - - switch (arg) { - case '--dialect': - parsed.dialect = readDialect(argv, ++i, config) - break - case '--out': - parsed.out = readValue(argv, ++i, '--out') - break - case '--stdout': - parsed.stdout = true - break - case '--force': - parsed.force = true - break - case '--name': - parsed.name = readMigrationName(argv, ++i) - break - case '--timestamp': { - const timestamp = readValue(argv, ++i, '--timestamp') - if (!/^\d{14}$/.test(timestamp)) { - throw new Error('Invalid --timestamp. Expected yyyymmddhhmmss.') - } - parsed.timestamp = timestamp - break - } - case '--help': - case '-h': - parsed.help = true - break - default: - throw new Error(`Unknown option: ${arg}`) - } - } - - return parsed -} - -async function writeMigrationFile( - out: string, - sql: string, - force: boolean, -): Promise { - try { - await writeFile(out, sql, { flag: force ? 'w' : 'wx' }) - } catch (error) { - if (hasCode(error, 'ENOENT')) { - await mkdir(dirname(out), { recursive: true }) - await writeFile(out, sql, { flag: force ? 'w' : 'wx' }) - return - } - if (hasCode(error, 'EEXIST')) { - throw new Error( - `Output file already exists: ${out}. Pass --force to overwrite.`, - ) - } - throw error - } -} - -function hasCode(error: unknown, code: string): boolean { - return Boolean( - error && - typeof error === 'object' && - 'code' in error && - error.code === code, - ) -} - -function readValue(argv: Array, index: number, flag: string): string { - const value = argv[index] - if (!value || value.startsWith('--')) { - throw new Error(`Missing value for ${flag}.`) - } - return value -} - -function readDialect( - argv: Array, - index: number, - config: MigrationCliConfig, -): TDialect { - const dialect = readValue(argv, index, '--dialect') - if (!isAllowedDialect(dialect, config.allowedDialects)) { - throw new Error( - `Invalid --dialect. Expected ${formatDialectList( - config.allowedDialects, - )}.`, - ) - } - return dialect -} - -function isAllowedDialect( - dialect: string, - allowedDialects: ReadonlyArray, -): dialect is TDialect { - return (allowedDialects as ReadonlyArray).includes(dialect) -} - -function readMigrationName(argv: Array, index: number): string { - const name = readValue(argv, index, '--name') - if (!/^[A-Za-z0-9_-]+$/.test(name)) { - throw new Error( - 'Invalid --name. Expected only letters, numbers, underscores, or hyphens.', - ) - } - return name -} - -function resolveOut(cwd: string, out: string): string { - return isAbsolute(out) ? out : join(cwd, out) -} - -function defaultTimestamp(now: Date): string { - const pad = (value: number) => String(value).padStart(2, '0') - return [ - now.getFullYear(), - pad(now.getMonth() + 1), - pad(now.getDate()), - pad(now.getHours()), - pad(now.getMinutes()), - pad(now.getSeconds()), - ].join('') -} - -function helpText(config: MigrationCliConfig): string { - return `${config.commandName} - -Generate a TanStack AI persistence SQL migration file. - -Options: - --dialect ${config.allowedDialects.join('|')} SQL dialect to generate. Defaults to ${config.defaultDialect}. - --out Write to a custom migration file path. - --stdout Print SQL instead of writing a file. - --timestamp Use a deterministic timestamp. - --name Safe migration name. Use letters, numbers, underscores, or hyphens. - --force Overwrite an existing output file. - --help, -h Show this help. -` -} - -function formatDialectList(dialects: ReadonlyArray): string { - if (dialects.length <= 1) return dialects.join('') - - const last = dialects[dialects.length - 1] - const initial = dialects.slice(0, -1) - return `${initial.join(', ')} or ${last}` -} diff --git a/packages/ai-persistence-sql/src/cli.ts b/packages/ai-persistence-sql/src/cli.ts deleted file mode 100644 index 4be4ea39a..000000000 --- a/packages/ai-persistence-sql/src/cli.ts +++ /dev/null @@ -1,9 +0,0 @@ -import { runMigrationCli, sqlMigrationCliConfig } from './cli-core' -import type { CliOptions, CliResult } from './cli-core' - -export function runCli( - argv: Array, - options?: CliOptions, -): Promise { - return runMigrationCli(argv, sqlMigrationCliConfig, options) -} diff --git a/packages/ai-persistence-sql/src/driver.ts b/packages/ai-persistence-sql/src/driver.ts deleted file mode 100644 index 0598dfd56..000000000 --- a/packages/ai-persistence-sql/src/driver.ts +++ /dev/null @@ -1,115 +0,0 @@ -/** - * Minimal SQL driver contract the shared stores are written against. Each - * backend (`-sqlite`, `-postgres`, `-cloudflare`, `-drizzle`, `-prisma`) - * implements this over its client; the stores never see a concrete driver. - */ - -export type Dialect = 'sqlite' | 'postgres' | 'mysql' - -export type SqlRow = Record - -export interface SqlDriver { - readonly dialect: Dialect - /** Run a statement that returns no rows (DDL, INSERT/UPDATE/DELETE). */ - exec: (sql: string, params?: ReadonlyArray) => Promise - /** Run a query and return all rows. */ - query: ( - sql: string, - params?: ReadonlyArray, - ) => Promise> - /** Run `fn` inside a transaction, passing a driver bound to the transaction. */ - transaction: (fn: (tx: SqlDriver) => Promise) => Promise -} - -/** - * Positional placeholder for the dialect: `?` for SQLite/MySQL, `$n` for Postgres. - * `index` is 1-based to match Postgres `$1`, `$2`, … - */ -export function param(dialect: Dialect, index: number): string { - return dialect === 'postgres' ? `$${index}` : '?' -} - -/** Build N placeholders starting at `start` (1-based), e.g. `?, ?, ?`. */ -export function params(dialect: Dialect, count: number, start = 1): string { - return Array.from({ length: count }, (_, i) => - param(dialect, start + i), - ).join(', ') -} - -/** Auto-incrementing integer primary-key column definition for the dialect. */ -export function autoIncrementPk(dialect: Dialect): string { - if (dialect === 'postgres') - return 'BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY' - if (dialect === 'mysql') return 'BIGINT AUTO_INCREMENT PRIMARY KEY' - return 'INTEGER PRIMARY KEY AUTOINCREMENT' -} - -/** JSON column type for the dialect. */ -export function jsonColumn(dialect: Dialect): string { - return dialect === 'postgres' ? 'JSONB' : 'TEXT' -} - -/** Large text payload column type for the dialect. */ -export function textColumn(dialect: Dialect): string { - return dialect === 'mysql' ? 'LONGTEXT' : 'TEXT' -} - -/** Binary/blob column type for the dialect. */ -export function blobColumn(dialect: Dialect): string { - return dialect === 'postgres' ? 'BYTEA' : 'BLOB' -} - -/** - * 64-bit integer column for the dialect. Epoch-ms timestamps overflow Postgres - * `INTEGER` (INT4, max ~2.1e9), so they must be `BIGINT`; SQLite `INTEGER` is - * already 64-bit. - */ -export function bigIntColumn(dialect: Dialect): string { - return dialect === 'sqlite' ? 'INTEGER' : 'BIGINT' -} - -/** Bounded text key column for dialects that cannot index unrestricted TEXT. */ -export function stringKeyColumn(dialect: Dialect): string { - return dialect === 'mysql' - ? 'VARCHAR(191) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin' - : 'TEXT' -} - -/** Quote an identifier for the target dialect. */ -export function identifier(dialect: Dialect, name: string): string { - if (dialect === 'mysql') return `\`${name.replaceAll('`', '``')}\`` - if (dialect === 'postgres') return `"${name.replaceAll('"', '""')}"` - return `"${name.replaceAll('"', '""')}"` -} - -/** Insert prefix for idempotent writes that should do nothing on unique conflicts. */ -export function insertDoNothingPrefix(dialect: Dialect): string { - void dialect - return 'INSERT INTO' -} - -/** Insert suffix for idempotent writes that should do nothing on unique conflicts. */ -export function insertDoNothingSuffix( - dialect: Dialect, - conflictColumns: ReadonlyArray, -): string { - if (dialect === 'mysql') { - const noopColumn = conflictColumns[0] - return ` ON DUPLICATE KEY UPDATE ${noopColumn} = ${noopColumn}` - } - return ` ON CONFLICT (${conflictColumns.join(', ')}) DO NOTHING` -} - -/** Upsert suffix for insert-or-update writes. */ -export function upsertUpdateSuffix( - dialect: Dialect, - conflictColumns: ReadonlyArray, - assignments: ReadonlyArray, -): string { - if (dialect === 'mysql') { - return ` ON DUPLICATE KEY UPDATE ${assignments.join(', ')}` - } - return ` ON CONFLICT (${conflictColumns.join( - ', ', - )}) DO UPDATE SET ${assignments.join(', ')}` -} diff --git a/packages/ai-persistence-sql/src/index.ts b/packages/ai-persistence-sql/src/index.ts deleted file mode 100644 index a658b4f57..000000000 --- a/packages/ai-persistence-sql/src/index.ts +++ /dev/null @@ -1,43 +0,0 @@ -// SQL driver contract + dialect helpers (the seam each backend implements). -export { - param, - params, - autoIncrementPk, - jsonColumn, - textColumn, - blobColumn, - bigIntColumn, - identifier, - stringKeyColumn, - insertDoNothingPrefix, - insertDoNothingSuffix, - upsertUpdateSuffix, -} from './driver' -export type { Dialect, SqlRow, SqlDriver } from './driver' - -// Schema migrations + raw DDL -export { migrate, ddl } from './migrations' - -// Migration CLI types. Runtime CLI helpers live on the ./cli-core subpath. -export type { - CliDialect, - CliOptions, - CliResult, - MigrationCliConfig, -} from './cli-core' - -// Shared SQL stores (one impl per store, dialect-parameterized) -export { - createMessageStore, - createRunStore, - createEventLog, - createLegacyEventLog, - createInternalEventStore, - createInterruptStore, - createApprovalStore, - createMetadataStore, -} from './stores' - -// Assemble a SqlPersistence from a driver -export { createSqlPersistence } from './sql-persistence' -export type { SqlPersistence, SqlPersistenceOptions } from './sql-persistence' diff --git a/packages/ai-persistence-sql/src/migrations.ts b/packages/ai-persistence-sql/src/migrations.ts deleted file mode 100644 index 6bc8870df..000000000 --- a/packages/ai-persistence-sql/src/migrations.ts +++ /dev/null @@ -1,125 +0,0 @@ -/** - * Versioned schema migrations. - * - * Raw drivers run these on first use (opt-out at the backend level). Each - * migration is applied at most once, tracked in `_tanstack_ai_migrations`. - * Idempotent: re-running `migrate` is a no-op once all versions are applied. - * - * JSON is stored as large text (TEXT, or LONGTEXT on MySQL) so reads don't - * depend on a driver's JSONB return shape; bytes are stored base64-encoded in - * a text column for the same portability reason. Epoch-ms timestamps use - * {@link bigIntColumn}. - */ -import { - bigIntColumn, - identifier, - insertDoNothingPrefix, - insertDoNothingSuffix, - param, - stringKeyColumn, - textColumn, -} from './driver' -import type { Dialect, SqlDriver } from './driver' - -interface Migration { - version: number - up: (dialect: Dialect) => Array -} - -/** DDL for schema v1. */ -function v1(dialect: Dialect): Array { - const ts = bigIntColumn(dialect) - const key = stringKeyColumn(dialect) - const text = textColumn(dialect) - const metadataKey = identifier(dialect, 'key') - const usageKey = identifier(dialect, 'usage') - return [ - `CREATE TABLE IF NOT EXISTS runs ( - run_id ${key} PRIMARY KEY, - thread_id ${key} NOT NULL, - status TEXT NOT NULL, - started_at ${ts} NOT NULL, - finished_at ${ts}, - error ${text}, - ${usageKey} ${text} - )`, - `CREATE TABLE IF NOT EXISTS public_events ( - run_id ${key} NOT NULL, - seq INTEGER NOT NULL, - event ${text} NOT NULL, - PRIMARY KEY (run_id, seq) - )`, - `CREATE TABLE IF NOT EXISTS internal_events ( - run_id ${key} NOT NULL, - namespace ${key} NOT NULL, - seq INTEGER NOT NULL, - type TEXT NOT NULL, - payload ${text} NOT NULL, - PRIMARY KEY (run_id, namespace, seq) - )`, - `CREATE TABLE IF NOT EXISTS messages ( - thread_id ${key} PRIMARY KEY, - messages ${text} NOT NULL - )`, - `CREATE TABLE IF NOT EXISTS interrupts ( - interrupt_id ${key} PRIMARY KEY, - run_id ${key} NOT NULL, - thread_id ${key} NOT NULL, - status TEXT NOT NULL, - requested_at ${ts} NOT NULL, - resolved_at ${ts}, - payload ${text} NOT NULL, - response ${text} - )`, - `CREATE TABLE IF NOT EXISTS metadata ( - scope ${key} NOT NULL, - ${metadataKey} ${key} NOT NULL, - value ${text} NOT NULL, - PRIMARY KEY (scope, ${metadataKey}) - )`, - ] -} - -const MIGRATIONS: ReadonlyArray = [{ version: 1, up: v1 }] - -/** Raw DDL strings for a dialect (every migration, in order) — for users who manage schema themselves. */ -export function ddl(dialect: Dialect): Array { - return MIGRATIONS.flatMap((m) => m.up(dialect)) -} - -/** - * Apply any not-yet-applied migrations. Idempotent and transactional per - * migration. Tracks applied versions in `_tanstack_ai_migrations`. - */ -export async function migrate(driver: SqlDriver): Promise { - await driver.exec( - `CREATE TABLE IF NOT EXISTS _tanstack_ai_migrations ( - version INTEGER PRIMARY KEY, - applied_at ${bigIntColumn(driver.dialect)} NOT NULL - )`, - ) - const rows = await driver.query<{ version: number }>( - 'SELECT version FROM _tanstack_ai_migrations', - ) - const applied = new Set(rows.map((r) => Number(r.version))) - - for (const migration of MIGRATIONS) { - if (applied.has(migration.version)) continue - await driver.transaction(async (tx) => { - for (const statement of migration.up(tx.dialect)) { - await tx.exec(statement) - } - await tx.exec( - `${insertDoNothingPrefix( - tx.dialect, - )} _tanstack_ai_migrations (version, applied_at) VALUES (${param( - tx.dialect, - 1, - )}, ${param(tx.dialect, 2)})${insertDoNothingSuffix(tx.dialect, [ - 'version', - ])}`, - [migration.version, Date.now()], - ) - }) - } -} diff --git a/packages/ai-persistence-sql/src/sql-persistence.ts b/packages/ai-persistence-sql/src/sql-persistence.ts deleted file mode 100644 index 7ec9494e7..000000000 --- a/packages/ai-persistence-sql/src/sql-persistence.ts +++ /dev/null @@ -1,106 +0,0 @@ -/** - * Assemble a {@link SqlPersistence} backed by SQL stores over a {@link SqlDriver}. - * - * By default the schema is not migrated; pass `migrate: true` to run lazy - * migrations on first use. If you manage the schema yourself, apply - * {@link migrate}/{@link ddl} outside this persistence wrapper. - * - * Migration runs lazily and exactly once, gated at the DRIVER level: the stores - * are built on a wrapper whose every `exec`/`query`/`transaction` first awaits - * the (idempotent, memoized) migration. This is why it composes correctly with - * `EventLog.read`, which returns an async iterable synchronously — the gate fires - * when the generator calls `query`, not when `read` is invoked. Migration itself - * runs on the RAW driver to avoid recursing through the gate. - */ -import { migrate as runMigrations } from './migrations' -import { - createApprovalStore, - createEventLog, - createInternalEventStore, - createInterruptStore, - createLegacyEventLog, - createMessageStore, - createMetadataStore, - createRunStore, -} from './stores' -import type { SqlDriver } from './driver' -import type { - AIPersistence, - ApprovalStore, - MessageStore, - PersistenceMode, - RunStore, -} from '@tanstack/ai-persistence' -import type { LegacyEventLog } from './stores' - -export interface SqlPersistenceOptions { - mode?: PersistenceMode - /** Run schema migrations on first use (default false). */ - migrate?: boolean -} - -export type SqlPersistence = AIPersistence & { - /** @deprecated Use stores.messages. */ - messages: MessageStore - /** @deprecated Use stores.runs. */ - runs: RunStore - /** @deprecated Use stores.publicEvents. */ - events: LegacyEventLog - /** @deprecated Use stores.interrupts. */ - approvals: ApprovalStore - /** @deprecated Use feature-based configuration instead of mode. */ - mode: PersistenceMode -} - -export function createSqlPersistence( - driver: SqlDriver, - opts?: SqlPersistenceOptions, -): SqlPersistence { - const shouldMigrate = opts?.migrate ?? false - - let migrated: Promise | undefined - const ensureSchema = (): Promise => { - if (!shouldMigrate) return Promise.resolve() - migrated ??= runMigrations(driver) // runs on the RAW driver — no recursion - return migrated - } - - const gated: SqlDriver = { - dialect: driver.dialect, - async exec(sql, params) { - await ensureSchema() - return driver.exec(sql, params) - }, - async query(sql, params) { - await ensureSchema() - return driver.query(sql, params) - }, - async transaction(fn) { - await ensureSchema() - return driver.transaction(fn) - }, - } - - const messages = createMessageStore(gated) - const runs = createRunStore(gated) - const publicEvents = createEventLog(gated) - const internalEvents = createInternalEventStore(gated) - const interrupts = createInterruptStore(gated) - const metadata = createMetadataStore(gated) - const approvals = createApprovalStore(gated) - return { - stores: { - messages, - runs, - publicEvents, - internalEvents, - interrupts, - metadata, - }, - messages, - runs, - events: createLegacyEventLog(gated), - approvals, - mode: opts?.mode ?? 'agent', - } -} diff --git a/packages/ai-persistence-sql/src/stores.ts b/packages/ai-persistence-sql/src/stores.ts deleted file mode 100644 index 46ad9cc26..000000000 --- a/packages/ai-persistence-sql/src/stores.ts +++ /dev/null @@ -1,700 +0,0 @@ -/** - * SQL implementations for the small durable persistence primitives. Dialect - * differences stay in placeholder syntax; JSON payloads are stored as TEXT for - * portability across SQLite, Postgres, D1, Drizzle, and Prisma raw drivers. - */ -import { AppendConflictError } from '@tanstack/ai-persistence' -import { - identifier, - insertDoNothingPrefix, - insertDoNothingSuffix, - param, - upsertUpdateSuffix, -} from './driver' -import type { - ApprovalRecord, - ApprovalStore, - InternalEventStore, - InterruptRecord, - InterruptStore, - MessageStore, - MetadataStore, - PersistedInternalEvent, - PersistedPublicEvent, - PublicEventStore, - RunRecord, - RunStatus, - RunStore, -} from '@tanstack/ai-persistence' -import type { ModelMessage, StreamChunk, TokenUsage } from '@tanstack/ai' -import type { SqlDriver } from './driver' - -const num = (v: unknown): number => Number(v) -const str = (v: unknown): string => String(v) - -function stableJsonValue(value: unknown): unknown { - if (Array.isArray(value)) { - return value.map((item) => stableJsonValue(item)) - } - if (value && typeof value === 'object') { - return Object.fromEntries( - Object.entries(value) - .sort(([a], [b]) => a.localeCompare(b)) - .map(([key, item]) => [key, stableJsonValue(item)]), - ) - } - return value -} - -const sameJson = (a: unknown, b: unknown): boolean => - JSON.stringify(stableJsonValue(a)) === JSON.stringify(stableJsonValue(b)) - -function encodeCursor(...parts: Array): string { - return Buffer.from(parts.join(':'), 'utf8').toString('base64url') -} - -export function createMessageStore(driver: SqlDriver): MessageStore { - const p = (i: number) => param(driver.dialect, i) - return { - async loadThread(threadId) { - const rows = await driver.query<{ messages: string }>( - `SELECT messages FROM messages WHERE thread_id = ${p(1)}`, - [threadId], - ) - const row = rows[0] - return row ? (JSON.parse(row.messages) as Array) : [] - }, - async saveThread(threadId, messages) { - const serialized = JSON.stringify(messages) - await driver.exec( - `INSERT INTO messages (thread_id, messages) VALUES (${p(1)}, ${p( - 2, - )})${upsertUpdateSuffix( - driver.dialect, - ['thread_id'], - [`messages = ${p(3)}`], - )}`, - [threadId, serialized, serialized], - ) - }, - } -} - -function mapRun(row: Record): RunRecord { - return { - runId: str(row.run_id), - threadId: str(row.thread_id), - status: str(row.status) as RunStatus, - startedAt: num(row.started_at), - ...(row.finished_at != null ? { finishedAt: num(row.finished_at) } : {}), - ...(row.error != null ? { error: str(row.error) } : {}), - ...(row.usage != null - ? { usage: JSON.parse(str(row.usage)) as TokenUsage } - : {}), - } -} - -export function createRunStore(driver: SqlDriver): RunStore { - const p = (i: number) => param(driver.dialect, i) - const usageKey = identifier(driver.dialect, 'usage') - const runColumns = [ - 'run_id', - 'thread_id', - 'status', - 'started_at', - 'finished_at', - 'error', - usageKey, - ].join(', ') - return { - async createOrResume(input) { - const existing = await this.get(input.runId) - if (existing) return existing - const record: RunRecord = { - runId: input.runId, - threadId: input.threadId, - status: input.status ?? 'running', - startedAt: input.startedAt, - } - await driver.exec( - `${insertDoNothingPrefix( - driver.dialect, - )} runs (run_id, thread_id, status, started_at) VALUES (${p( - 1, - )}, ${p(2)}, ${p(3)}, ${p(4)})${insertDoNothingSuffix(driver.dialect, [ - 'run_id', - ])}`, - [record.runId, record.threadId, record.status, record.startedAt], - ) - return (await this.get(input.runId)) ?? record - }, - async update(runId, patch) { - const sets: Array = [] - const values: Array = [] - let i = 1 - if (patch.status !== undefined) { - sets.push(`status = ${p(i++)}`) - values.push(patch.status) - } - if (patch.finishedAt !== undefined) { - sets.push(`finished_at = ${p(i++)}`) - values.push(patch.finishedAt) - } - if (patch.error !== undefined) { - sets.push(`error = ${p(i++)}`) - values.push(patch.error) - } - if (patch.usage !== undefined) { - sets.push(`${usageKey} = ${p(i++)}`) - values.push(JSON.stringify(patch.usage)) - } - if (!sets.length) return - values.push(runId) - await driver.exec( - `UPDATE runs SET ${sets.join(', ')} WHERE run_id = ${p(i)}`, - values, - ) - }, - async get(runId) { - const rows = await driver.query( - `SELECT ${runColumns} FROM runs WHERE run_id = ${p(1)}`, - [runId], - ) - const row = rows[0] - return row ? mapRun(row) : null - }, - } -} - -function publicEventFromRow( - runId: string, - row: { seq: number; event: string }, -): PersistedPublicEvent { - const seq = num(row.seq) - return { - seq, - event: JSON.parse(str(row.event)) as StreamChunk, - cursor: encodeCursor(seq, runId), - } -} - -export function createEventLog(driver: SqlDriver): PublicEventStore { - const p = (i: number) => param(driver.dialect, i) - async function readAt( - tx: SqlDriver, - runId: string, - seq: number, - ): Promise { - const rows = await tx.query<{ seq: number; event: string }>( - `SELECT seq, event FROM public_events WHERE run_id = ${p( - 1, - )} AND seq = ${p(2)}`, - [runId, seq], - ) - const row = rows[0] - return row ? publicEventFromRow(runId, row) : null - } - async function latest(tx: SqlDriver, runId: string): Promise { - const rows = await tx.query<{ max_seq: number | null }>( - `SELECT MAX(seq) AS max_seq FROM public_events WHERE run_id = ${p(1)}`, - [runId], - ) - const max = rows[0]?.max_seq - return max == null ? 0 : num(max) - } - return { - append(input) { - return driver.transaction(async (tx) => { - const targetSeq = input.expectedSeq + 1 - const existingAtTarget = await readAt(tx, input.runId, targetSeq) - if (existingAtTarget) { - if (sameJson(existingAtTarget.event, input.event)) { - return existingAtTarget - } - throw new AppendConflictError( - `Public event append conflict for run ${input.runId} at seq ${targetSeq}`, - ) - } - - const current = await latest(tx, input.runId) - if (current !== input.expectedSeq) { - throw new AppendConflictError( - `Public event append conflict for run ${input.runId}: expected latest seq ${input.expectedSeq}, got ${current}`, - ) - } - - await tx.exec( - `${insertDoNothingPrefix( - tx.dialect, - )} public_events (run_id, seq, event) VALUES (${p( - 1, - )}, ${p(2)}, ${p(3)})${insertDoNothingSuffix(tx.dialect, [ - 'run_id', - 'seq', - ])}`, - [input.runId, targetSeq, JSON.stringify(input.event)], - ) - const persisted = await readAt(tx, input.runId, targetSeq) - if (!persisted || !sameJson(persisted.event, input.event)) { - throw new AppendConflictError( - `Public event append conflict for run ${input.runId} at seq ${targetSeq}`, - ) - } - return persisted - }) - }, - read(runId, opts) { - const after = opts?.afterSeq - const sql = - after === undefined - ? `SELECT seq, event FROM public_events WHERE run_id = ${p( - 1, - )} ORDER BY seq ASC` - : `SELECT seq, event FROM public_events WHERE run_id = ${p( - 1, - )} AND seq > ${p(2)} ORDER BY seq ASC` - const values = after === undefined ? [runId] : [runId, after] - return (async function* (): AsyncIterable { - const rows = await driver.query<{ seq: number; event: string }>( - sql, - values, - ) - for (const row of rows) yield publicEventFromRow(runId, row) - })() - }, - async hasRun(runId) { - const rows = await driver.query( - `SELECT 1 AS one FROM public_events WHERE run_id = ${p(1)} LIMIT 1`, - [runId], - ) - return rows.length > 0 - }, - latestSeq(runId) { - return latest(driver, runId) - }, - } -} - -export type LegacyEventLog = Omit & { - append: { - (runId: string, seq: number, event: StreamChunk): Promise - (input: { - runId: string - expectedSeq: number - event: StreamChunk - }): Promise - } -} - -export function createLegacyEventLog(driver: SqlDriver): LegacyEventLog { - const store = createEventLog(driver) - const p = (i: number) => param(driver.dialect, i) - function append(runId: string, seq: number, event: StreamChunk): Promise - function append(input: { - runId: string - expectedSeq: number - event: StreamChunk - }): Promise - async function append( - inputOrRunId: - | string - | { runId: string; expectedSeq: number; event: StreamChunk }, - seq?: number, - event?: StreamChunk, - ): Promise { - if (typeof inputOrRunId !== 'string') return store.append(inputOrRunId) - if (seq === undefined || event === undefined) { - throw new Error(`Legacy event append requires runId, seq, and event.`) - } - await driver.exec( - `${insertDoNothingPrefix( - driver.dialect, - )} public_events (run_id, seq, event) VALUES (${p( - 1, - )}, ${p(2)}, ${p(3)})${insertDoNothingSuffix(driver.dialect, [ - 'run_id', - 'seq', - ])}`, - [inputOrRunId, seq, JSON.stringify(event)], - ) - } - return { - append, - read: store.read, - hasRun: store.hasRun, - latestSeq: store.latestSeq, - } -} - -function internalEventFromRow( - runId: string, - row: { seq: number; namespace: string; type: string; payload: string }, -): PersistedInternalEvent { - const seq = num(row.seq) - const namespace = str(row.namespace) - return { - seq, - namespace, - type: str(row.type), - payload: JSON.parse(str(row.payload)), - cursor: encodeCursor(seq, runId, namespace), - } -} - -export function createInternalEventStore( - driver: SqlDriver, -): InternalEventStore { - const p = (i: number) => param(driver.dialect, i) - async function readAt( - tx: SqlDriver, - runId: string, - namespace: string, - seq: number, - ): Promise { - const rows = await tx.query<{ - seq: number - namespace: string - type: string - payload: string - }>( - `SELECT seq, namespace, type, payload FROM internal_events WHERE run_id = ${p( - 1, - )} AND namespace = ${p(2)} AND seq = ${p(3)}`, - [runId, namespace, seq], - ) - const row = rows[0] - return row ? internalEventFromRow(runId, row) : null - } - async function latest( - tx: SqlDriver, - runId: string, - namespace?: string, - ): Promise { - const sql = - namespace === undefined - ? `SELECT MAX(seq) AS max_seq FROM internal_events WHERE run_id = ${p( - 1, - )}` - : `SELECT MAX(seq) AS max_seq FROM internal_events WHERE run_id = ${p( - 1, - )} AND namespace = ${p(2)}` - const values = namespace === undefined ? [runId] : [runId, namespace] - const rows = await tx.query<{ max_seq: number | null }>(sql, values) - const max = rows[0]?.max_seq - return max == null ? 0 : num(max) - } - return { - append(input) { - return driver.transaction(async (tx) => { - const targetSeq = input.expectedSeq + 1 - const existingAtTarget = await readAt( - tx, - input.runId, - input.namespace, - targetSeq, - ) - if (existingAtTarget) { - if ( - existingAtTarget.type === input.type && - sameJson(existingAtTarget.payload, input.payload) - ) { - return existingAtTarget - } - throw new AppendConflictError( - `Internal event append conflict for run ${input.runId}, namespace ${input.namespace}, seq ${targetSeq}`, - ) - } - - const current = await latest(tx, input.runId, input.namespace) - if (current !== input.expectedSeq) { - throw new AppendConflictError( - `Internal event append conflict for run ${input.runId}, namespace ${input.namespace}: expected latest seq ${input.expectedSeq}, got ${current}`, - ) - } - - await tx.exec( - `${insertDoNothingPrefix( - tx.dialect, - )} internal_events (run_id, namespace, seq, type, payload) VALUES (${p( - 1, - )}, ${p(2)}, ${p(3)}, ${p(4)}, ${p(5)})${insertDoNothingSuffix( - tx.dialect, - ['run_id', 'namespace', 'seq'], - )}`, - [ - input.runId, - input.namespace, - targetSeq, - input.type, - JSON.stringify(input.payload), - ], - ) - const persisted = await readAt( - tx, - input.runId, - input.namespace, - targetSeq, - ) - if ( - !persisted || - persisted.type !== input.type || - !sameJson(persisted.payload, input.payload) - ) { - throw new AppendConflictError( - `Internal event append conflict for run ${input.runId}, namespace ${input.namespace}, seq ${targetSeq}`, - ) - } - return persisted - }) - }, - read(runId, opts) { - const values: Array = [runId] - const filters = [`run_id = ${p(1)}`] - let i = 2 - if (opts?.namespace !== undefined) { - filters.push(`namespace = ${p(i++)}`) - values.push(opts.namespace) - } - if (opts?.afterSeq !== undefined) { - filters.push(`seq > ${p(i++)}`) - values.push(opts.afterSeq) - } - const sql = `SELECT seq, namespace, type, payload FROM internal_events WHERE ${filters.join( - ' AND ', - )} ORDER BY namespace ASC, seq ASC` - return (async function* (): AsyncIterable { - const rows = await driver.query<{ - seq: number - namespace: string - type: string - payload: string - }>(sql, values) - for (const row of rows) yield internalEventFromRow(runId, row) - })() - }, - latestSeq(runId, namespace) { - return latest(driver, runId, namespace) - }, - } -} - -function mapInterrupt(row: Record): InterruptRecord { - return { - interruptId: str(row.interrupt_id), - runId: str(row.run_id), - threadId: str(row.thread_id), - status: str(row.status) as InterruptRecord['status'], - requestedAt: num(row.requested_at), - ...(row.resolved_at != null ? { resolvedAt: num(row.resolved_at) } : {}), - payload: JSON.parse(str(row.payload)) as Record, - ...(row.response != null - ? { response: JSON.parse(str(row.response)) } - : {}), - } -} - -export function createInterruptStore(driver: SqlDriver): InterruptStore { - const p = (i: number) => param(driver.dialect, i) - return { - async create(record) { - await driver.exec( - `${insertDoNothingPrefix( - driver.dialect, - )} interrupts (interrupt_id, run_id, thread_id, status, requested_at, payload, response) - VALUES (${p(1)}, ${p(2)}, ${p(3)}, ${p(4)}, ${p(5)}, ${p(6)}, ${p(7)}) - ${insertDoNothingSuffix(driver.dialect, ['interrupt_id'])}`, - [ - record.interruptId, - record.runId, - record.threadId, - record.status, - record.requestedAt, - JSON.stringify(record.payload), - record.response === undefined - ? null - : JSON.stringify(record.response), - ], - ) - }, - async resolve(interruptId, response) { - await driver.exec( - `UPDATE interrupts SET status = ${p(1)}, resolved_at = ${p( - 2, - )}, response = ${p(3)} WHERE interrupt_id = ${p(4)}`, - [ - 'resolved', - Date.now(), - response === undefined ? null : JSON.stringify(response), - interruptId, - ], - ) - }, - async cancel(interruptId) { - await driver.exec( - `UPDATE interrupts SET status = ${p(1)}, resolved_at = ${p( - 2, - )} WHERE interrupt_id = ${p(3)}`, - ['cancelled', Date.now(), interruptId], - ) - }, - async get(interruptId) { - const rows = await driver.query( - `SELECT * FROM interrupts WHERE interrupt_id = ${p(1)}`, - [interruptId], - ) - const row = rows[0] - return row ? mapInterrupt(row) : null - }, - async list(threadId) { - const rows = await driver.query( - `SELECT * FROM interrupts WHERE thread_id = ${p( - 1, - )} ORDER BY requested_at ASC`, - [threadId], - ) - return rows.map(mapInterrupt) - }, - async listPending(threadId) { - const rows = await driver.query( - `SELECT * FROM interrupts WHERE thread_id = ${p( - 1, - )} AND status = 'pending' ORDER BY requested_at ASC`, - [threadId], - ) - return rows.map(mapInterrupt) - }, - async listByRun(runId) { - const rows = await driver.query( - `SELECT * FROM interrupts WHERE run_id = ${p( - 1, - )} ORDER BY requested_at ASC`, - [runId], - ) - return rows.map(mapInterrupt) - }, - async listPendingByRun(runId) { - const rows = await driver.query( - `SELECT * FROM interrupts WHERE run_id = ${p( - 1, - )} AND status = 'pending' ORDER BY requested_at ASC`, - [runId], - ) - return rows.map(mapInterrupt) - }, - } -} - -function mapApproval(row: Record): ApprovalRecord { - const status = str(row.status) - return { - approvalId: str(row.interrupt_id), - runId: str(row.run_id), - threadId: str(row.thread_id), - status: - status === 'resolved' - ? 'granted' - : status === 'cancelled' - ? 'denied' - : 'pending', - requestedAt: num(row.requested_at), - ...(row.resolved_at != null ? { resolvedAt: num(row.resolved_at) } : {}), - payload: JSON.parse(str(row.payload)) as Record, - } -} - -export function createApprovalStore(driver: SqlDriver): ApprovalStore { - const p = (i: number) => param(driver.dialect, i) - return { - async create(record) { - await driver.exec( - `${insertDoNothingPrefix( - driver.dialect, - )} interrupts (interrupt_id, run_id, thread_id, status, requested_at, payload, response) - VALUES (${p(1)}, ${p(2)}, ${p(3)}, ${p(4)}, ${p(5)}, ${p(6)}, ${p(7)}) - ${insertDoNothingSuffix(driver.dialect, ['interrupt_id'])}`, - [ - record.approvalId, - record.runId, - record.threadId, - record.status === 'pending' - ? 'pending' - : record.status === 'granted' - ? 'resolved' - : 'cancelled', - record.requestedAt, - JSON.stringify({ ...record.payload, compatibility: 'approval' }), - null, - ], - ) - }, - async resolve(approvalId, granted) { - await driver.exec( - `UPDATE interrupts SET status = ${p(1)}, resolved_at = ${p( - 2, - )} WHERE interrupt_id = ${p(3)}`, - [granted ? 'resolved' : 'cancelled', Date.now(), approvalId], - ) - }, - async get(approvalId) { - const rows = await driver.query( - `SELECT * FROM interrupts WHERE interrupt_id = ${p(1)}`, - [approvalId], - ) - const row = rows[0] - return row ? mapApproval(row) : null - }, - async decisionsForThread(threadId) { - const rows = await driver.query( - `SELECT interrupt_id, status, payload FROM interrupts WHERE thread_id = ${p( - 1, - )} AND status <> 'pending'`, - [threadId], - ) - const decisions = new Map() - for (const row of rows) { - const payload = JSON.parse(str(row.payload)) as Record - if (payload.compatibility !== 'approval') continue - decisions.set(str(row.interrupt_id), str(row.status) === 'resolved') - } - return decisions - }, - } -} - -export function createMetadataStore(driver: SqlDriver): MetadataStore { - const p = (i: number) => param(driver.dialect, i) - const metadataKey = identifier(driver.dialect, 'key') - return { - async get(scope, key) { - const rows = await driver.query<{ value: string }>( - `SELECT value FROM metadata WHERE scope = ${p( - 1, - )} AND ${metadataKey} = ${p(2)}`, - [scope, key], - ) - const row = rows[0] - return row ? JSON.parse(str(row.value)) : null - }, - async set(scope, key, value) { - const serialized = JSON.stringify(value) - await driver.exec( - `INSERT INTO metadata (scope, ${metadataKey}, value) VALUES (${p( - 1, - )}, ${p(2)}, ${p(3)})${upsertUpdateSuffix( - driver.dialect, - ['scope', metadataKey], - [`value = ${p(4)}`], - )}`, - [scope, key, serialized, serialized], - ) - }, - async delete(scope, key) { - await driver.exec( - `DELETE FROM metadata WHERE scope = ${p( - 1, - )} AND ${metadataKey} = ${p(2)}`, - [scope, key], - ) - }, - } -} diff --git a/packages/ai-persistence-sql/tests/cli.test.ts b/packages/ai-persistence-sql/tests/cli.test.ts deleted file mode 100644 index 1b2ff4bdf..000000000 --- a/packages/ai-persistence-sql/tests/cli.test.ts +++ /dev/null @@ -1,129 +0,0 @@ -import { mkdtemp, readFile, writeFile } from 'node:fs/promises' -import { tmpdir } from 'node:os' -import { join } from 'node:path' -import { describe, expect, it } from 'vitest' -import { runCli } from '../src/cli' - -describe('SQL persistence CLI', () => { - it('writes the default generic migration path', async () => { - const cwd = await mkdtemp(join(tmpdir(), 'tanstack-ai-sql-')) - - const result = await runCli( - ['--dialect', 'postgres', '--timestamp', '20260706123456'], - { cwd }, - ) - - expect(result).toEqual({ - kind: 'file', - path: join( - cwd, - 'migrations', - '20260706123456_tanstack_ai_persistence.sql', - ), - }) - if (result.kind === 'file') { - await expect(readFile(result.path, 'utf8')).resolves.toContain( - 'TanStack AI persistence migration for postgres', - ) - } - }) - - it('writes a custom output path with a custom migration name', async () => { - const cwd = await mkdtemp(join(tmpdir(), 'tanstack-ai-sql-')) - const out = join(cwd, 'db', 'migrations', 'custom.sql') - - const result = await runCli( - [ - '--dialect', - 'sqlite', - '--timestamp', - '20260706123456', - '--name', - 'custom_name', - '--out', - out, - ], - { cwd }, - ) - - expect(result).toEqual({ kind: 'file', path: out }) - await expect(readFile(out, 'utf8')).resolves.toContain( - 'TanStack AI persistence migration for sqlite', - ) - }) - - it('prints SQL to stdout for every supported dialect', async () => { - for (const dialect of ['sqlite', 'postgres', 'mysql']) { - const written: Array = [] - const result = await runCli(['--dialect', dialect, '--stdout'], { - writeStdout: (value) => written.push(value), - }) - - expect(result.kind).toBe('stdout') - if (result.kind === 'stdout') { - expect(result.sql).toContain( - `TanStack AI persistence migration for ${dialect}`, - ) - expect(result.sql).toContain('CREATE TABLE IF NOT EXISTS runs') - expect(written).toEqual([result.sql]) - } - } - }) - - it('shows help text', async () => { - const result = await runCli(['--help']) - - expect(result.kind).toBe('help') - if (result.kind === 'help') { - expect(result.text).toContain('tanstack-ai-persistence-sql') - expect(result.text).toContain('--dialect sqlite|postgres|mysql') - expect(result.text).toContain('--stdout') - } - }) - - it('rejects invalid dialects', async () => { - await expect(runCli(['--dialect', 'mssql'])).rejects.toThrow( - 'Invalid --dialect', - ) - }) - - it('rejects missing option values', async () => { - await expect(runCli(['--out'])).rejects.toThrow('Missing value for --out') - await expect(runCli(['--dialect'])).rejects.toThrow( - 'Missing value for --dialect', - ) - await expect(runCli(['--timestamp'])).rejects.toThrow( - 'Missing value for --timestamp', - ) - await expect(runCli(['--name'])).rejects.toThrow('Missing value for --name') - }) - - it('rejects path traversal migration names', async () => { - await expect(runCli(['--name', '../escape'])).rejects.toThrow( - 'Invalid --name', - ) - }) - - it('refuses to overwrite an existing file without --force', async () => { - const cwd = await mkdtemp(join(tmpdir(), 'tanstack-ai-sql-')) - const out = join(cwd, 'migration.sql') - await writeFile(out, 'existing') - - await expect(runCli(['--dialect', 'mysql', '--out', out])).rejects.toThrow( - 'already exists', - ) - }) - - it('overwrites an existing file with --force', async () => { - const cwd = await mkdtemp(join(tmpdir(), 'tanstack-ai-sql-')) - const out = join(cwd, 'migration.sql') - await writeFile(out, 'existing') - - const result = await runCli(['--dialect', 'mysql', '--out', out, '--force']) - - expect(result).toEqual({ kind: 'file', path: out }) - await expect(readFile(out, 'utf8')).resolves.toContain( - 'TanStack AI persistence migration for mysql', - ) - }) -}) diff --git a/packages/ai-persistence-sql/tests/sql-persistence.test.ts b/packages/ai-persistence-sql/tests/sql-persistence.test.ts deleted file mode 100644 index 320867d34..000000000 --- a/packages/ai-persistence-sql/tests/sql-persistence.test.ts +++ /dev/null @@ -1,765 +0,0 @@ -import { describe, expect, it } from 'vitest' -import { EventType } from '@tanstack/ai' -import { - AppendConflictError, - createApprovalController, -} from '@tanstack/ai-persistence' -import type { StreamChunk } from '@tanstack/ai' -import { createSqlPersistence } from '../src/sql-persistence' -import { ddl, migrate } from '../src/migrations' -import { createTestSqliteDriver } from './sqlite-driver' -import type { SqlDriver } from '../src/driver' - -const text = (delta: string): StreamChunk => ({ - type: EventType.TEXT_MESSAGE_CONTENT, - messageId: 'm1', - delta, - timestamp: 1, -}) - -const textWithMetadata = ( - delta: string, - metadata: Record, -): StreamChunk => Object.assign(text(delta), { metadata }) - -function createInsertRaceDriver(opts: { - table: 'public_events' | 'internal_events' - persistedPayload: unknown -}): SqlDriver { - const base = createTestSqliteDriver() - let raced = false - let aborted = false - const makeDriver = (inTransaction: boolean): SqlDriver => ({ - dialect: base.dialect, - async query(sql, params) { - if (inTransaction && aborted) { - throw new Error('current transaction is aborted') - } - return base.query(sql, params) - }, - transaction: async (fn) => { - aborted = false - try { - return await fn(makeDriver(true)) - } finally { - aborted = false - } - }, - async exec(sql, params = []) { - if (inTransaction && aborted) { - throw new Error('current transaction is aborted') - } - if (!raced && sql.includes(`INSERT INTO ${opts.table}`)) { - raced = true - if (opts.table === 'public_events') { - await base.exec( - 'INSERT INTO public_events (run_id, seq, event) VALUES (?, ?, ?)', - [params[0], params[1], JSON.stringify(opts.persistedPayload)], - ) - } else { - await base.exec( - 'INSERT INTO internal_events (run_id, namespace, seq, type, payload) VALUES (?, ?, ?, ?, ?)', - [ - params[0], - params[1], - params[2], - params[3], - JSON.stringify(opts.persistedPayload), - ], - ) - } - if (!sql.includes('ON CONFLICT')) { - aborted = true - throw new Error('simulated unique constraint violation') - } - } - await base.exec(sql, params) - }, - }) - return makeDriver(false) -} - -function createRecordingDriver(): SqlDriver & { - statements: Array - queueQueryRows: (rows: Array>) => void -} { - const statements: Array = [] - const queuedRows: Array>> = [] - const driver: SqlDriver & { - statements: Array - queueQueryRows: (rows: Array>) => void - } = { - dialect: 'mysql' as SqlDriver['dialect'], - statements, - queueQueryRows(rows) { - queuedRows.push(rows) - }, - async exec(sql) { - statements.push(sql) - }, - async query(sql) { - statements.push(sql) - return (queuedRows.shift() ?? []) as Array - }, - async transaction(fn) { - return fn(driver) - }, - } - return driver -} - -function createMigratingTestPersistence() { - return createSqlPersistence(createTestSqliteDriver(), { migrate: true }) -} - -describe('migrate', () => { - it('is idempotent (re-running applies nothing new)', async () => { - const driver = createTestSqliteDriver() - await migrate(driver) - await migrate(driver) - const rows = await driver.query<{ version: number }>( - 'SELECT version FROM _tanstack_ai_migrations', - ) - expect(rows.map((r) => Number(r.version))).toEqual([1]) - }) - - it('creates only the small durable primitive tables', async () => { - const driver = createTestSqliteDriver() - await migrate(driver) - - const rows = await driver.query<{ name: string }>( - "SELECT name FROM sqlite_master WHERE type = 'table' AND name NOT LIKE 'sqlite_%' ORDER BY name", - ) - expect(rows.map((r) => r.name)).toEqual([ - '_tanstack_ai_migrations', - 'internal_events', - 'interrupts', - 'messages', - 'metadata', - 'public_events', - 'runs', - ]) - expect(rows.map((r) => r.name)).not.toEqual( - expect.arrayContaining([ - 'approvals', - 'artifacts', - 'checkpoints', - 'mcp_credentials', - 'mcp_sessions', - 'message_threads', - 'run_events', - 'sandbox_processes', - 'workflow_checkpoints', - ]), - ) - }) - - it('handles concurrent migration marker insertion idempotently', async () => { - const base = createTestSqliteDriver() - let migrationReads = 0 - let releaseReads: (() => void) | undefined - const bothReadsReached = new Promise((resolve) => { - releaseReads = resolve - }) - const driver: SqlDriver = { - dialect: base.dialect, - exec: base.exec, - transaction: base.transaction, - async query(sql, params) { - if (sql === 'SELECT version FROM _tanstack_ai_migrations') { - migrationReads += 1 - if (migrationReads === 2) releaseReads?.() - await bothReadsReached - } - return base.query(sql, params) - }, - } - - await Promise.all([migrate(driver), migrate(driver)]) - - const rows = await base.query<{ version: number }>( - 'SELECT version FROM _tanstack_ai_migrations', - ) - expect(rows.map((row) => Number(row.version))).toEqual([1]) - }) -}) - -describe('mysql dialect SQL generation', () => { - it('uses MySQL-safe binary key columns and LONGTEXT payload columns in DDL', () => { - const statements = ddl('mysql' as SqlDriver['dialect']) - const schema = statements.join('\n') - const mysqlKey = 'VARCHAR(191) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin' - - expect(schema).toContain(`run_id ${mysqlKey} PRIMARY KEY`) - expect(schema).toContain(`thread_id ${mysqlKey} PRIMARY KEY`) - expect(schema).toContain(`interrupt_id ${mysqlKey} PRIMARY KEY`) - expect(schema).toContain(`namespace ${mysqlKey} NOT NULL`) - expect(schema).toContain(`\`key\` ${mysqlKey} NOT NULL`) - expect(schema).toContain('PRIMARY KEY (scope, `key`)') - expect(schema).toContain('error LONGTEXT') - expect(schema).toContain('`usage` LONGTEXT') - expect(schema).toContain('event LONGTEXT NOT NULL') - expect(schema).toContain('payload LONGTEXT NOT NULL') - expect(schema).toContain('messages LONGTEXT NOT NULL') - expect(schema).toContain('response LONGTEXT') - expect(schema).toContain('value LONGTEXT NOT NULL') - expect(schema).not.toContain('TEXT PRIMARY KEY') - expect(schema).not.toContain(' key VARCHAR(191)') - }) - - it('uses no-op upsert for MySQL migration markers', async () => { - const driver = createRecordingDriver() - await migrate(driver) - - expect(driver.statements.join('\n')).toContain( - 'INSERT INTO _tanstack_ai_migrations', - ) - expect(driver.statements.join('\n')).toContain( - 'ON DUPLICATE KEY UPDATE version = version', - ) - expect(driver.statements.join('\n')).not.toContain('INSERT IGNORE') - expect(driver.statements.join('\n')).not.toContain('ON CONFLICT') - }) - - it('emits MySQL-compatible no-op insert and upsert SQL for stores', async () => { - const driver = createRecordingDriver() - const persistence = createSqlPersistence(driver, { migrate: false }) - - driver.queueQueryRows([]) - driver.queueQueryRows([ - { - run_id: 'r1', - thread_id: 't1', - status: 'running', - started_at: 1, - }, - ]) - await persistence.stores.runs!.createOrResume({ - runId: 'r1', - threadId: 't1', - startedAt: 1, - }) - await persistence.stores.runs!.update('r1', { - usage: { promptTokens: 1, completionTokens: 2, totalTokens: 3 }, - }) - - await persistence.stores.messages!.saveThread('t1', [ - { role: 'user', content: 'hi' }, - ]) - await persistence.stores.metadata!.set('scope', 'key', { ok: true }) - await persistence.stores.metadata!.get('scope', 'key') - await persistence.stores.metadata!.delete('scope', 'key') - await persistence.stores.interrupts!.create({ - interruptId: 'i1', - runId: 'r1', - threadId: 't1', - status: 'pending', - requestedAt: 1, - payload: { kind: 'input' }, - }) - await persistence.approvals!.create({ - approvalId: 'a1', - runId: 'r1', - threadId: 't1', - status: 'pending', - requestedAt: 1, - payload: { kind: 'approval' }, - }) - await persistence.events!.append('r1', 1, text('legacy')) - - driver.queueQueryRows([]) - driver.queueQueryRows([{ max_seq: 0 }]) - driver.queueQueryRows([ - { - seq: 1, - event: JSON.stringify(text('public')), - }, - ]) - await persistence.stores.publicEvents!.append({ - runId: 'r1', - expectedSeq: 0, - event: text('public'), - }) - - driver.queueQueryRows([]) - driver.queueQueryRows([{ max_seq: 0 }]) - driver.queueQueryRows([ - { - seq: 1, - namespace: 'agent', - type: 'step', - payload: JSON.stringify({ ok: true }), - }, - ]) - await persistence.stores.internalEvents!.append({ - runId: 'r1', - namespace: 'agent', - expectedSeq: 0, - type: 'step', - payload: { ok: true }, - }) - - const sql = driver.statements.join('\n') - expect(sql).toContain('INSERT INTO runs') - expect(sql).toContain( - 'SELECT run_id, thread_id, status, started_at, finished_at, error, `usage` FROM runs WHERE run_id = ?', - ) - expect(sql).toContain('UPDATE runs SET `usage` = ? WHERE run_id = ?') - expect(sql).toContain('ON DUPLICATE KEY UPDATE run_id = run_id') - expect(sql).toContain('INSERT INTO interrupts') - expect(sql).toContain('ON DUPLICATE KEY UPDATE interrupt_id = interrupt_id') - expect(sql).toContain('INSERT INTO public_events') - expect(sql).toContain('INSERT INTO internal_events') - expect(sql).toContain('INSERT INTO metadata (scope, `key`, value) VALUES') - expect(sql).toContain( - 'SELECT value FROM metadata WHERE scope = ? AND `key` = ?', - ) - expect(sql).toContain('DELETE FROM metadata WHERE scope = ? AND `key` = ?') - expect(sql).toContain('ON DUPLICATE KEY UPDATE messages = ?') - expect(sql).toContain('ON DUPLICATE KEY UPDATE value = ?') - expect(sql).not.toContain('INSERT IGNORE') - expect(sql).not.toContain('ON CONFLICT') - }) -}) - -describe('createSqlPersistence (sqlite dialect)', () => { - it('does not run migrations by default unless explicitly enabled', async () => { - const defaultDriver = createRecordingDriver() - const defaultPersistence = createSqlPersistence(defaultDriver) - - await defaultPersistence.stores.metadata!.get('scope', 'key') - - expect(defaultDriver.statements.join('\n')).not.toContain( - '_tanstack_ai_migrations', - ) - - const optInDriver = createRecordingDriver() - const optInPersistence = createSqlPersistence(optInDriver, { - migrate: true, - }) - - await optInPersistence.stores.metadata!.get('scope', 'key') - - expect(optInDriver.statements.join('\n')).toContain( - 'CREATE TABLE IF NOT EXISTS _tanstack_ai_migrations', - ) - }) - - it('migrates lazily on first use and round-trips runs', async () => { - const p = createMigratingTestPersistence() - const run = await p.runs!.createOrResume({ - runId: 'r1', - threadId: 't1', - startedAt: 100, - }) - expect(run.status).toBe('running') - // Idempotent resume returns the same record. - const again = await p.runs!.createOrResume({ - runId: 'r1', - threadId: 't1', - startedAt: 999, - }) - expect(again.startedAt).toBe(100) - - await p.runs!.update('r1', { - status: 'completed', - finishedAt: 200, - usage: { promptTokens: 1, completionTokens: 2, totalTokens: 3 }, - }) - const got = await p.runs!.get('r1') - expect(got?.status).toBe('completed') - expect(got?.finishedAt).toBe(200) - expect(got?.usage?.totalTokens).toBe(3) - }) - - it('appends events and replays after a sequence', async () => { - const p = createMigratingTestPersistence() - await p.events!.append('r1', 1, text('a')) - await p.events!.append('r1', 2, text('b')) - await p.events!.append('r1', 3, text('c')) - expect(await p.events!.hasRun('r1')).toBe(true) - expect(await p.events!.latestSeq('r1')).toBe(3) - - const seen: Array<{ seq: number; delta: string }> = [] - for await (const e of p.events!.read('r1', { afterSeq: 1 })) { - if (e.event.type === 'TEXT_MESSAGE_CONTENT') { - seen.push({ seq: e.seq, delta: e.event.delta }) - } - } - expect(seen).toEqual([ - { seq: 2, delta: 'b' }, - { seq: 3, delta: 'c' }, - ]) - }) - - it('append is idempotent on (runId, seq)', async () => { - const p = createMigratingTestPersistence() - await p.events!.append('r1', 1, text('a')) - await p.events!.append('r1', 1, text('a-again')) - expect(await p.events!.latestSeq('r1')).toBe(1) - }) - - it('public event append detects conflicting target rows', async () => { - const p = createMigratingTestPersistence() - await p.stores.publicEvents!.append({ - runId: 'r1', - expectedSeq: 0, - event: text('a'), - }) - - await expect( - p.stores.publicEvents!.append({ - runId: 'r1', - expectedSeq: 0, - event: text('different'), - }), - ).rejects.toBeInstanceOf(AppendConflictError) - }) - - it('public event append enforces expected sequence and permits identical retry', async () => { - const p = createMigratingTestPersistence() - const first = await p.stores.publicEvents!.append({ - runId: 'r1', - expectedSeq: 0, - event: text('a'), - }) - const retry = await p.stores.publicEvents!.append({ - runId: 'r1', - expectedSeq: 0, - event: text('a'), - }) - expect(retry).toEqual(first) - - await expect( - p.stores.publicEvents!.append({ - runId: 'r1', - expectedSeq: 0, - event: text('b'), - }), - ).rejects.toBeInstanceOf(AppendConflictError) - await expect( - p.stores.publicEvents!.append({ - runId: 'r1', - expectedSeq: 2, - event: text('c'), - }), - ).rejects.toBeInstanceOf(AppendConflictError) - }) - - it('public event idempotent retry ignores JSON object key order', async () => { - const p = createMigratingTestPersistence() - const first = await p.stores.publicEvents!.append({ - runId: 'r1', - expectedSeq: 0, - event: textWithMetadata('a', { a: 1, b: 2 }), - }) - const retry = await p.stores.publicEvents!.append({ - runId: 'r1', - expectedSeq: 0, - event: textWithMetadata('a', { b: 2, a: 1 }), - }) - - expect(retry).toEqual(first) - }) - - it('public event append normalizes insert races to CAS outcomes', async () => { - const event = text('winner') - const idempotent = createSqlPersistence( - createInsertRaceDriver({ - table: 'public_events', - persistedPayload: event, - }), - { migrate: true }, - ) - - await expect( - idempotent.stores.publicEvents!.append({ - runId: 'r1', - expectedSeq: 0, - event, - }), - ).resolves.toMatchObject({ seq: 1, event }) - - const conflicting = createSqlPersistence( - createInsertRaceDriver({ - table: 'public_events', - persistedPayload: text('other'), - }), - { migrate: true }, - ) - await expect( - conflicting.stores.publicEvents!.append({ - runId: 'r1', - expectedSeq: 0, - event, - }), - ).rejects.toBeInstanceOf(AppendConflictError) - }) - - it('internal event append enforces expected sequence per run and namespace', async () => { - const p = createMigratingTestPersistence() - const first = await p.stores.internalEvents!.append({ - runId: 'r1', - namespace: 'agent', - expectedSeq: 0, - type: 'step', - payload: { n: 1 }, - }) - const retry = await p.stores.internalEvents!.append({ - runId: 'r1', - namespace: 'agent', - expectedSeq: 0, - type: 'step', - payload: { n: 1 }, - }) - expect(retry).toEqual(first) - - await p.stores.internalEvents!.append({ - runId: 'r1', - namespace: 'tools', - expectedSeq: 0, - type: 'call', - payload: { tool: 'search' }, - }) - expect(await p.stores.internalEvents!.latestSeq('r1', 'agent')).toBe(1) - expect(await p.stores.internalEvents!.latestSeq('r1', 'tools')).toBe(1) - expect(await p.stores.internalEvents!.latestSeq('r1')).toBe(1) - - await expect( - p.stores.internalEvents!.append({ - runId: 'r1', - namespace: 'agent', - expectedSeq: 0, - type: 'step', - payload: { n: 2 }, - }), - ).rejects.toBeInstanceOf(AppendConflictError) - - const agentEvents = [] - for await (const event of p.stores.internalEvents!.read('r1', { - namespace: 'agent', - })) { - agentEvents.push(event) - } - expect(agentEvents.map((event) => event.payload)).toEqual([{ n: 1 }]) - }) - - it('internal event idempotent retry ignores JSON object key order', async () => { - const p = createMigratingTestPersistence() - const first = await p.stores.internalEvents!.append({ - runId: 'r1', - namespace: 'agent', - expectedSeq: 0, - type: 'step', - payload: { a: 1, b: 2 }, - }) - const retry = await p.stores.internalEvents!.append({ - runId: 'r1', - namespace: 'agent', - expectedSeq: 0, - type: 'step', - payload: { b: 2, a: 1 }, - }) - - expect(retry).toEqual(first) - }) - - it('internal event append normalizes insert races to CAS outcomes', async () => { - const idempotent = createSqlPersistence( - createInsertRaceDriver({ - table: 'internal_events', - persistedPayload: { n: 1 }, - }), - { migrate: true }, - ) - - await expect( - idempotent.stores.internalEvents!.append({ - runId: 'r1', - namespace: 'agent', - expectedSeq: 0, - type: 'step', - payload: { n: 1 }, - }), - ).resolves.toMatchObject({ - seq: 1, - namespace: 'agent', - type: 'step', - payload: { n: 1 }, - }) - - const conflicting = createSqlPersistence( - createInsertRaceDriver({ - table: 'internal_events', - persistedPayload: { n: 2 }, - }), - { migrate: true }, - ) - await expect( - conflicting.stores.internalEvents!.append({ - runId: 'r1', - namespace: 'agent', - expectedSeq: 0, - type: 'step', - payload: { n: 1 }, - }), - ).rejects.toBeInstanceOf(AppendConflictError) - }) - - it('round-trips the thread transcript', async () => { - const p = createMigratingTestPersistence() - expect(await p.messages!.loadThread('t1')).toEqual([]) - await p.messages!.saveThread('t1', [{ role: 'user', content: 'hi' }]) - await p.messages!.saveThread('t1', [ - { role: 'user', content: 'hi' }, - { role: 'assistant', content: 'hello' }, - ]) - expect(await p.messages!.loadThread('t1')).toEqual([ - { role: 'user', content: 'hi' }, - { role: 'assistant', content: 'hello' }, - ]) - }) - - it('persists and resolves approvals with thread decisions', async () => { - const p = createMigratingTestPersistence() - await p.approvals!.create({ - approvalId: 'a1', - runId: 'r1', - threadId: 't1', - status: 'pending', - requestedAt: 1, - payload: { command: 'rm' }, - }) - await p.approvals!.resolve('a1', true) - expect((await p.approvals!.get('a1'))?.status).toBe('granted') - expect((await p.approvals!.decisionsForThread('t1')).get('a1')).toBe(true) - }) - - it('passes SQL approval compatibility store into the deprecated approval controller', async () => { - const p = createMigratingTestPersistence() - const controller = createApprovalController({ store: p.approvals }) - - await controller.request({ - approvalId: 'a1', - runId: 'r1', - threadId: 't1', - requestedAt: 1, - payload: { command: 'rm' }, - }) - await controller.resolve('a1', false) - - expect((await controller.decisionsForThread('t1')).get('a1')).toBe(false) - }) - - it('keeps generic interrupts out of legacy approval decisions and maps responses', async () => { - const p = createMigratingTestPersistence() - await p.stores.interrupts!.create({ - interruptId: 'interrupt-1', - runId: 'r1', - threadId: 't1', - status: 'pending', - requestedAt: 1, - payload: { kind: 'generic' }, - }) - await p.stores.interrupts!.resolve('interrupt-1', { ok: true }) - - expect((await p.stores.interrupts!.get('interrupt-1'))?.response).toEqual({ - ok: true, - }) - expect(await p.approvals!.decisionsForThread('t1')).toEqual(new Map()) - }) - - it('creates, lists, resolves, cancels, and queries blocking interrupts', async () => { - const p = createMigratingTestPersistence() - await p.stores.interrupts!.create({ - interruptId: 'i1', - runId: 'r1', - threadId: 't1', - status: 'pending', - requestedAt: 1, - payload: { kind: 'approval' }, - }) - await p.stores.interrupts!.create({ - interruptId: 'i2', - runId: 'r2', - threadId: 't1', - status: 'pending', - requestedAt: 2, - payload: { kind: 'input' }, - }) - - expect((await p.stores.interrupts!.list('t1')).map((i) => i.runId)).toEqual( - ['r1', 'r2'], - ) - expect( - (await p.stores.interrupts!.listPending('t1')).map((i) => i.interruptId), - ).toEqual(['i1', 'i2']) - - await p.stores.interrupts!.resolve('i1', { approved: true }) - await p.stores.interrupts!.cancel('i2') - expect(await p.stores.interrupts!.listPending('t1')).toEqual([]) - expect((await p.stores.interrupts!.get('i1'))?.response).toEqual({ - approved: true, - }) - expect((await p.stores.interrupts!.get('i2'))?.status).toBe('cancelled') - }) - - it('lists interrupts and pending interrupts by run', async () => { - const p = createMigratingTestPersistence() - await p.stores.interrupts!.create({ - interruptId: 'i1', - runId: 'r1', - threadId: 't1', - status: 'pending', - requestedAt: 1, - payload: { kind: 'approval' }, - }) - await p.stores.interrupts!.create({ - interruptId: 'i2', - runId: 'r1', - threadId: 't2', - status: 'pending', - requestedAt: 2, - payload: { kind: 'input' }, - }) - await p.stores.interrupts!.create({ - interruptId: 'i3', - runId: 'r2', - threadId: 't1', - status: 'pending', - requestedAt: 3, - payload: { kind: 'other' }, - }) - - await p.stores.interrupts!.cancel('i2') - - expect( - (await p.stores.interrupts!.listByRun('r1')).map( - (interrupt) => interrupt.interruptId, - ), - ).toEqual(['i1', 'i2']) - expect( - (await p.stores.interrupts!.listPendingByRun('r1')).map( - (interrupt) => interrupt.interruptId, - ), - ).toEqual(['i1']) - }) - - it('stores app-owned namespaced metadata', async () => { - const p = createMigratingTestPersistence() - await p.stores.metadata!.set('app:user:u1', 'theme', 'dark') - await p.stores.metadata!.set('app:user:u2', 'theme', 'light') - - expect(await p.stores.metadata!.get('app:user:u1', 'theme')).toBe('dark') - expect(await p.stores.metadata!.get('app:user:u2', 'theme')).toBe('light') - - await p.stores.metadata!.delete('app:user:u1', 'theme') - expect(await p.stores.metadata!.get('app:user:u1', 'theme')).toBeNull() - expect(await p.stores.metadata!.get('app:user:u2', 'theme')).toBe('light') - }) - - it('does not expose artifact compatibility as a base SQL store', () => { - const p = createMigratingTestPersistence() - expect(p.stores.artifacts).toBeUndefined() - }) -}) diff --git a/packages/ai-persistence-sql/tests/sqlite-driver.ts b/packages/ai-persistence-sql/tests/sqlite-driver.ts deleted file mode 100644 index f00695762..000000000 --- a/packages/ai-persistence-sql/tests/sqlite-driver.ts +++ /dev/null @@ -1,32 +0,0 @@ -/** - * Test-only SqlDriver over Node's built-in `node:sqlite` (Node 22+). Lets the - * shared SQL stores be runtime-verified here without depending on the - * `@tanstack/ai-persistence-sqlite` package (which productionizes this same - * adapter). Uses an in-memory database by default. - */ -import { DatabaseSync } from 'node:sqlite' -import type { SqlDriver, SqlRow } from '../src/driver' - -export function createTestSqliteDriver(path = ':memory:'): SqlDriver { - const db = new DatabaseSync(path) - const driver: SqlDriver = { - dialect: 'sqlite', - exec(sql, params = []) { - db.prepare(sql).run(...(params as Array)) - return Promise.resolve() - }, - query( - sql: string, - params: ReadonlyArray = [], - ) { - const rows = db.prepare(sql).all(...(params as Array)) - return Promise.resolve(rows as Array) - }, - // node:sqlite has no async transaction API; stores use constraints plus - // reconciliation for CAS when this is a pass-through. - transaction(fn) { - return fn(driver) - }, - } - return driver -} diff --git a/packages/ai-persistence-sql/tsconfig.json b/packages/ai-persistence-sql/tsconfig.json deleted file mode 100644 index c38689f4e..000000000 --- a/packages/ai-persistence-sql/tsconfig.json +++ /dev/null @@ -1,8 +0,0 @@ -{ - "extends": "../../tsconfig.base.json", - "compilerOptions": { - "outDir": "dist" - }, - "include": ["src", "tests"], - "exclude": ["node_modules", "dist"] -} diff --git a/packages/ai-persistence-sql/tsup.bin.config.ts b/packages/ai-persistence-sql/tsup.bin.config.ts deleted file mode 100644 index ba12a08a3..000000000 --- a/packages/ai-persistence-sql/tsup.bin.config.ts +++ /dev/null @@ -1,12 +0,0 @@ -import { defineConfig } from 'tsup' - -export default defineConfig({ - entry: { bin: 'src/bin.ts' }, - outDir: 'dist/bin', - format: ['esm'], - target: 'node20', - platform: 'node', - sourcemap: true, - clean: false, - banner: { js: '#!/usr/bin/env node' }, -}) diff --git a/packages/ai-persistence-sql/vite.config.ts b/packages/ai-persistence-sql/vite.config.ts deleted file mode 100644 index ba8f405e4..000000000 --- a/packages/ai-persistence-sql/vite.config.ts +++ /dev/null @@ -1,37 +0,0 @@ -import { defineConfig, mergeConfig } from 'vitest/config' -import { tanstackViteConfig } from '@tanstack/vite-config' -import packageJson from './package.json' - -const config = defineConfig({ - test: { - name: packageJson.name, - dir: './', - watch: false, - - globals: true, - environment: 'node', - include: ['tests/**/*.test.ts'], - coverage: { - provider: 'v8', - reporter: ['text', 'json', 'html', 'lcov'], - exclude: [ - 'node_modules/', - 'dist/', - 'tests/', - '**/*.test.ts', - '**/*.config.ts', - '**/types.ts', - ], - include: ['src/**/*.ts'], - }, - }, -}) - -export default mergeConfig( - config, - tanstackViteConfig({ - entry: ['./src/index.ts', './src/cli-core.ts'], - srcDir: './src', - cjs: false, - }), -) diff --git a/packages/ai-persistence-sqlite/package.json b/packages/ai-persistence-sqlite/package.json deleted file mode 100644 index b9b94b774..000000000 --- a/packages/ai-persistence-sqlite/package.json +++ /dev/null @@ -1,54 +0,0 @@ -{ - "name": "@tanstack/ai-persistence-sqlite", - "version": "0.1.0", - "description": "SQLite backend for TanStack AI persistence — durable runs, messages, events, interrupts, and metadata over node:sqlite (zero-dep, Node 22+) or a bring-your-own better-sqlite3 handle.", - "author": "", - "license": "MIT", - "repository": { - "type": "git", - "url": "git+https://github.com/TanStack/ai.git", - "directory": "packages/ai-persistence-sqlite" - }, - "keywords": [ - "ai", - "tanstack", - "persistence", - "sqlite", - "node-sqlite", - "better-sqlite3" - ], - "type": "module", - "module": "./dist/esm/index.js", - "types": "./dist/esm/index.d.ts", - "exports": { - ".": { - "types": "./dist/esm/index.d.ts", - "import": "./dist/esm/index.js" - } - }, - "files": [ - "dist", - "src" - ], - "scripts": { - "build": "vite build", - "clean": "premove ./build ./dist", - "lint:fix": "eslint ./src --fix", - "test:build": "publint --strict", - "test:eslint": "eslint ./src", - "test:lib": "vitest", - "test:lib:dev": "pnpm test:lib --watch", - "test:types": "tsc" - }, - "peerDependencies": { - "@tanstack/ai": "workspace:^", - "@tanstack/ai-persistence": "workspace:^", - "@tanstack/ai-persistence-sql": "workspace:^" - }, - "devDependencies": { - "@tanstack/ai": "workspace:*", - "@tanstack/ai-persistence": "workspace:*", - "@tanstack/ai-persistence-sql": "workspace:*", - "@vitest/coverage-v8": "4.0.14" - } -} diff --git a/packages/ai-persistence-sqlite/src/index.ts b/packages/ai-persistence-sqlite/src/index.ts deleted file mode 100644 index 538c2e692..000000000 --- a/packages/ai-persistence-sqlite/src/index.ts +++ /dev/null @@ -1,95 +0,0 @@ -/** - * SQLite backend. Wraps a SQLite client in the shared `SqlDriver` and assembles - * a `SqlPersistence` via `@tanstack/ai-persistence-sql`. - * - * Two client paths: - * - convenience: `sqlitePersistence({ path })` lazily opens a `node:sqlite` - * `DatabaseSync` (built in to Node 22+, zero dependency). - * - BYO: pass `{ db }` — an existing `node:sqlite` `DatabaseSync` OR a - * `better-sqlite3` `Database` (both expose `prepare(sql).run/all(...params)`). - */ -import { createRequire } from 'node:module' -import { createSqlPersistence } from '@tanstack/ai-persistence-sql' -import type { - SqlDriver, - SqlPersistence, - SqlRow, -} from '@tanstack/ai-persistence-sql' -import type { PersistenceMode } from '@tanstack/ai-persistence' - -/** The subset of a SQLite client the driver needs (node:sqlite & better-sqlite3 both satisfy it). */ -interface SqliteClient { - prepare: (sql: string) => { - run: (...params: Array) => unknown - all: (...params: Array) => Array - } -} - -export interface SqliteDriverOptions { - /** File path (or ':memory:'). Used when `db` is not provided. */ - path?: string - /** Bring-your-own SQLite handle (node:sqlite DatabaseSync or better-sqlite3 Database). */ - db?: SqliteClient -} - -type DatabaseSyncConstructor = new (path: string) => SqliteClient - -function loadNodeSqliteDatabaseSync(): DatabaseSyncConstructor { - try { - const require = createRequire(import.meta.url) - const sqlite = require('node:sqlite') as { - DatabaseSync?: DatabaseSyncConstructor - } - if (typeof sqlite.DatabaseSync === 'function') { - return sqlite.DatabaseSync - } - } catch { - // Throw the package-level guidance below. - } - throw new Error( - 'sqlitePersistence({ path }) requires node:sqlite, available in Node 22+. Pass sqlitePersistence({ db }) with a compatible SQLite handle on older Node versions.', - ) -} - -/** Build a {@link SqlDriver} backed by SQLite. */ -export function createSqliteDriver(opts?: SqliteDriverOptions): SqlDriver { - const client: SqliteClient = - opts?.db ?? new (loadNodeSqliteDatabaseSync())(opts?.path ?? ':memory:') - - const driver: SqlDriver = { - dialect: 'sqlite', - exec(sql, params = []) { - client.prepare(sql).run(...params) - return Promise.resolve() - }, - query( - sql: string, - params: ReadonlyArray = [], - ) { - return Promise.resolve(client.prepare(sql).all(...params) as Array) - }, - // SQLite clients here are synchronous; statements already run in order, so a - // plain pass-through is sufficient (no async interleaving to isolate). - transaction(fn) { - return fn(driver) - }, - } - return driver -} - -export interface SqlitePersistenceOptions extends SqliteDriverOptions { - mode?: PersistenceMode - /** Run migrations on first use (default false). */ - migrate?: boolean -} - -/** SQLite-backed {@link SqlPersistence}. */ -export function sqlitePersistence( - opts?: SqlitePersistenceOptions, -): SqlPersistence { - const driver = createSqliteDriver({ path: opts?.path, db: opts?.db }) - return createSqlPersistence(driver, { - mode: opts?.mode, - migrate: opts?.migrate, - }) -} diff --git a/packages/ai-persistence-sqlite/tests/sqlite.test.ts b/packages/ai-persistence-sqlite/tests/sqlite.test.ts deleted file mode 100644 index 592ddd581..000000000 --- a/packages/ai-persistence-sqlite/tests/sqlite.test.ts +++ /dev/null @@ -1,78 +0,0 @@ -import { describe, expect, it } from 'vitest' -import { EventType } from '@tanstack/ai' -import type { StreamChunk } from '@tanstack/ai' -import { createResumeSource } from '@tanstack/ai-persistence' -import { createSqliteDriver, sqlitePersistence } from '../src/index' - -const text = (delta: string): StreamChunk => ({ - type: EventType.TEXT_MESSAGE_CONTENT, - messageId: 'm1', - delta, - timestamp: 1, -}) - -describe('sqlitePersistence', () => { - it('uses a supplied db handle without opening node:sqlite', async () => { - const calls: Array<{ sql: string; params: Array }> = [] - const db = { - prepare: (sql: string) => ({ - run: (...params: Array) => { - calls.push({ sql, params }) - }, - all: (...params: Array) => { - calls.push({ sql, params }) - return [{ value: 'ok' }] - }, - }), - } - - const driver = createSqliteDriver({ db }) - await driver.exec('insert into test values (?)', ['a']) - const rows = await driver.query('select value from test where id = ?', [1]) - - expect(rows).toEqual([{ value: 'ok' }]) - expect(calls).toEqual([ - { sql: 'insert into test values (?)', params: ['a'] }, - { sql: 'select value from test where id = ?', params: [1] }, - ]) - }) - - it('does not create schema by default', async () => { - const p = sqlitePersistence() - - await expect( - p.runs!.createOrResume({ runId: 'r1', threadId: 't1', startedAt: 1 }), - ).rejects.toThrow(/no such table: runs/) - }) - - it('round-trips a run, events, and transcript on an in-memory db', async () => { - const p = sqlitePersistence({ migrate: true }) - await p.runs!.createOrResume({ runId: 'r1', threadId: 't1', startedAt: 1 }) - await p.events!.append('r1', 1, text('a')) - await p.events!.append('r1', 2, text('b')) - await p.messages!.saveThread('t1', [{ role: 'user', content: 'hi' }]) - - expect(await p.events!.latestSeq('r1')).toBe(2) - expect((await p.runs!.get('r1'))?.status).toBe('running') - expect(await p.messages!.loadThread('t1')).toEqual([ - { role: 'user', content: 'hi' }, - ]) - }) - - it('drives the core ResumeSource: replays the tail after a cursor', async () => { - const p = sqlitePersistence({ migrate: true }) - await p.runs!.createOrResume({ runId: 'r1', threadId: 't1', startedAt: 1 }) - await p.events!.append('r1', 1, text('a')) - await p.events!.append('r1', 2, text('b')) - await p.events!.append('r1', 3, text('c')) - - const source = createResumeSource(p.events!, p.runs) - expect(await source.hasRun('r1')).toBe(true) - - const deltas: Array = [] - for await (const chunk of source.replay('r1', undefined)) { - if (chunk.type === 'TEXT_MESSAGE_CONTENT') deltas.push(chunk.delta) - } - expect(deltas).toEqual(['a', 'b', 'c']) - }) -}) diff --git a/packages/ai-persistence-sqlite/tsconfig.json b/packages/ai-persistence-sqlite/tsconfig.json deleted file mode 100644 index c38689f4e..000000000 --- a/packages/ai-persistence-sqlite/tsconfig.json +++ /dev/null @@ -1,8 +0,0 @@ -{ - "extends": "../../tsconfig.base.json", - "compilerOptions": { - "outDir": "dist" - }, - "include": ["src", "tests"], - "exclude": ["node_modules", "dist"] -} diff --git a/packages/ai-persistence-sqlite/vite.config.ts b/packages/ai-persistence-sqlite/vite.config.ts deleted file mode 100644 index 11f5b20b7..000000000 --- a/packages/ai-persistence-sqlite/vite.config.ts +++ /dev/null @@ -1,37 +0,0 @@ -import { defineConfig, mergeConfig } from 'vitest/config' -import { tanstackViteConfig } from '@tanstack/vite-config' -import packageJson from './package.json' - -const config = defineConfig({ - test: { - name: packageJson.name, - dir: './', - watch: false, - - globals: true, - environment: 'node', - include: ['tests/**/*.test.ts'], - coverage: { - provider: 'v8', - reporter: ['text', 'json', 'html', 'lcov'], - exclude: [ - 'node_modules/', - 'dist/', - 'tests/', - '**/*.test.ts', - '**/*.config.ts', - '**/types.ts', - ], - include: ['src/**/*.ts'], - }, - }, -}) - -export default mergeConfig( - config, - tanstackViteConfig({ - entry: ['./src/index.ts'], - srcDir: './src', - cjs: false, - }), -) From 369965667b0ff9f4202b4083ebd849009dc3a33f Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Thu, 9 Jul 2026 18:52:26 +0200 Subject: [PATCH 5/9] feat(client, frameworks)!: resumable SSE + interrupt-only resume; drop cursor & generation delivery-resume --- packages/ai-angular/src/inject-chat.ts | 25 +- .../ai-angular/src/inject-generate-audio.ts | 2 +- .../ai-angular/src/inject-generate-video.ts | 52 +- packages/ai-angular/src/inject-generation.ts | 61 +- packages/ai-angular/tests/inject-chat.test.ts | 52 +- .../tests/inject-generation.test.ts | 141 +--- packages/ai-client/src/chat-client.ts | 258 +++--- packages/ai-client/src/connection-adapters.ts | 212 ++++- packages/ai-client/src/generation-client.ts | 80 +- packages/ai-client/src/generation-types.ts | 27 +- packages/ai-client/src/index.ts | 2 + packages/ai-client/src/types.ts | 15 +- .../ai-client/src/video-generation-client.ts | 79 +- .../tests/chat-client-resume.test.ts | 760 ++++-------------- .../connection-adapters-resumable.test.ts | 155 ++++ .../tests/connection-adapters-xhr.test.ts | 4 +- .../tests/connection-adapters.test.ts | 10 +- .../ai-client/tests/generation-client.test.ts | 36 - .../tests/generation-resume-state.test.ts | 16 +- packages/ai-preact/src/types.ts | 1 - packages/ai-preact/src/use-chat.ts | 78 +- packages/ai-preact/tests/use-chat.test.ts | 130 +-- packages/ai-react/src/types.ts | 1 - packages/ai-react/src/use-chat.ts | 78 +- packages/ai-react/src/use-generate-audio.ts | 8 +- packages/ai-react/src/use-generate-image.ts | 8 +- packages/ai-react/src/use-generate-speech.ts | 8 +- packages/ai-react/src/use-generate-video.ts | 45 +- packages/ai-react/src/use-generation.ts | 49 +- packages/ai-react/src/use-summarize.ts | 8 +- packages/ai-react/src/use-transcription.ts | 8 +- packages/ai-react/tests/use-chat.test.ts | 236 +----- .../ai-react/tests/use-generation.test.ts | 288 +------ packages/ai-solid/src/types.ts | 1 - packages/ai-solid/src/use-chat.ts | 48 +- packages/ai-solid/src/use-generate-audio.ts | 2 +- packages/ai-solid/src/use-generate-image.ts | 2 +- packages/ai-solid/src/use-generate-speech.ts | 2 +- packages/ai-solid/src/use-generate-video.ts | 46 +- packages/ai-solid/src/use-generation.ts | 46 +- packages/ai-solid/src/use-summarize.ts | 2 +- packages/ai-solid/src/use-transcription.ts | 2 +- .../ai-solid/tests/use-generation.test.ts | 42 +- packages/ai-svelte/src/create-chat.svelte.ts | 47 +- .../src/create-generate-audio.svelte.ts | 3 +- .../src/create-generate-image.svelte.ts | 3 +- .../src/create-generate-speech.svelte.ts | 3 +- .../src/create-generate-video.svelte.ts | 40 +- .../ai-svelte/src/create-generation.svelte.ts | 40 +- .../ai-svelte/src/create-summarize.svelte.ts | 3 +- .../src/create-transcription.svelte.ts | 3 +- packages/ai-svelte/src/types.ts | 1 - .../ai-svelte/tests/create-generation.test.ts | 136 +--- packages/ai-svelte/tests/use-chat.test.ts | 65 -- packages/ai-vue/src/types.ts | 1 - packages/ai-vue/src/use-chat.ts | 46 +- packages/ai-vue/src/use-generate-audio.ts | 2 +- packages/ai-vue/src/use-generate-image.ts | 2 +- packages/ai-vue/src/use-generate-speech.ts | 2 +- packages/ai-vue/src/use-generate-video.ts | 49 +- packages/ai-vue/src/use-generation.ts | 49 +- packages/ai-vue/src/use-summarize.ts | 2 +- packages/ai-vue/src/use-transcription.ts | 2 +- packages/ai-vue/tests/use-generation.test.ts | 180 +---- 64 files changed, 895 insertions(+), 2910 deletions(-) create mode 100644 packages/ai-client/tests/connection-adapters-resumable.test.ts diff --git a/packages/ai-angular/src/inject-chat.ts b/packages/ai-angular/src/inject-chat.ts index a61666a72..70c961681 100644 --- a/packages/ai-angular/src/inject-chat.ts +++ b/packages/ai-angular/src/inject-chat.ts @@ -80,14 +80,6 @@ export function injectChat< ? { connection: options.connection } : { fetcher: options.fetcher } - let lifecycleOwned = false - const resumeIfLifecycleOwned = () => { - if (!lifecycleOwned) { - return - } - void client.maybeAutoResume() - } - const client = new ChatClient({ devtoolsBridgeFactory: createChatDevtoolsBridge, ...transport, @@ -119,7 +111,6 @@ export function injectChat< onError: (err) => options.onError?.(err), onResumeStateChange: (resumeState, pendingInterrupts) => { options.onResumeStateChange?.(resumeState, pendingInterrupts) - resumeIfLifecycleOwned() }, tools: options.tools, onCustomEvent: (eventType, data, context) => @@ -173,25 +164,15 @@ export function injectChat< afterNextRender( () => { - lifecycleOwned = true client.mountDevtools() - void client.maybeAutoResume() - if (typeof window !== 'undefined') { - window.addEventListener('online', handleOnline) - } + // Delivery-durability resume is transparent: the resumable SSE + // connection adapter reattaches via the browser's native Last-Event-ID + // on reconnect. No client-side auto-resume wiring is needed. }, { injector }, ) - const handleOnline = () => { - void client.maybeAutoResume() - } - destroyRef.onDestroy(() => { - lifecycleOwned = false - if (typeof window !== 'undefined') { - window.removeEventListener('online', handleOnline) - } if (liveSource?.()) { client.unsubscribe() } else { diff --git a/packages/ai-angular/src/inject-generate-audio.ts b/packages/ai-angular/src/inject-generate-audio.ts index 302fa3424..fbb57607c 100644 --- a/packages/ai-angular/src/inject-generate-audio.ts +++ b/packages/ai-angular/src/inject-generate-audio.ts @@ -24,7 +24,7 @@ export interface InjectGenerateAudioOptions< TOutput = AudioGenerationResult, > extends Pick< InjectGenerationOptions, - 'persistence' | 'autoResume' | 'initialResumeSnapshot' | 'resumeState' + 'persistence' | 'initialResumeSnapshot' > { /** Connect-based adapter for streaming transport (SSE, HTTP stream, custom) */ connection?: ConnectConnectionAdapter diff --git a/packages/ai-angular/src/inject-generate-video.ts b/packages/ai-angular/src/inject-generate-video.ts index 0e7801d05..4e75df991 100644 --- a/packages/ai-angular/src/inject-generate-video.ts +++ b/packages/ai-angular/src/inject-generate-video.ts @@ -38,9 +38,7 @@ export interface InjectGenerateVideoOptions { body?: ReactiveOption> devtools?: AIDevtoolsDisplayOptions persistence?: GenerationPersistenceOptions - autoResume?: boolean initialResumeSnapshot?: GenerationResumeSnapshot - resumeState?: ReactiveOption onResult?: (result: VideoGenerateResult) => TOutput | null | void onError?: (error: Error) => void onProgress?: (progress: number, message?: string) => void @@ -63,7 +61,6 @@ export interface InjectGenerateVideoResult { resumeState: Signal pendingArtifacts: Signal> resultArtifacts: Signal> - resume: (state?: GenerationResumeState) => Promise } // `TTransformed` infers from the `onResult` return position so the callback @@ -119,10 +116,6 @@ export function injectGenerateVideo( const bodySource = options.body !== undefined ? toReactive(options.body) : undefined - const resumeStateSource = - options.resumeState !== undefined - ? toReactive(options.resumeState) - : undefined const baseOptions = { id: clientId, @@ -130,13 +123,9 @@ export function injectGenerateVideo( ...(options.persistence !== undefined && { persistence: options.persistence, }), - ...(options.autoResume !== undefined && { autoResume: options.autoResume }), ...(options.initialResumeSnapshot !== undefined && { initialResumeSnapshot: options.initialResumeSnapshot, }), - ...(resumeStateSource?.() !== undefined && { - resumeState: resumeStateSource(), - }), devtoolsBridgeFactory: createVideoDevtoolsBridge, devtools: { ...options.devtools, @@ -208,45 +197,17 @@ export function injectGenerateVideo( () => { client.updateOptions({ body: bodySource(), - ...(resumeStateSource?.() !== undefined && { - resumeState: resumeStateSource(), - }), - }) - }, - { injector }, - ) - } - - if (resumeStateSource && !bodySource) { - effect( - () => { - const nextResumeState = resumeStateSource() - client.updateOptions({ - ...(nextResumeState !== undefined && { - resumeState: nextResumeState, - }), }) }, { injector }, ) } + // Mount devtools only. Generation runs are never auto-started after render — + // persisted state is read-only for display. afterNextRender( () => { client.mountDevtools() - void client - .maybeAutoResume() - .catch((err: unknown) => { - if (disposed) return - const nextError = err instanceof Error ? err : new Error(String(err)) - options.onError?.(nextError) - error.set(nextError) - status.set('error') - }) - .finally(() => { - if (disposed) return - setResumeSnapshotState(client.getResumeSnapshot()) - }) }, { injector }, ) @@ -255,14 +216,6 @@ export function injectGenerateVideo( client.dispose() }) - const resume = async (state?: GenerationResumeState) => { - const didResume = await client.resume(state) - if (!disposed) { - setResumeSnapshotState(client.getResumeSnapshot()) - } - return didResume - } - return { generate: (input: VideoGenerateInput) => client.generate(input), result: result.asReadonly(), @@ -277,6 +230,5 @@ export function injectGenerateVideo( resumeState: resumeState.asReadonly(), pendingArtifacts: pendingArtifacts.asReadonly(), resultArtifacts: resultArtifacts.asReadonly(), - resume, } } diff --git a/packages/ai-angular/src/inject-generation.ts b/packages/ai-angular/src/inject-generation.ts index 48d4edbe7..542f2c6e6 100644 --- a/packages/ai-angular/src/inject-generation.ts +++ b/packages/ai-angular/src/inject-generation.ts @@ -40,14 +40,10 @@ export interface InjectGenerationOptions { body?: ReactiveOption> /** Display options for TanStack AI Devtools. */ devtools?: AIDevtoolsDisplayOptions - /** Server-side lightweight resume state persistence. */ + /** Server-side lightweight generation state persistence. */ persistence?: GenerationPersistenceOptions - /** Whether to resume a persisted run after render. Defaults to true. */ - autoResume?: boolean - /** Initial lightweight resume snapshot restored by the app. */ + /** Initial lightweight resume snapshot restored by the app (read-only state). */ initialResumeSnapshot?: GenerationResumeSnapshot - /** Explicit run/cursor state to use for the next resume/generation request. */ - resumeState?: ReactiveOption /** * Callback when a result is received. Can optionally return a transformed value. * @@ -81,14 +77,12 @@ export interface InjectGenerationResult { reset: () => void /** Lightweight generation resume snapshot, if one is available */ resumeSnapshot: Signal - /** Current resumable run/cursor state, if one is available */ + /** 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> - /** Resume the current/initial generation run, if resumable */ - resume: (state?: GenerationResumeState) => Promise } // `TTransformed` infers from the `onResult` return position (a covariant @@ -146,10 +140,6 @@ export function injectGeneration< const bodySource = options.body !== undefined ? toReactive(options.body) : undefined - const resumeStateSource = - options.resumeState !== undefined - ? toReactive(options.resumeState) - : undefined const clientOptions: GenerationClientOptions = { id: clientId, @@ -157,13 +147,9 @@ export function injectGeneration< ...(options.persistence !== undefined && { persistence: options.persistence, }), - ...(options.autoResume !== undefined && { autoResume: options.autoResume }), ...(options.initialResumeSnapshot !== undefined && { initialResumeSnapshot: options.initialResumeSnapshot, }), - ...(resumeStateSource?.() !== undefined && { - resumeState: resumeStateSource(), - }), devtoolsBridgeFactory: createGenerationDevtoolsBridge, devtools: { ...options.devtools, @@ -222,45 +208,17 @@ export function injectGeneration< () => { client.updateOptions({ body: bodySource(), - ...(resumeStateSource?.() !== undefined && { - resumeState: resumeStateSource(), - }), - }) - }, - { injector }, - ) - } - - if (resumeStateSource && !bodySource) { - effect( - () => { - const nextResumeState = resumeStateSource() - client.updateOptions({ - ...(nextResumeState !== undefined && { - resumeState: nextResumeState, - }), }) }, { injector }, ) } + // Mount devtools only. Generation runs are never auto-started after render — + // persisted state is read-only for display. afterNextRender( () => { client.mountDevtools() - void client - .maybeAutoResume() - .catch((err: unknown) => { - if (disposed) return - const nextError = err instanceof Error ? err : new Error(String(err)) - options.onError?.(nextError) - error.set(nextError) - status.set('error') - }) - .finally(() => { - if (disposed) return - setResumeSnapshotState(client.getResumeSnapshot()) - }) }, { injector }, ) @@ -269,14 +227,6 @@ export function injectGeneration< client.dispose() }) - const resume = async (state?: GenerationResumeState) => { - const didResume = await client.resume(state) - if (!disposed) { - setResumeSnapshotState(client.getResumeSnapshot()) - } - return didResume - } - return { generate: ((input: TInput) => client.generate(input)) as ( input: Record, @@ -291,6 +241,5 @@ export function injectGeneration< resumeState: resumeState.asReadonly(), pendingArtifacts: pendingArtifacts.asReadonly(), resultArtifacts: resultArtifacts.asReadonly(), - resume, } } diff --git a/packages/ai-angular/tests/inject-chat.test.ts b/packages/ai-angular/tests/inject-chat.test.ts index 8f51a3729..6855074af 100644 --- a/packages/ai-angular/tests/inject-chat.test.ts +++ b/packages/ai-angular/tests/inject-chat.test.ts @@ -3,7 +3,7 @@ import { z } from 'zod' import { EventType } from '@tanstack/ai/client' import { Component, signal } from '@angular/core' import { TestBed } from '@angular/core/testing' -import { ChatClient, type ConnectionAdapter } from '@tanstack/ai-client' +import { ChatClient } from '@tanstack/ai-client' import { injectChat } from '../src/inject-chat' import { createMockConnectionAdapter, @@ -137,14 +137,16 @@ describe('injectChat — resume', () => { messageId: 'msg-1', timestamp: Date.now(), delta: 'Hi', - cursor: 'cursor-1', }, { type: EventType.RUN_FINISHED, runId: 'run-1', threadId: 'thread-1', timestamp: Date.now(), - outcome: { type: 'success' }, + outcome: { + type: 'interrupt', + interrupts: [{ id: 'interrupt-1', reason: 'client_tool_input' }], + }, }, ], }) @@ -156,54 +158,16 @@ describe('injectChat — resume', () => { await result.sendMessage('Hi') + // A run that pauses on an interrupt forwards interrupt (state) resume — + // the thread/run ids to target on a follow-up. No delivery cursor. expect(onResumeStateChange).toHaveBeenCalledWith( expect.objectContaining({ threadId: 'thread-1', runId: expect.any(String), - cursor: 'cursor-1', }), - [], + expect.arrayContaining([expect.objectContaining({ id: 'interrupt-1' })]), ) }) - - it('forwards initialResumeSnapshot to ChatClient auto-resume', async () => { - type RunContext = Parameters< - Extract< - ConnectionAdapter, - { - connect: unknown - } - >['connect'] - >[3] - const contexts: Array = [] - const adapter: Extract = { - async *connect(_messages, _data, _signal, runContext) { - contexts.push(runContext) - }, - } - const { flush } = renderInjectChat({ - connection: adapter, - initialResumeSnapshot: { - resumeState: { - threadId: 'thread-1', - runId: 'run-1', - cursor: 'cursor-1', - }, - pendingInterrupts: [], - }, - }) - - flush() - await tick() - - expect(contexts).toEqual([ - expect.objectContaining({ - threadId: 'thread-1', - runId: 'run-1', - cursor: 'cursor-1', - }), - ]) - }) }) describe('injectChat — structured output', () => { diff --git a/packages/ai-angular/tests/inject-generation.test.ts b/packages/ai-angular/tests/inject-generation.test.ts index 8d1433a6e..2ee11685c 100644 --- a/packages/ai-angular/tests/inject-generation.test.ts +++ b/packages/ai-angular/tests/inject-generation.test.ts @@ -65,51 +65,10 @@ const videoResumeSnapshot: GenerationResumeSnapshot = { resumeState: { threadId: 'thread-resume', runId: 'run-resume', - cursor: 'cursor-resume', }, status: 'running', } -function createDeferred() { - let resolve!: (value: T) => void - let reject!: (reason?: unknown) => void - const promise = new Promise((res, rej) => { - resolve = res - reject = rej - }) - return { promise, resolve, reject } -} - -function createReplayVideoChunks(): Array { - return [ - { - type: 'RUN_STARTED', - runId: 'run-resume', - threadId: 'thread-resume', - cursor: 'cursor-start', - timestamp: Date.now(), - }, - { - type: 'CUSTOM', - name: 'generation:result', - value: { - jobId: 'job-replay', - status: 'completed', - url: 'https://example.com/video.mp4', - }, - cursor: 'cursor-result', - timestamp: Date.now(), - }, - { - type: 'RUN_FINISHED', - runId: 'run-resume', - threadId: 'thread-resume', - cursor: 'cursor-finished', - timestamp: Date.now(), - }, - ] as unknown as Array -} - function createRunContextCaptureAdapter(chunks: Array): { adapter: ConnectConnectionAdapter connect: ReturnType @@ -161,92 +120,56 @@ describe('injectGeneration', () => { expect(result.status()).toBe('success') }) - it('ignores auto-resume rejection after destroy', async () => { - const deferred = createDeferred() - const getItem = vi.fn(() => deferred.promise) - const onError = vi.fn() - const { result, destroy } = renderInjectGeneration({ - fetcher: async () => ({ value: 1 }), + 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: { - server: { - getItem, - setItem: vi.fn(), - removeItem: vi.fn(), - }, + server: { getItem, setItem: vi.fn(), removeItem: vi.fn() }, }, - onError, + initialResumeSnapshot: snapshot, }) - await vi.waitFor(() => expect(getItem).toHaveBeenCalled()) - destroy() - deferred.reject(new Error('resume failed')) - await deferred.promise.catch(() => {}) await Promise.resolve() - expect(onError).not.toHaveBeenCalled() - expect(result.error()).toBeUndefined() + 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('explicitly resumes from the current snapshot', async () => { - const { adapter, connect, runContexts } = createRunContextCaptureAdapter( - createReplayVideoChunks(), - ) - const { result, flush } = renderInjectGenerateVideo({ + 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, - initialResumeSnapshot: videoResumeSnapshot, - autoResume: false, - }) - - const didResume = await result.resume() - flush() - - expect(didResume).toBe(true) - expect(connect).toHaveBeenCalledTimes(1) - expect(runContexts[0]).toEqual(videoResumeSnapshot.resumeState) - expect(result.resumeSnapshot()).toEqual( - expect.objectContaining({ - status: 'complete', - resumeState: null, - }), - ) - expect(result.result()).toEqual( - expect.objectContaining({ - jobId: 'job-replay', - }), - ) - }) - - it('ignores video auto-resume rejection after destroy', async () => { - const deferred = createDeferred() - const getItem = vi.fn(() => deferred.promise) - const onError = vi.fn() - const { result, destroy } = renderInjectGenerateVideo({ - fetcher: async () => ({ - jobId: 'job-1', - status: 'completed', - url: 'https://example.com/video.mp4', - }), persistence: { - server: { - getItem, - setItem: vi.fn(), - removeItem: vi.fn(), - }, + server: { getItem, setItem: vi.fn(), removeItem: vi.fn() }, }, - onError, + initialResumeSnapshot: videoResumeSnapshot, }) - await vi.waitFor(() => expect(getItem).toHaveBeenCalled()) - destroy() - deferred.reject(new Error('video resume failed')) - await deferred.promise.catch(() => {}) await Promise.resolve() - expect(onError).not.toHaveBeenCalled() - expect(result.error()).toBeUndefined() + 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/chat-client.ts b/packages/ai-client/src/chat-client.ts index c088196c5..440668df4 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -129,10 +129,16 @@ export class ChatClient< private readonly persistor?: ChatPersistor private readonly serverPersistence?: ChatServerPersistence private serverPersistenceGeneration = 0 + // Whether the consumer supplied an `onError` callback. Server-persistence is + // best-effort and must never break chat, but its failures should still be + // observable: they are routed to the consumer's `onError` when one exists, + // otherwise surfaced via `console.warn`. + private hasUserOnError = false private disposed = false private currentRunId: string | null = null - // Resume tracking: the latest in-band cursor seen for the active run, so a - // reconnect can replay events after it. Cleared when the run terminates. + // Interrupt-resume tracking: the run/thread of the most recent interrupted + // run, so approvals/client-tool results can be sent back. Cleared when the + // run terminates. This is STATE (interrupt) resume, not delivery/cursor. private lastResume: ChatResumeState | null = null private pendingInterrupts: Array = [] private pendingInterruptRunId: string | null = null @@ -140,12 +146,10 @@ export class ChatClient< string, RunAgentResumeItem >() - private readonly autoResume: boolean - // When set, the next streamResponse() resumes this run/cursor instead of - // starting a fresh run (consumed once). + // When set, the next streamResponse() continues this interrupted run instead + // of starting a fresh run (consumed once). private pendingResumeRunId: string | null = null private pendingResumeThreadId: string | null = null - private pendingResumeCursor: string | null = null private pendingResumeItems: Array | null = null private activeResumeThreadId: string | null = null private activeResumeRunId: string | null = null @@ -223,7 +227,6 @@ export class ChatClient< constructor(options: ChatClientOptions) { this.uniqueId = options.id || this.generateUniqueId('chat') this.threadId = options.threadId || this.generateUniqueId('thread') - this.autoResume = options.autoResume ?? true const persistence = normalizePersistence(options.persistence) if (persistence.client) { this.persistor = new ChatPersistor( @@ -233,6 +236,7 @@ export class ChatClient< ) } this.serverPersistence = persistence.server + this.hasUserOnError = typeof options.onError === 'function' // Both `body` (deprecated) and `forwardedProps` populate the AG-UI // `RunAgentInput.forwardedProps` wire field. They are stored // separately so `updateOptions` can replace one without touching the @@ -506,6 +510,26 @@ export class ChatClient< } } + /** + * Route a caught server-persistence error to an observable sink. Server + * persistence is best-effort and must never break chat, so callers still + * swallow the failure — but it is no longer silent: it goes to the + * consumer's `onError` callback when one was provided, otherwise to + * `console.warn`. + */ + private reportServerPersistenceError(error: unknown): void { + const normalized = + error instanceof Error ? error : new Error(String(error)) + if (this.hasUserOnError) { + this.callbacksRef.current.onError(normalized) + } else { + console.warn( + '[TanStack AI] Server persistence adapter error (non-fatal):', + normalized, + ) + } + } + private readInitialResumeSnapshot(): | ChatResumeSnapshot | null @@ -513,7 +537,8 @@ export class ChatClient< | Promise { try { return this.serverPersistence?.getItem(this.threadId) - } catch { + } catch (error) { + this.reportServerPersistenceError(error) return undefined } } @@ -534,8 +559,10 @@ export class ChatClient< this.notifyResumeStateChange() } }) - .catch(() => { - // Persistence adapters are best-effort and must not break chat setup. + .catch((error: unknown) => { + // Persistence adapters are best-effort and must not break chat setup, + // but the failure is surfaced to an observable sink. + this.reportServerPersistenceError(error) }) } @@ -629,109 +656,69 @@ export class ChatClient< } /** - * Observe the in-band resume cursor on each chunk so a reconnect can replay - * after the last seen event. Cleared when the run reaches a terminal event. + * Track interrupt state off the stream's terminal events. A RUN_FINISHED with + * an interrupt outcome records the pending interrupts + the run/thread to + * resume; any other terminal event for the tracked/current run clears that + * state. This is interrupt (state) resume — there is no delivery cursor. */ - private observeResumeCursor(chunk: StreamChunk): void { - if (chunk.type === 'RUN_FINISHED' || chunk.type === 'RUN_ERROR') { - // A server-signaled terminal event completes the run — drop its resume - // state. (A stream that merely ends without a terminal is an interruption - // and keeps its resume state so it can be continued.) - const runId = getChunkRunId(chunk) - const threadId = - 'threadId' in chunk && typeof chunk.threadId === 'string' - ? chunk.threadId - : this.activeResumeThreadId - const cursor = - 'cursor' in chunk && typeof chunk.cursor === 'string' - ? chunk.cursor - : undefined - if ( - chunk.type === 'RUN_FINISHED' && - chunk.outcome?.type === 'interrupt' - ) { - const interruptedRunId = - this.currentRunId && - (cursor || this.lastResume?.runId === this.currentRunId) - ? this.currentRunId - : (runId ?? this.activeResumeRunId ?? this.currentRunId ?? '') - const resumeTarget = - this.lastResume?.runId === interruptedRunId ? this.lastResume : null - if (cursor) { - this.lastResume = { - threadId: threadId ?? this.threadId, - runId: interruptedRunId, - cursor, - } - } else if (!resumeTarget) { - this.pendingInterrupts = [] - this.pendingInterruptRunId = null - this.pendingInterruptResumeItems.clear() - this.notifyResumeStateChange() - return - } - this.pendingInterruptRunId = interruptedRunId - this.pendingInterrupts = [...chunk.outcome.interrupts] - this.pendingInterruptResumeItems.clear() - this.notifyResumeStateChange() - return - } - - const isRunlessSessionError = chunk.type === 'RUN_ERROR' && !runId - const isTrackedRunTerminal = Boolean( - runId && this.lastResume?.runId === runId, - ) - const isPendingInterruptRunTerminal = Boolean( - runId && this.pendingInterruptRunId === runId, - ) - const isCurrentRunTerminal = Boolean( - (runId && this.currentRunId === runId) || - (this.currentRunId && this.lastResume?.runId === this.currentRunId), - ) - const isCurrentStreamTerminal = - this.isLoading && chunk.type === 'RUN_FINISHED' && !runId - if ( - isRunlessSessionError || - isTrackedRunTerminal || - isPendingInterruptRunTerminal || - isCurrentRunTerminal || - isCurrentStreamTerminal - ) { - this.lastResume = null - } - if ( - isRunlessSessionError || - isTrackedRunTerminal || - isPendingInterruptRunTerminal || - isCurrentRunTerminal || - isCurrentStreamTerminal - ) { - this.pendingInterrupts = [] - this.pendingInterruptRunId = null - this.pendingInterruptResumeItems.clear() - } - this.notifyResumeStateChange() + private observeInterruptState(chunk: StreamChunk): void { + if (chunk.type !== 'RUN_FINISHED' && chunk.type !== 'RUN_ERROR') { return } - const cursor = - 'cursor' in chunk && typeof chunk.cursor === 'string' - ? chunk.cursor - : undefined - if (cursor && this.currentRunId) { + const runId = getChunkRunId(chunk) + const threadId = + 'threadId' in chunk && typeof chunk.threadId === 'string' + ? chunk.threadId + : this.activeResumeThreadId + + if (chunk.type === 'RUN_FINISHED' && chunk.outcome?.type === 'interrupt') { + // Track the REQUEST run id (what the client sent) so a resume targets the + // same run even when provider events carry their own run id. + const interruptedRunId = + this.currentRunId ?? runId ?? this.activeResumeRunId ?? '' this.lastResume = { - threadId: this.activeResumeThreadId ?? this.threadId, - runId: this.currentRunId, - cursor, + threadId: threadId ?? this.threadId, + runId: interruptedRunId, } + this.pendingInterruptRunId = interruptedRunId + this.pendingInterrupts = [...chunk.outcome.interrupts] + this.pendingInterruptResumeItems.clear() this.notifyResumeStateChange() + return } + + const isRunlessSessionError = chunk.type === 'RUN_ERROR' && !runId + const isTrackedRunTerminal = Boolean( + runId && this.lastResume?.runId === runId, + ) + const isPendingInterruptRunTerminal = Boolean( + runId && this.pendingInterruptRunId === runId, + ) + const isCurrentRunTerminal = Boolean( + (runId && this.currentRunId === runId) || + (this.currentRunId && this.lastResume?.runId === this.currentRunId), + ) + const isCurrentStreamTerminal = + this.isLoading && chunk.type === 'RUN_FINISHED' && !runId + if ( + isRunlessSessionError || + isTrackedRunTerminal || + isPendingInterruptRunTerminal || + isCurrentRunTerminal || + isCurrentStreamTerminal + ) { + this.lastResume = null + this.pendingInterrupts = [] + this.pendingInterruptRunId = null + this.pendingInterruptResumeItems.clear() + } + this.notifyResumeStateChange() } /** - * The resume state for the active/interrupted run (the run id plus the last - * cursor seen), or null when there is nothing to resume. Apps can persist this - * to resume across a full reload; in-session reconnects use it automatically - * via {@link maybeAutoResume}. + * The interrupt-resume state for the active/interrupted run (its run/thread + * ids), or null when there is nothing to resume. Apps can persist this to + * resume interrupts across a full reload. */ getResumeState(): ChatResumeState | null { return this.lastResume ? { ...this.lastResume } : null @@ -741,21 +728,6 @@ export class ChatClient< return [...this.pendingInterrupts] } - /** - * Resume a run by replaying its persisted events after the last cursor, then - * continuing live — without re-sending messages. Uses the supplied state, or - * the tracked in-session state. No-op (returns false) when there is nothing to - * resume or a stream is already in flight. - */ - resume(state?: ChatResumeState): Promise { - const target = state ?? this.lastResume - if (!target || this.isLoading) return Promise.resolve(false) - this.pendingResumeThreadId = target.threadId - this.pendingResumeRunId = target.runId - this.pendingResumeCursor = target.cursor - return this.streamResponse() - } - resumeInterrupts( resume: Array, state?: ChatResumeState, @@ -764,28 +736,10 @@ export class ChatClient< if (!target || this.isLoading) return Promise.resolve(false) this.pendingResumeThreadId = target.threadId this.pendingResumeRunId = target.runId - this.pendingResumeCursor = target.cursor this.pendingResumeItems = [...resume] return this.streamResponse() } - /** - * Auto-resume hook for framework integrations to call on mount / when the tab - * comes back online. Honors the `autoResume` option (default true) and only - * fires when an interrupted run is tracked and no stream is in flight. - */ - maybeAutoResume(): Promise { - if ( - !this.autoResume || - this.isLoading || - !this.lastResume || - this.pendingInterrupts.length > 0 - ) { - return Promise.resolve(false) - } - return this.resume() - } - private generateUniqueId(prefix: string): string { return `${prefix}-${Date.now()}-${Math.random().toString(36).substring(7)}` } @@ -849,8 +803,9 @@ export class ChatClient< : this.serverPersistence.removeItem(this.threadId) if (result instanceof Promise) { result - .catch(() => { - // Persistence adapters are best-effort and must not break chat. + .catch((error: unknown) => { + // Best-effort: must not break chat, but is surfaced to a sink. + this.reportServerPersistenceError(error) }) .finally(() => { if (generation !== this.serverPersistenceGeneration) { @@ -858,8 +813,9 @@ export class ChatClient< } }) } - } catch { - // Persistence adapters are best-effort and must not break chat. + } catch (error) { + // Best-effort: must not break chat, but is surfaced to a sink. + this.reportServerPersistenceError(error) } } @@ -879,8 +835,9 @@ export class ChatClient< : this.serverPersistence.removeItem(this.threadId) if (result instanceof Promise) { result - .catch(() => { - // Persistence adapters are best-effort and must not break chat. + .catch((error: unknown) => { + // Best-effort: must not break chat, but is surfaced to a sink. + this.reportServerPersistenceError(error) }) .finally(() => { if (generation !== this.serverPersistenceGeneration) { @@ -888,8 +845,9 @@ export class ChatClient< } }) } - } catch { - // Persistence adapters are best-effort and must not break chat. + } catch (error) { + // Best-effort: must not break chat, but is surfaced to a sink. + this.reportServerPersistenceError(error) } } @@ -1066,7 +1024,7 @@ export class ChatClient< // per-run error only clears that run, while a runId-less RUN_ERROR is // treated as a session-level error that clears every active run. this.updateRunLifecycle(chunk) - this.observeResumeCursor(chunk) + this.observeInterruptState(chunk) // Yield control back to event loop for UI updates await new Promise((resolve) => setTimeout(resolve, 0)) } @@ -1235,17 +1193,16 @@ export class ChatClient< // Track generation so a superseded stream's cleanup doesn't clobber the new one const generation = ++this.streamGeneration - // Resuming reuses the original runId so the server replays that run's events. + // Resuming an interrupt reuses the original runId so the server continues + // that run with the resume decisions. const resumeThreadId = this.pendingResumeThreadId const resumeRunId = this.pendingResumeRunId - const resumeCursor = this.pendingResumeCursor const resumeItems = this.pendingResumeItems const isResumeRequest = Boolean( - resumeThreadId || resumeRunId || resumeCursor || resumeItems, + resumeThreadId || resumeRunId || resumeItems, ) this.pendingResumeThreadId = null this.pendingResumeRunId = null - this.pendingResumeCursor = null this.pendingResumeItems = null const runId = resumeRunId ?? @@ -1342,7 +1299,6 @@ export class ChatClient< : { type: 'object' }, })), forwardedProps: { ...mergedBody }, - ...(resumeCursor ? { cursor: resumeCursor } : {}), ...(resumeItems ? { resume: resumeItems } : {}), } this.devtoolsBridge.beginRun(runContext.runId, runContext.threadId) @@ -1600,7 +1556,6 @@ export class ChatClient< this.pendingInterruptResumeItems.clear() this.pendingResumeThreadId = null this.pendingResumeRunId = null - this.pendingResumeCursor = null this.pendingResumeItems = null this.notifyResumeStateChange() this.setError(undefined) @@ -1988,6 +1943,7 @@ export class ChatClient< } if (options.onError !== undefined) { this.callbacksRef.current.onError = options.onError + this.hasUserOnError = true } if (options.onSubscriptionChange !== undefined) { this.callbacksRef.current.onSubscriptionChange = diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index e427e7a1b..d9154d0c5 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -48,6 +48,20 @@ export class StreamTruncatedError extends Error { } } +/** + * Thrown when a durable (id-tagged) run's stream ends with no terminal event + * and a reconnect makes no forward progress — the run cannot complete, so the + * consumer must not be left silently hanging on a stream that just stops. + */ +export class DurableStreamIncompleteError extends Error { + constructor() { + super( + 'Durable run ended without a terminal event and could not resume — the run did not complete.', + ) + this.name = 'DurableStreamIncompleteError' + } +} + function generateRunId(prefix: string): string { return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 8)}` } @@ -142,10 +156,16 @@ async function* readStreamLines( * * A JSON parse failure throws — the consumer surfaces it as an error. */ -async function* responseToSSEChunks( +/** + * Yield StreamChunks parsed from an SSE Response body, each paired with the + * `id:` offset of the SSE event it arrived on (if any). Delivery durability + * tags every event with `id: `; carrying that id up lets the + * resumable adapter track the last offset (for reconnect) and de-dupe replays. + */ +async function* responseToSSEEvents( response: Response, abortSignal?: AbortSignal, -): AsyncGenerator { +): AsyncGenerator<{ chunk: StreamChunk; id?: string }> { if (!response.ok) { throw new Error( `HTTP error! status: ${response.status} ${response.statusText}`, @@ -155,11 +175,15 @@ async function* responseToSSEChunks( let lastThreadId: string | undefined let lastRunId: string | undefined let lastModel: string | undefined + let pendingId: string | undefined for await (const line of readStreamLines(reader, abortSignal)) { + if (line.startsWith('id:')) { + pendingId = line.slice(3).trim() + continue + } if ( line.startsWith(':') || line.startsWith('event:') || - line.startsWith('id:') || line.startsWith('retry:') ) { continue @@ -174,7 +198,7 @@ async function* responseToSSEChunks( timestamp: Date.now(), finishReason: 'stop', } - yield synthetic + yield { chunk: synthetic } return } const chunk = JSON.parse(data) as StreamChunk @@ -187,10 +211,110 @@ async function* responseToSSEChunks( if ('model' in chunk && typeof chunk.model === 'string') { lastModel = chunk.model } + const id = pendingId + pendingId = undefined + yield { chunk, ...(id !== undefined ? { id } : {}) } + } +} + +async function* responseToSSEChunks( + response: Response, + abortSignal?: AbortSignal, +): AsyncGenerator { + for await (const { chunk } of responseToSSEEvents(response, abortSignal)) { yield chunk } } +/** + * Iterate an SSE endpoint with native-style resumability. Each SSE event's + * `id:` (a `runId@seq` delivery-durability offset) is remembered; if the + * connection drops or ends before a terminal event, the request is re-issued + * with a `Last-Event-ID` header so the server replays strictly after the last + * offset. Already-seen offsets are de-duped, so an overlapping replay is safe. + * + * When the server does NOT tag events (no durability), no offset is ever seen, + * so no reconnect happens — behaviour is identical to a plain single fetch. + */ +async function* resumableServerSentEvents( + fetchClient: typeof globalThis.fetch, + url: string, + requestInit: RequestInit, + abortSignal?: AbortSignal, +): AsyncGenerator { + const seen = new Set() + let lastEventId: string | undefined + + for (;;) { + if (abortSignal?.aborted) return + const headers: Record = { + ...(requestInit.headers as Record | undefined), + } + if (lastEventId !== undefined) { + headers['Last-Event-ID'] = lastEventId + } + const response = await fetchClient(url, { + ...requestInit, + headers, + ...(abortSignal ? { signal: abortSignal } : {}), + }) + + let sawTerminal = false + let progressed = false + try { + for await (const { chunk, id } of responseToSSEEvents( + response, + abortSignal, + )) { + if (id !== undefined) { + if (seen.has(id)) continue + seen.add(id) + lastEventId = id + } + progressed = true + if (chunk.type === 'RUN_FINISHED' || chunk.type === 'RUN_ERROR') { + sawTerminal = true + } + yield chunk + if (sawTerminal) return + } + } catch (error) { + if (abortSignal?.aborted) return + // A truncated connection is resumable only if we have an offset and made + // forward progress; otherwise surface the failure. + if ( + error instanceof StreamTruncatedError && + lastEventId !== undefined && + progressed + ) { + continue + } + throw error + } + + if (abortSignal?.aborted) return + + if (lastEventId !== undefined) { + // A durable (id-tagged) run. + if (progressed) { + // Clean end WITHOUT a terminal event but we advanced — the producer is + // still going (or the socket rolled over). Reconnect from the last + // offset. + continue + } + // Ended without a terminal event AND made no forward progress on this + // pass: the run cannot complete. Surface an error rather than returning + // silently, which would leave the consumer with neither a terminal event + // nor a failure. + throw new DurableStreamIncompleteError() + } + + // A non-durable (untagged) stream that ended cleanly. Legitimate — the + // upper layer synthesizes a terminal event. Stop. + return + } +} + /** * Per-send context provided by the chat client to the connection adapter. * The adapter combines this with serialized messages to build a full @@ -200,12 +324,6 @@ export interface RunAgentInputContext { threadId: string runId: string parentRunId?: string - /** - * Resume cursor. When set, the request resumes `runId` — the server replays - * persisted events after this cursor (see `chat({ cursor })`). On a resume the - * client sends no new messages. - */ - cursor?: string /** AG-UI interrupt resume entries returned to the server on a follow-up run. */ resume?: Array /** Client-declared tools to advertise in the request payload. */ @@ -230,6 +348,23 @@ export interface ConnectConnectionAdapter { ) => AsyncIterable } +/** + * A {@link ConnectConnectionAdapter} that also supports joining an existing run + * (a second tab, or re-attaching after a full reload) via `joinRun`, replaying + * the ordered stream from the start off the server's delivery-durability sink. + */ +export interface ResumableConnectConnectionAdapter + extends ConnectConnectionAdapter { + /** + * Join an in-flight or finished run by id, replaying from the start + * (`?offset=-1`). Read-only — sends no messages. + */ + joinRun: ( + runId: string, + abortSignal?: AbortSignal, + ) => AsyncIterable +} + export interface SubscribeConnectionAdapter { /** * Subscribe to stream chunks. @@ -452,7 +587,6 @@ function buildRunAgentInputBody( ...(runContext?.parentRunId !== undefined && { parentRunId: runContext.parentRunId, }), - ...(runContext?.cursor !== undefined && { cursor: runContext.cursor }), ...(runContext?.resume !== undefined && { resume: runContext.resume }), state: {}, messages: wireMessages, @@ -503,7 +637,7 @@ export function fetchServerSentEvents( options: | FetchConnectionOptions | (() => FetchConnectionOptions | Promise) = {}, -): ConnectConnectionAdapter { +): ResumableConnectConnectionAdapter { return { async *connect(messages, data, abortSignal, runContext) { // Resolve URL and options if they are functions @@ -535,15 +669,52 @@ export function fetchServerSentEvents( // under `exactOptionalPropertyTypes`), so spread it conditionally // rather than passing `undefined` explicitly. const signal = abortSignal || resolvedOptions.signal - const response = await fetchClient(resolvedUrl, { - method: 'POST', - headers: requestHeaders, - body: JSON.stringify(requestBody), - credentials: resolvedOptions.credentials || 'same-origin', - ...(signal ? { signal } : {}), - }) - yield* responseToSSEChunks(response, abortSignal) + // Resumable SSE: if the server tags events with `id:` offsets (delivery + // durability), a dropped/rolled-over connection auto-reconnects with a + // `Last-Event-ID` header and de-dupes the replayed prefix. With no tags, + // this is a single plain fetch. + yield* resumableServerSentEvents( + fetchClient, + resolvedUrl, + { + method: 'POST', + headers: requestHeaders, + body: JSON.stringify(requestBody), + credentials: resolvedOptions.credentials || 'same-origin', + }, + signal, + ) + }, + async *joinRun(runId, abortSignal) { + // Read an in-flight or finished run from the start. `?offset=-1` tells the + // server's delivery-durability sink to replay from the beginning; `runId` + // identifies which run. This is a read-only GET — no messages are sent. + const resolvedUrl = typeof url === 'function' ? url() : url + const resolvedOptions = + typeof options === 'function' ? await options() : options + + const separator = resolvedUrl.includes('?') ? '&' : '?' + const joinUrl = `${resolvedUrl}${separator}offset=-1&runId=${encodeURIComponent( + runId, + )}` + + const requestHeaders: Record = { + ...mergeHeaders(resolvedOptions.headers), + } + const fetchClient = resolvedOptions.fetchClient ?? fetch + const signal = abortSignal || resolvedOptions.signal + + yield* resumableServerSentEvents( + fetchClient, + joinUrl, + { + method: 'GET', + headers: requestHeaders, + credentials: resolvedOptions.credentials || 'same-origin', + }, + signal, + ) }, } } @@ -994,7 +1165,6 @@ export function fetcherToConnectionAdapter( data, threadId: runContext.threadId, runId: runContext.runId, - ...(runContext.cursor !== undefined && { cursor: runContext.cursor }), ...(runContext.resume !== undefined && { resume: runContext.resume }), }, { signal: abortSignal }, diff --git a/packages/ai-client/src/generation-client.ts b/packages/ai-client/src/generation-client.ts index acec24ea1..6eb0cf4c6 100644 --- a/packages/ai-client/src/generation-client.ts +++ b/packages/ai-client/src/generation-client.ts @@ -20,7 +20,6 @@ import type { GenerationClientState, GenerationFetcher, GenerationResumeSnapshot, - GenerationResumeState, GenerationServerPersistence, } from './generation-types' @@ -90,7 +89,6 @@ export class GenerationClient< private readonly devtoolsMetadata: AIDevtoolsClientMetadata private readonly devtoolsBridge: GenerationDevtoolsBridge private readonly threadId: string - private readonly autoResume: boolean private readonly serverPersistence: GenerationServerPersistence | undefined private body: Record private result: TOutput | null = null @@ -100,9 +98,7 @@ export class GenerationClient< private error: Error | undefined = undefined private status: GenerationClientState = 'idle' private resumeSnapshot: GenerationResumeSnapshot | undefined - private resumeState: GenerationResumeState | undefined private resumeSnapshotPersistenceQueue: Promise = Promise.resolve() - private resumeLifecycleToken = 0 private resumePersistenceError: Error | undefined = undefined private abortController: AbortController | null = null private readonly callbacksRef: GenerationCallbacks @@ -124,13 +120,8 @@ export class GenerationClient< this.connection = options.connection this.fetcher = options.fetcher this.body = options.body ?? {} - this.autoResume = options.autoResume ?? true this.serverPersistence = options.persistence?.server this.resumeSnapshot = options.initialResumeSnapshot - this.resumeState = - options.resumeState ?? - options.initialResumeSnapshot?.resumeState ?? - undefined this.callbacksRef = { onResult: options.onResult, @@ -201,10 +192,7 @@ export class GenerationClient< try { if (this.fetcher) { // Direct fetch path - const result = await this.fetcher(input, { - signal, - ...(this.resumeState ? { resumeState: this.resumeState } : {}), - }) + const result = await this.fetcher(input, { signal }) if (signal.aborted) return if (result instanceof Response) { // Server function returned SSE Response — parse stream @@ -266,41 +254,6 @@ export class GenerationClient< } } - async resume(state?: GenerationResumeState): Promise { - const resumeToken = this.resumeLifecycleToken - if (state) { - this.resumeState = state - this.resumeSnapshot = { - ...(this.resumeSnapshot ?? { status: 'running' }), - resumeState: state, - } - this.callbacksRef.onResumeSnapshotChange?.(this.resumeSnapshot) - } else { - await this.hydrateResumeSnapshot() - } - - if (this.disposed || resumeToken !== this.resumeLifecycleToken) { - return false - } - - if (!this.resumeState) { - return false - } - - // Resume requests are identified by resumeState/runContext. They do not - // need user media input and must not require callers to retain large input - // payloads across a reload. - await this.generate({} as TInput) - return true - } - - async maybeAutoResume(): Promise { - if (!this.autoResume || this.isLoading) { - return false - } - return this.resume() - } - /** * Process a stream of AG-UI events from the streaming connection adapter. */ @@ -368,7 +321,6 @@ export class GenerationClient< * Abort any in-flight generation request. */ stop(): void { - this.resumeLifecycleToken++ const runId = this.devtoolsBridge.getActiveRunId() if (this.abortController) { this.abortController.abort() @@ -409,7 +361,6 @@ export class GenerationClient< | 'onError' | 'onProgress' | 'onChunk' - | 'resumeState' > >, ): void { @@ -428,9 +379,6 @@ export class GenerationClient< if (options.onChunk !== undefined) { this.callbacksRef.onChunk = options.onChunk } - if (options.resumeState !== undefined) { - this.resumeState = options.resumeState - } } dispose(): void { @@ -575,13 +523,6 @@ export class GenerationClient< } private createRunContext(runId: string): RunAgentInputContext { - if (this.resumeState) { - return { - threadId: this.resumeState.threadId, - runId: this.resumeState.runId, - cursor: this.resumeState.cursor, - } - } return { threadId: this.threadId, runId, @@ -593,18 +534,16 @@ export class GenerationClient< this.resumeSnapshot, chunk, ) - this.resumeState = this.resumeSnapshot.resumeState ?? undefined this.callbacksRef.onResumeSnapshotChange?.(this.resumeSnapshot) void this.persistResumeSnapshot(this.resumeSnapshot) } private completePlainFetcherResumeSnapshot(): void { - if (!this.resumeState && !this.resumeSnapshot) { + if (!this.resumeSnapshot) { return } - this.resumeState = undefined this.resumeSnapshot = { - ...(this.resumeSnapshot ?? {}), + ...this.resumeSnapshot, resumeState: null, status: 'complete', } @@ -612,19 +551,6 @@ export class GenerationClient< void this.persistResumeSnapshot(this.resumeSnapshot) } - private async hydrateResumeSnapshot(): Promise { - if (this.resumeSnapshot || !this.serverPersistence) { - return - } - const snapshot = await this.serverPersistence.getItem(this.threadId) - if (!snapshot) { - return - } - this.resumeSnapshot = snapshot - this.resumeState = snapshot.resumeState ?? undefined - this.callbacksRef.onResumeSnapshotChange?.(snapshot) - } - private async persistResumeSnapshot( snapshot: GenerationResumeSnapshot, ): Promise { diff --git a/packages/ai-client/src/generation-types.ts b/packages/ai-client/src/generation-types.ts index 6f792de02..53194e3a0 100644 --- a/packages/ai-client/src/generation-types.ts +++ b/packages/ai-client/src/generation-types.ts @@ -67,7 +67,6 @@ export type GenerationResumeStatus = 'idle' | 'running' | 'complete' | 'error' export interface GenerationResumeState { threadId: string runId: string - cursor: string } export type GenerationPendingArtifact = PersistedArtifactRef @@ -150,8 +149,6 @@ export const GENERATION_EVENTS = { export interface GenerationFetcherOptions { /** AbortSignal that is triggered when the user calls `stop()` */ signal: AbortSignal - /** Explicit persisted run/cursor metadata for direct server-function resume */ - resumeState?: GenerationResumeState } /** @@ -199,33 +196,20 @@ export interface GenerationClientOptions<_TInput, TResult, TOutput = TResult> { /** Metadata used to register this generation hook with TanStack AI Devtools */ devtools?: Partial - /** - * Whether framework integrations should attempt to resume a persisted run on - * mount. Low-level clients expose the flag for hooks; they do not auto-resume - * by themselves. Defaults to `true` in framework hooks. - */ - autoResume?: boolean - /** * Initial lightweight resume snapshot restored by framework hooks. Contains - * only cursor metadata, errors, and persisted artifact refs. + * only observed run metadata, errors, and persisted artifact refs. This is + * read-only state for display; it does not trigger any generation. */ initialResumeSnapshot?: GenerationResumeSnapshot /** - * Optional persistence adapters for lightweight generation resume state. + * Optional persistence adapters for lightweight generation state. * Generation hooks only support `server` persistence; generated media bytes * are never written into browser storage by this client. */ persistence?: GenerationPersistenceOptions - /** - * Explicit run/cursor state to send on the next generation request. This is - * only a resume input; explicit `stop()` still aborts the local connection and - * does not model a durable server-side cancel. - */ - resumeState?: GenerationResumeState - /** * Factory that constructs the devtools bridge. Default is a no-op * factory; the real implementation lives in `@tanstack/ai-client/devtools`. @@ -267,7 +251,6 @@ export function updateGenerationResumeSnapshot( ): GenerationResumeSnapshot { const threadId = stringField(chunk, 'threadId') const runId = stringField(chunk, 'runId') - const cursor = stringField(chunk, 'cursor') const previousArtifacts = previous?.pendingArtifacts ?? [] const next: GenerationResumeSnapshot = { resumeState: previous?.resumeState ?? null, @@ -281,8 +264,8 @@ export function updateGenerationResumeSnapshot( lastEvent: createGenerationEventSnapshot(chunk), } - if (threadId && runId && cursor) { - next.resumeState = { threadId, runId, cursor } + if (threadId && runId) { + next.resumeState = { threadId, runId } next.status = 'running' } else if (chunk.type === 'RUN_STARTED') { next.status = 'running' diff --git a/packages/ai-client/src/index.ts b/packages/ai-client/src/index.ts index 3dfb4f428..53fefb69e 100644 --- a/packages/ai-client/src/index.ts +++ b/packages/ai-client/src/index.ts @@ -116,9 +116,11 @@ export { stream, rpcStream, StreamTruncatedError, + DurableStreamIncompleteError, type ConnectConnectionAdapter, type ConnectionAdapter, type FetchConnectionOptions, + type ResumableConnectConnectionAdapter, type RunAgentInputContext, type SubscribeConnectionAdapter, type XhrConnectionOptions, diff --git a/packages/ai-client/src/types.ts b/packages/ai-client/src/types.ts index c73ce1ed1..52ee48fbb 100644 --- a/packages/ai-client/src/types.ts +++ b/packages/ai-client/src/types.ts @@ -24,7 +24,6 @@ export type { StructuredOutputPart } export interface ChatResumeState { threadId: string runId: string - cursor: string } export type ChatPendingInterrupt = Interrupt @@ -46,7 +45,6 @@ export interface ChatFetcherInput { data?: Record threadId: string runId: string - cursor?: string resume?: Array } @@ -429,19 +427,10 @@ export interface ChatClientBaseOptions< */ threadId?: string - /** - * Whether to auto-resume an interrupted run when {@link maybeAutoResume} is - * called (e.g. by a framework integration on mount / when the tab comes back - * online). Requires server-side persistence so the run's events can be - * replayed by `runId + cursor`. Defaults to `true`; set `false` to opt out. - */ - autoResume?: boolean - /** * Initial resumable run state, useful when rehydrating a persisted client - * after a full page reload. The server still owns replay/resume decisions; - * this only restores the client-side interrupt descriptors needed to send - * AG-UI resume entries. + * after a full page reload. This restores the client-side interrupt + * descriptors needed to send AG-UI resume entries. */ initialResumeSnapshot?: ChatResumeSnapshot diff --git a/packages/ai-client/src/video-generation-client.ts b/packages/ai-client/src/video-generation-client.ts index 04c6eae6f..30624519f 100644 --- a/packages/ai-client/src/video-generation-client.ts +++ b/packages/ai-client/src/video-generation-client.ts @@ -19,7 +19,6 @@ import type { GenerationClientState, GenerationFetcher, GenerationResumeSnapshot, - GenerationResumeState, GenerationServerPersistence, VideoGenerateInput, VideoGenerateResult, @@ -97,7 +96,6 @@ export class VideoGenerationClient { private readonly devtoolsMetadata: AIDevtoolsClientMetadata private readonly devtoolsBridge: VideoDevtoolsBridge private readonly threadId: string - private readonly autoResume: boolean private readonly serverPersistence: GenerationServerPersistence | undefined private body: Record @@ -110,9 +108,7 @@ export class VideoGenerationClient { private error: Error | undefined = undefined private status: GenerationClientState = 'idle' private resumeSnapshot: GenerationResumeSnapshot | undefined - private resumeState: GenerationResumeState | undefined private resumeSnapshotPersistenceQueue: Promise = Promise.resolve() - private resumeLifecycleToken = 0 private resumePersistenceError: Error | undefined = undefined private abortController: AbortController | null = null private readonly callbacksRef: VideoCallbacks @@ -134,13 +130,8 @@ export class VideoGenerationClient { this.connection = options.connection this.fetcher = options.fetcher this.body = options.body ?? {} - this.autoResume = options.autoResume ?? true this.serverPersistence = options.persistence?.server this.resumeSnapshot = options.initialResumeSnapshot - this.resumeState = - options.resumeState ?? - options.initialResumeSnapshot?.resumeState ?? - undefined this.callbacksRef = { onResult: options.onResult, @@ -259,40 +250,6 @@ export class VideoGenerationClient { } } - async resume(state?: GenerationResumeState): Promise { - const resumeToken = this.resumeLifecycleToken - if (state) { - this.resumeState = state - this.resumeSnapshot = { - ...(this.resumeSnapshot ?? { status: 'running' }), - resumeState: state, - } - this.callbacksRef.onResumeSnapshotChange?.(this.resumeSnapshot) - } else { - await this.hydrateResumeSnapshot() - } - - if (this.disposed || resumeToken !== this.resumeLifecycleToken) { - return false - } - - if (!this.resumeState) { - return false - } - - // Resume is driven by run/cursor metadata. Do not require callers to keep - // large prompt/media input payloads in browser state across a refresh. - await this.generate({} as VideoGenerateInput) - return true - } - - async maybeAutoResume(): Promise { - if (!this.autoResume || this.isLoading) { - return false - } - return this.resume() - } - /** * Direct fetcher mode: call fetcher and set result. */ @@ -304,10 +261,7 @@ export class VideoGenerationClient { if (!this.fetcher) return // Fetcher returns a completed result directly, or a Response with SSE body - const result = await this.fetcher(input, { - signal, - ...(this.resumeState ? { resumeState: this.resumeState } : {}), - }) + const result = await this.fetcher(input, { signal }) if (signal.aborted) return if (result instanceof Response) { @@ -400,7 +354,6 @@ export class VideoGenerationClient { * Abort any in-flight generation or polling. */ stop(): void { - this.resumeLifecycleToken++ const runId = this.devtoolsBridge.getActiveRunId() if (this.abortController) { this.abortController.abort() @@ -445,7 +398,6 @@ export class VideoGenerationClient { | 'onChunk' | 'onJobCreated' | 'onStatusUpdate' - | 'resumeState' > >, ): void { @@ -470,9 +422,6 @@ export class VideoGenerationClient { if (options.onStatusUpdate !== undefined) { this.callbacksRef.onStatusUpdate = options.onStatusUpdate } - if (options.resumeState !== undefined) { - this.resumeState = options.resumeState - } } dispose(): void { @@ -656,13 +605,6 @@ export class VideoGenerationClient { } private createRunContext(runId: string): RunAgentInputContext { - if (this.resumeState) { - return { - threadId: this.resumeState.threadId, - runId: this.resumeState.runId, - cursor: this.resumeState.cursor, - } - } return { threadId: this.threadId, runId, @@ -674,18 +616,16 @@ export class VideoGenerationClient { this.resumeSnapshot, chunk, ) - this.resumeState = this.resumeSnapshot.resumeState ?? undefined this.callbacksRef.onResumeSnapshotChange?.(this.resumeSnapshot) void this.persistResumeSnapshot(this.resumeSnapshot) } private completePlainFetcherResumeSnapshot(): void { - if (!this.resumeState && !this.resumeSnapshot) { + if (!this.resumeSnapshot) { return } - this.resumeState = undefined this.resumeSnapshot = { - ...(this.resumeSnapshot ?? {}), + ...this.resumeSnapshot, resumeState: null, status: 'complete', } @@ -693,19 +633,6 @@ export class VideoGenerationClient { void this.persistResumeSnapshot(this.resumeSnapshot) } - private async hydrateResumeSnapshot(): Promise { - if (this.resumeSnapshot || !this.serverPersistence) { - return - } - const snapshot = await this.serverPersistence.getItem(this.threadId) - if (!snapshot) { - return - } - this.resumeSnapshot = snapshot - this.resumeState = snapshot.resumeState ?? undefined - this.callbacksRef.onResumeSnapshotChange?.(snapshot) - } - private async persistResumeSnapshot( snapshot: GenerationResumeSnapshot, ): Promise { diff --git a/packages/ai-client/tests/chat-client-resume.test.ts b/packages/ai-client/tests/chat-client-resume.test.ts index 4d51ebcf4..f2e7e4839 100644 --- a/packages/ai-client/tests/chat-client-resume.test.ts +++ b/packages/ai-client/tests/chat-client-resume.test.ts @@ -1,10 +1,6 @@ import { describe, expect, it, vi } from 'vitest' import { EventType } from '@tanstack/ai/client' import { ChatClient } from '../src/chat-client' -import { - createApprovalToolCallChunks, - createToolCallChunks, -} from './test-utils' import type { ConnectConnectionAdapter, RunAgentInputContext, @@ -40,12 +36,11 @@ function recordingAdapter(scripts: Array