@tanstack/ai-persistence 0.1.5 → 0.4.0

package/dist/esm/middleware.js

@@ -1,8 +1,8 @@
 import { validateChatPersistenceStores, validateGenerationPersistenceStores } from "./types.js";
-import { InterruptsCapability, PersistenceCapability, provideInterrupts, providePersistence } from "./capabilities.js";
+import { InterruptsCapability, PersistenceCapability, PersistenceCompletionCapability, provideInterrupts, providePersistence, providePersistenceCompletion } from "./capabilities.js";
 import { artifactBlobKey } from "./retrieve.js";
-import { defineChatMiddleware, getDetachableRun, wasCancelRequested } from "@tanstack/ai";
-import { providePendingTurn } from "@tanstack/ai/adapter-internals";
+import { InterruptResumeValidationError, defineChatMiddleware, getDetachableRun, readInterruptBinding, validateInterruptResumeBatch, wasCancelRequested } from "@tanstack/ai";
+import { createInterruptBinding, getGenericInterruptDefinitionRegistry, providePendingTurn, rehydrateInterruptRequest } from "@tanstack/ai/adapter-internals";
 import { base64ToUint8Array } from "@tanstack/ai-utils";
 //#region src/middleware.ts
 /**
@@ -23,30 +23,116 @@ var DEFAULT_ARTIFACT_FETCH_TIMEOUT_MS = 3e4;
 var DEFAULT_MAX_ARTIFACT_BYTES = 1073741824;
 var runState = /* @__PURE__ */ new WeakMap();
 var validResumeStatuses = /* @__PURE__ */ new Set(["resolved", "cancelled"]);
-function validatePendingResumes(pending, resume) {
+function mergeMaps(left, right) {
+	if (!left && !right) return void 0;
+	return new Map([...left ?? [], ...right ?? []]);
+}
+function mergeSets(left, right) {
+	if (!left && !right) return void 0;
+	return /* @__PURE__ */ new Set([...left ?? [], ...right ?? []]);
+}
+function mergeResumeToolState(left, right) {
+	if (!left) return right;
+	if (!right) return left;
+	return {
+		approvals: mergeMaps(left.approvals, right.approvals),
+		clientToolResults: mergeMaps(left.clientToolResults, right.clientToolResults),
+		genericInterrupts: mergeMaps(left.genericInterrupts, right.genericInterrupts),
+		genericInterruptRequests: mergeMaps(left.genericInterruptRequests, right.genericInterruptRequests),
+		deniedToolResults: mergeMaps(left.deniedToolResults, right.deniedToolResults),
+		cancelledToolCallIds: mergeSets(left.cancelledToolCallIds, right.cancelledToolCallIds)
+	};
+}
+function rejectMixedRunPending(pending, ctx) {
+	if (new Set(pending.map((interrupt) => interrupt.runId)).size <= 1) return;
+	throw new InterruptResumeValidationError([{
+		scope: "batch",
+		threadId: ctx.threadId,
+		interruptedRunId: ctx.runId,
+		generation: 0,
+		interruptIds: pending.map((interrupt) => interrupt.interruptId),
+		code: "stale",
+		message: "Thread has pending interrupts from more than one run.",
+		source: "server",
+		retryable: false
+	}]);
+}
+function validatePendingResumes(pending, resume, ctx) {
+	const interruptedRunId = pending[0]?.runId ?? ctx.runId;
+	const failure = (interruptId, code, message) => {
+		throw new InterruptResumeValidationError([{
+			scope: "item",
+			threadId: ctx.threadId,
+			interruptedRunId,
+			generation: 0,
+			interruptId,
+			code,
+			message,
+			source: "client",
+			retryable: false
+		}, {
+			scope: "batch",
+			threadId: ctx.threadId,
+			interruptedRunId,
+			generation: 0,
+			interruptIds: pending.map((interrupt) => interrupt.interruptId),
+			code: code === "conflict" ? "conflict" : "incomplete-batch",
+			message: "Resume entries must resolve or cancel the complete interrupt batch.",
+			source: "client",
+			retryable: false
+		}]);
+	};
 	const pendingInterruptIds = new Set(pending.map((interrupt) => interrupt.interruptId));
-	const resumeByInterruptId = new Map((resume ?? []).map((entry) => [entry.interruptId, entry]));
+	const resumeByInterruptId = /* @__PURE__ */ new Map();
+	for (const entry of resume ?? []) {
+		if (resumeByInterruptId.has(entry.interruptId)) return failure(entry.interruptId, "conflict", `Interrupt ${entry.interruptId} has duplicate resume entries.`);
+		resumeByInterruptId.set(entry.interruptId, entry);
+	}
 	if (pending.length === 0) {
 		const staleEntry = resume?.[0];
-		if (staleEntry) throw new Error(`Resume entry references non-pending interrupt ${staleEntry.interruptId}.`);
+		if (staleEntry) return failure(staleEntry.interruptId, "unknown-interrupt", `Resume entry references non-pending interrupt ${staleEntry.interruptId}.`);
 		return resumeByInterruptId;
 	}
-	if (!resume || resume.length === 0) throw new Error(`Thread has pending interrupts; resume is required before accepting new input.`);
+	const firstPending = pending[0];
+	if (firstPending === void 0) return resumeByInterruptId;
+	if (!resume || resume.length === 0) return failure(firstPending.interruptId, "unknown-interrupt", `Thread has pending interrupts; resume is required before accepting new input.`);
 	for (const interrupt of pending) {
 		const entry = resumeByInterruptId.get(interrupt.interruptId);
-		if (!entry) throw new Error(`Missing resume entry for pending interrupt ${interrupt.interruptId}.`);
-		if (!validResumeStatuses.has(entry.status)) throw new Error(`Invalid resume status for pending interrupt ${interrupt.interruptId}: ${entry.status}.`);
+		if (!entry) return failure(interrupt.interruptId, "unknown-interrupt", `Missing resume entry for pending interrupt ${interrupt.interruptId}.`);
+		if (!validResumeStatuses.has(entry.status)) return failure(interrupt.interruptId, "unknown-interrupt", `Invalid resume status for pending interrupt ${interrupt.interruptId}: ${entry.status}.`);
 	}
-	for (const entry of resume) if (!pendingInterruptIds.has(entry.interruptId)) throw new Error(`Resume entry references non-pending interrupt ${entry.interruptId}.`);
+	for (const entry of resume) if (!pendingInterruptIds.has(entry.interruptId)) return failure(entry.interruptId, "unknown-interrupt", `Resume entry references non-pending interrupt ${entry.interruptId}.`);
 	return resumeByInterruptId;
 }
 async function applyPendingResumes(pending, resumeByInterruptId, interrupts) {
+	const entries = [];
 	for (const interrupt of pending) {
 		const entry = resumeByInterruptId.get(interrupt.interruptId);
 		if (!entry) continue;
-		if (entry.status === "resolved") await interrupts.resolve(interrupt.interruptId, entry.payload);
-		else await interrupts.cancel(interrupt.interruptId);
+		if (entry.status === "resolved") entries.push({
+			interruptId: interrupt.interruptId,
+			status: "resolved",
+			response: entry.payload
+		});
+		else entries.push({
+			interruptId: interrupt.interruptId,
+			status: "cancelled"
+		});
+	}
+	if (interrupts.commitBatch) {
+		await interrupts.commitBatch(entries);
+		return;
+	}
+	const ids = /* @__PURE__ */ new Set();
+	for (const entry of entries) {
+		if (ids.has(entry.interruptId)) throw new Error(`Interrupt batch contains duplicate id: ${entry.interruptId}.`);
+		ids.add(entry.interruptId);
+		const existing = await interrupts.get(entry.interruptId);
+		if (!existing) throw new Error(`Interrupt batch references missing id: ${entry.interruptId}.`);
+		if (existing.status !== "pending") throw new Error(`Interrupt batch references non-pending id: ${entry.interruptId}.`);
 	}
+	for (const entry of entries) if (entry.status === "resolved") await interrupts.resolve(entry.interruptId, entry.response);
+	else await interrupts.cancel(entry.interruptId);
 }
 /**
 * Commit the resumes stashed in `onConfig`, marking each resumed interrupt
@@ -72,6 +158,147 @@ function interruptKind(interrupt) {
 	const metadata = objectValue(interrupt.payload.metadata);
 	return metadata ? stringField(metadata, "kind") : void 0;
 }
+function hasReservedInterruptBinding(payload) {
+	const metadata = objectValue(objectValue(payload)?.metadata);
+	return !!metadata && "tanstack:interruptBinding" in metadata;
+}
+function isPersistedInterruptDescriptor(value) {
+	const record = objectValue(value);
+	return !!record && typeof record.id === "string" && typeof record.reason === "string" && typeof record.message === "string";
+}
+/**
+* Does this pending record belong to the TanStack chat resume protocol?
+*
+* An external system can persist an AG-UI descriptor in the same durable
+* thread. A descriptor without a TanStack binding or legacy tool marker stays
+* pending for its owner, but it does not make this resume incomplete. Older
+* opaque records remain owned because their provenance cannot be known.
+*/
+function isChatOwnedPendingInterrupt(interrupt) {
+	const kind = interruptKind(interrupt);
+	return !isPersistedInterruptDescriptor(interrupt.payload) || stringField(interrupt.payload, "toolCallId") !== void 0 || kind === "approval" || kind === "client_tool" || hasReservedInterruptBinding(interrupt.payload);
+}
+function durableGenericFailure(ctx, persisted, message) {
+	return new InterruptResumeValidationError([{
+		scope: "item",
+		threadId: ctx.threadId,
+		interruptedRunId: persisted.runId || ctx.runId,
+		generation: 0,
+		interruptId: persisted.interruptId,
+		code: "stale",
+		message,
+		source: "server",
+		retryable: false
+	}, {
+		scope: "batch",
+		threadId: ctx.threadId,
+		interruptedRunId: persisted.runId || ctx.runId,
+		generation: 0,
+		interruptIds: [persisted.interruptId],
+		code: "item-validation-failed",
+		message: "One or more persisted interrupt records are invalid.",
+		source: "server",
+		retryable: false
+	}]);
+}
+async function durableGenericResumeState(ctx, pending, resume, tools) {
+	const registry = getGenericInterruptDefinitionRegistry(ctx, { optional: true });
+	const records = [];
+	for (const persisted of pending) {
+		if (!isPersistedInterruptDescriptor(persisted.payload)) {
+			if (hasReservedInterruptBinding(persisted.payload)) throw durableGenericFailure(ctx, persisted, `Persisted interrupt ${persisted.interruptId} has an invalid binding descriptor.`);
+			continue;
+		}
+		const descriptor = persisted.payload;
+		const binding = readInterruptBinding(descriptor);
+		if (!binding) {
+			if (hasReservedInterruptBinding(descriptor)) throw durableGenericFailure(ctx, persisted, `Persisted interrupt ${persisted.interruptId} has an invalid or incomplete binding.`);
+			continue;
+		}
+		if (descriptor.id !== persisted.interruptId || binding.interruptId !== persisted.interruptId || binding.interruptedRunId !== persisted.runId || binding.generation !== 0) throw durableGenericFailure(ctx, persisted, `Persisted interrupt ${persisted.interruptId} has stale correlation metadata.`);
+		if (binding.kind !== "generic") {
+			records.push({
+				interruptId: persisted.interruptId,
+				payload: descriptor,
+				binding
+			});
+			continue;
+		}
+		if (!binding.definitionId || !binding.key || binding.batchIndex === void 0) {
+			records.push({
+				interruptId: persisted.interruptId,
+				payload: descriptor,
+				binding
+			});
+			continue;
+		}
+		if (!registry) throw durableGenericFailure(ctx, persisted, `Persisted generic interrupt ${persisted.interruptId} cannot be restored because no interrupt registry is available.`);
+		const definition = registry.definitions.get(binding.definitionId);
+		if (!definition) throw durableGenericFailure(ctx, persisted, `Persisted generic interrupt definition ${binding.definitionId} is unavailable.`);
+		const payload = objectValue(descriptor.metadata)?.["tanstack:interruptPayload"];
+		let request;
+		try {
+			request = rehydrateInterruptRequest(definition, {
+				key: binding.key,
+				reason: descriptor.reason,
+				message: descriptor.message,
+				...descriptor.expiresAt !== void 0 ? { expiresAt: descriptor.expiresAt } : {},
+				...payload !== void 0 ? { payload } : {}
+			});
+		} catch (error) {
+			throw durableGenericFailure(ctx, persisted, `Persisted generic interrupt ${persisted.interruptId} is invalid: ${error instanceof Error ? error.message : String(error)}`);
+		}
+		const emitted = createInterruptBinding(request, { batchIndex: binding.batchIndex });
+		if (emitted.descriptor.responseSchemaHash !== binding.responseSchemaHash || emitted.descriptor.payloadSchemaHash !== binding.payloadSchemaHash || binding.interruptId !== persisted.interruptId) throw durableGenericFailure(ctx, persisted, `Persisted generic interrupt ${persisted.interruptId} is stale.`);
+		records.push({
+			interruptId: persisted.interruptId,
+			payload: descriptor,
+			binding,
+			genericRequest: request
+		});
+	}
+	const firstRecord = records[0];
+	if (firstRecord === void 0) return void 0;
+	const interruptedRunId = firstRecord.binding.interruptedRunId;
+	const generation = firstRecord.binding.generation;
+	const validated = await validateInterruptResumeBatch({
+		threadId: ctx.threadId,
+		interruptedRunId,
+		generation,
+		pending: records,
+		resume: resume.filter((entry) => records.some((record) => record.interruptId === entry.interruptId)),
+		tools
+	});
+	if (validated.errors.length > 0 || !validated.resumeToolState) throw new InterruptResumeValidationError(validated.errors);
+	const isGenericRecord = (record) => record.binding.kind === "generic" && record.genericRequest !== void 0;
+	const genericRecords = [];
+	const batchIndexes = /* @__PURE__ */ new Set();
+	for (const record of records) {
+		if (!isGenericRecord(record)) continue;
+		const batchIndex = record.binding.batchIndex;
+		if (batchIndex === void 0 || batchIndexes.has(batchIndex)) throw new InterruptResumeValidationError([{
+			scope: "batch",
+			threadId: ctx.threadId,
+			interruptedRunId,
+			generation,
+			interruptIds: records.map((item) => item.interruptId),
+			code: "stale",
+			message: "Persisted generic interrupts have duplicate or invalid batch indexes.",
+			source: "server",
+			retryable: false
+		}]);
+		batchIndexes.add(batchIndex);
+		genericRecords.push({
+			record,
+			batchIndex
+		});
+	}
+	genericRecords.sort((left, right) => left.batchIndex - right.batchIndex);
+	return {
+		...validated.resumeToolState,
+		genericInterruptRequests: new Map(genericRecords.flatMap(({ record }) => record.genericRequest ? [[record.interruptId, record.genericRequest]] : []))
+	};
+}
 function resolvedApprovalDecision(entry) {
 	if (entry.status === "cancelled") return false;
 	const payload = objectValue(entry.payload);
@@ -109,30 +336,6 @@ function resumeToolStateFromPending(pending, resumeByInterruptId) {
 		cancelledToolCallIds
 	};
 }
-/**
-* Build the transcript to persist when a run finishes successfully.
-*
-* The chat engine appends an assistant message to the middleware message list
-* only when that turn carries tool calls (to feed the agent loop); a run's
-* terminal *text* reply is never appended. So `ctx.messages` at `onFinish` is
-* missing the assistant's final answer. Reattach it from the finish info —
-* `info.content` is the last turn's accumulated text (reset each cycle) — so
-* the stored thread is the complete conversation a server-authoritative client
-* hydrates on load. A guard avoids duplicating a terminal assistant turn should
-* the engine ever start appending it itself.
-*/
-function finishedTranscript(messages, info, messageId, createdAt) {
-	const transcript = [...messages];
-	const last = transcript[transcript.length - 1];
-	const alreadyPresent = last?.role === "assistant" && last.toolCalls === void 0 && last.content === info.content;
-	if (info.content && !alreadyPresent) transcript.push({
-		role: "assistant",
-		content: info.content,
-		...messageId ? { id: messageId } : {},
-		...createdAt ? { createdAt } : {}
-	});
-	return transcript;
-}
 function interruptPayload(interrupt) {
 	return interrupt && typeof interrupt === "object" ? { ...interrupt } : { value: interrupt };
 }
@@ -664,6 +867,7 @@ function accumulateTokenUsage(current, next) {
 	const providerUsageDetails = next.providerUsageDetails ?? current.providerUsageDetails;
 	const durationSeconds = sumOptionalNumber(current.durationSeconds, next.durationSeconds);
 	const unitsBilled = sumOptionalNumber(current.unitsBilled, next.unitsBilled);
+	const billed = accumulateBilled(current.billed, next.billed);
 	const cost = sumOptionalNumber(current.cost, next.cost);
 	return {
 		...current,
@@ -675,11 +879,25 @@ function accumulateTokenUsage(current, next) {
 		...completionTokensDetails ? { completionTokensDetails } : {},
 		...durationSeconds !== void 0 ? { durationSeconds } : {},
 		...unitsBilled !== void 0 ? { unitsBilled } : {},
+		...billed !== void 0 ? { billed } : {},
 		...cost !== void 0 ? { cost } : {},
 		...costDetails ? { costDetails } : {},
 		...providerUsageDetails ? { providerUsageDetails } : {}
 	};
 }
+/**
+* Sum billed quantities when both reports use the same unit. Different units
+* cannot be added, so the later report wins.
+*/
+function accumulateBilled(current, next) {
+	if (!current) return next;
+	if (!next) return current;
+	if (current.unit !== next.unit) return next;
+	return {
+		quantity: current.quantity + next.quantity,
+		unit: current.unit
+	};
+}
 async function completeRun(runs, runId, usage) {
 	await runs?.update(runId, {
 		status: "completed",
@@ -745,16 +963,33 @@ function withPersistence(persistence, options = {}) {
 	const { wantsInterrupts, runs } = resolvePersistencePlan(persistence);
 	const messageStore = persistence.stores.messages;
 	if (!messageStore) throw new Error("Chat persistence requires stores.messages.");
-	const provides = [PersistenceCapability, ...wantsInterrupts ? [InterruptsCapability] : []];
+	const provides = [
+		PersistenceCapability,
+		PersistenceCompletionCapability,
+		...wantsInterrupts ? [InterruptsCapability] : []
+	];
 	return defineChatMiddleware({
 		name: "chat-persistence",
 		provides,
 		setup(ctx) {
 			providePersistence(ctx, persistence);
+			let resolveCompletion = () => void 0;
+			let rejectCompletion = () => void 0;
+			const completion = new Promise((resolve, reject) => {
+				resolveCompletion = resolve;
+				rejectCompletion = reject;
+			});
+			completion.catch(() => void 0);
 			runState.set(ctx, {
 				merged: false,
-				interrupted: false
+				interrupted: false,
+				completion: {
+					promise: completion,
+					resolve: resolveCompletion,
+					reject: rejectCompletion
+				}
 			});
+			providePersistenceCompletion(ctx, { waitForRunCompletion: () => completion });
 			if (wantsInterrupts && persistence.stores.interrupts) provideInterrupts(ctx, persistence.stores.interrupts);
 			providePendingTurn(ctx, { snapshot: async () => {
 				const stored = await messageStore.loadThread(ctx.threadId);
@@ -766,16 +1001,18 @@ function withPersistence(persistence, options = {}) {
 			if (ctx.phase !== "init") return;
 			const patch = {};
 			if (wantsInterrupts && persistence.stores.interrupts) {
-				const pending = await persistence.stores.interrupts.listPending(ctx.threadId);
-				const resumeByInterruptId = validatePendingResumes(pending, config.resume);
+				const ownedPending = (await persistence.stores.interrupts.listPending(ctx.threadId)).filter(isChatOwnedPendingInterrupt);
+				rejectMixedRunPending(ownedPending, ctx);
+				const resumeByInterruptId = validatePendingResumes(ownedPending, config.resume, ctx);
 				if ((config.resume?.length ?? 0) > 0) {
-					const resumeToolState = resumeToolStateFromPending(pending, resumeByInterruptId);
+					const resumeToolState = resumeToolStateFromPending(ownedPending, resumeByInterruptId);
+					const genericResumeState = await durableGenericResumeState(ctx, ownedPending, config.resume ?? [], config.tools);
 					patch.resume = [];
-					if (resumeToolState) patch.resumeToolState = resumeToolState;
+					if (resumeToolState || genericResumeState) patch.resumeToolState = mergeResumeToolState(resumeToolState, genericResumeState);
 				}
 				const state = runState.get(ctx);
-				if (state && pending.length > 0) state.pendingResumes = {
-					pending,
+				if (state && ownedPending.length > 0) state.pendingResumes = {
+					pending: ownedPending,
 					resumeByInterruptId
 				};
 			}
@@ -795,7 +1032,7 @@ function withPersistence(persistence, options = {}) {
 			} catch {}
 		},
 		async onChunk(ctx, chunk) {
-			if (ctx.phase === "modelStream") {
+			if (snapshotStreaming && ctx.phase === "modelStream") {
 				const s = runState.get(ctx);
 				if (s && chunk.type === "TEXT_MESSAGE_START") {
 					s.streamingMessageId = typeof chunk.messageId === "string" && chunk.messageId !== "" ? chunk.messageId : void 0;
@@ -851,19 +1088,35 @@ function withPersistence(persistence, options = {}) {
 		async onFinish(ctx, info) {
 			const state = runState.get(ctx);
 			if (state?.interrupted) return;
-			await messageStore.saveThread(ctx.threadId, finishedTranscript(ctx.messages, info, state?.streamingMessageId, state?.streamingMessageCreatedAt));
-			await completeRun(runs, ctx.runId, state?.usage ?? info.usage);
-			await commitPendingResumes(state, persistence.stores.interrupts);
+			try {
+				await messageStore.saveThread(ctx.threadId, [...ctx.messages]);
+				await commitPendingResumes(state, persistence.stores.interrupts);
+				await completeRun(runs, ctx.runId, state?.usage ?? info.usage);
+				state?.completion?.resolve();
+			} catch (error) {
+				try {
+					await failRun(runs, ctx.runId, error, state?.usage);
+				} finally {
+					state?.completion?.reject(error);
+				}
+				throw error;
+			}
 		},
 		async onError(ctx, info) {
-			await failRun(runs, ctx.runId, info.error, runState.get(ctx)?.usage);
+			try {
+				await failRun(runs, ctx.runId, info.error, runState.get(ctx)?.usage);
+			} finally {
+				runState.get(ctx)?.completion?.reject(info.error);
+			}
 		},
 		async onAbort(ctx, info) {
-			const cancelled = info.cancelRequested === true || runs !== void 0 && await wasCancelRequested(runs, ctx.runId);
 			const state = runState.get(ctx);
-			if (cancelled || !detachableRun(ctx) && state?.interrupted !== true) {
-				await abortRun(runs, ctx.runId, state?.usage);
-				return;
+			let terminal = false;
+			try {
+				terminal = info.cancelRequested === true || runs !== void 0 && await wasCancelRequested(runs, ctx.runId) || !detachableRun(ctx) && state?.interrupted !== true;
+				if (terminal) await abortRun(runs, ctx.runId, state?.usage);
+			} finally {
+				if (terminal) state?.completion?.reject(info.reason);
 			}
 		}
 	});