@assistant-ui/ai-sdk 0.0.8 → 0.0.10

package/dist/runtime/useAISDKRuntime.js

@@ -15,6 +15,7 @@ import { useExternalStoreRuntime, useRuntimeAdapters } from "@assistant-ui/core/
 import { getExternalStoreMessages, pickExternalStoreSharedOptions } from "@assistant-ui/core";
 import { c } from "@assistant-ui/tap/react-shim/compiler-runtime";
 import { useCallback, useEffect, useInsertionEffect, useMemo, useRef, useState } from "@assistant-ui/tap/react-shim";
+import { useReplaySafeEffect } from "@assistant-ui/store/internal";
 //#region src/runtime/useAISDKRuntime.ts
 const toUIMessage = (createMessage, fallbackRole) => ({
 	...createMessage,
@@ -137,11 +138,97 @@ const useGeneratedSuggestions = (suggestionAdapter, messages, isRunning) => {
 		t6 = $[11];
 		t7 = $[12];
 	}
-	useEffect(t6, t7);
+	useReplaySafeEffect(t6, t7);
 	return suggestions;
 };
 const NO_CANCELLED_MESSAGE_IDS = /* @__PURE__ */ new Set();
+const NO_SUPERSEDED_APPROVAL_PROJECTION = Object.freeze({
+	approvalIds: Object.freeze(/* @__PURE__ */ new Set()),
+	statusMessageIds: Object.freeze(/* @__PURE__ */ new Set())
+});
 const NO_TOOL_APPROVAL_RESPONSES = /* @__PURE__ */ new Map();
+const hostToolApprovalsByChat = /* @__PURE__ */ new WeakMap();
+const toApprovalResponses = (owned) => owned && owned.size > 0 ? new Map([...owned].map(([id, entry]) => [id, entry.response])) : NO_TOOL_APPROVAL_RESPONSES;
+/**
+* The answers live on the owner, but each mounted runtime renders them from
+* its own state, so a write has to be announced: the runtime that performed it
+* may already be unmounted (a rollback resolving after a remount), and another
+* runtime may be mounted over the same owner.
+*/
+const hostApprovalListenersByChat = /* @__PURE__ */ new WeakMap();
+const subscribeToHostApprovals = (owner, listener) => {
+	const listeners = hostApprovalListenersByChat.get(owner) ?? /* @__PURE__ */ new Set();
+	hostApprovalListenersByChat.set(owner, listeners);
+	listeners.add(listener);
+	return () => {
+		listeners.delete(listener);
+	};
+};
+const notifyHostApprovals = (owner) => {
+	for (const listener of [...hostApprovalListenersByChat.get(owner) ?? []]) listener();
+};
+const getSupersededApprovalProjection = (messages, hostApprovalIds, joinStrategy, isRunning) => {
+	const approvalIds = /* @__PURE__ */ new Set();
+	const statusMessageIds = /* @__PURE__ */ new Set();
+	const lastIndex = messages.length - 1;
+	let lastAssistant;
+	let hasSupersededApproval = false;
+	let previousWasVoice = false;
+	const flush = () => {
+		const hasOpenToolPart = lastAssistant === messages[lastIndex] && lastAssistant?.parts?.some((part) => {
+			if (!isToolUIPart(part) || part.state === "output-available" || part.state === "output-error" || part.state === "output-denied") return false;
+			const approval = part.approval;
+			const resolution = approval?.resolution ?? (approval?.descriptor)?.resolution;
+			return resolution !== "cancelled" && resolution !== "expired";
+		});
+		if (lastAssistant && hasSupersededApproval && !hasOpenToolPart && !(isRunning && lastAssistant === messages[lastIndex])) statusMessageIds.add(lastAssistant.id);
+		lastAssistant = void 0;
+		hasSupersededApproval = false;
+		previousWasVoice = false;
+	};
+	for (let index = 0; index < messages.length; index++) {
+		const message = messages[index];
+		if (message.role !== "assistant") flush();
+		else {
+			const isVoice = message.metadata?.modality === "voice";
+			if (isVoice || previousWasVoice || joinStrategy === "none") flush();
+			lastAssistant = message;
+			previousWasVoice = isVoice;
+		}
+		for (const part of message.parts ?? []) {
+			if (!isToolUIPart(part) || part.state !== "approval-requested") continue;
+			const approval = part.approval;
+			if (!approval) continue;
+			const approvalId = approval.id;
+			if (hostApprovalIds.has(approvalId)) continue;
+			const resolution = approval.resolution ?? approval.descriptor?.resolution;
+			if (resolution === "cancelled" || resolution === "expired") continue;
+			if (index !== lastIndex) {
+				approvalIds.add(approvalId);
+				if (message.role === "assistant") hasSupersededApproval = true;
+			}
+		}
+	}
+	flush();
+	if (approvalIds.size === 0 && statusMessageIds.size === 0) return NO_SUPERSEDED_APPROVAL_PROJECTION;
+	return {
+		approvalIds,
+		statusMessageIds
+	};
+};
+const findRawToolMessageIndex = (messages, messageId, toolCallId, joinStrategy) => {
+	const containsToolCall = (message) => message.parts?.some((part) => isToolUIPart(part) && part.toolCallId === toolCallId) === true;
+	const messageIndex = messages.findIndex((message) => message.id === messageId);
+	if (messageIndex === -1) return -1;
+	if (containsToolCall(messages[messageIndex])) return messageIndex;
+	if (joinStrategy === "none" || messages[messageIndex]?.role !== "assistant") return -1;
+	let start = messageIndex;
+	while (start > 0 && messages[start - 1]?.role === "assistant") start--;
+	let end = messageIndex;
+	while (end + 1 < messages.length && messages[end + 1]?.role === "assistant") end++;
+	for (let index = start; index <= end; index++) if (containsToolCall(messages[index])) return index;
+	return -1;
+};
 const toChatError = (error) => {
 	const code = error.code;
 	return {
@@ -155,11 +242,42 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 	const contextAdapters = useRuntimeAdapters();
 	const [toolStatuses, setToolStatuses] = useState({});
 	const [cancelledMessages, setCancelledMessages] = useState(null);
-	const [toolApprovalResponses, setToolApprovalResponses] = useState(NO_TOOL_APPROVAL_RESPONSES);
+	const approvalOwner = adapter.unstable_hostApprovalOwner;
+	const ownedApprovals = approvalOwner ? hostToolApprovalsByChat.get(approvalOwner) ?? (() => {
+		const created = /* @__PURE__ */ new Map();
+		hostToolApprovalsByChat.set(approvalOwner, created);
+		return created;
+	})() : void 0;
+	const [toolApprovalResponses, setToolApprovalResponses] = useState(() => toApprovalResponses(ownedApprovals));
 	const [toolArtifactEpoch, setToolArtifactEpoch] = useState(0);
 	const [toolInteractionEpoch, setToolInteractionEpoch] = useState(0);
-	const hostApprovalIdsRef = useRef(/* @__PURE__ */ new Set());
-	const toolApprovalResponsesRef = useRef(/* @__PURE__ */ new Map());
+	const toolApprovalResponsesRef = useRef(new Map(toApprovalResponses(ownedApprovals)));
+	const ownedApprovalIdsRef = useRef(new Set(ownedApprovals?.keys()));
+	const hostApprovalIdsRef = useRef(new Set(ownedApprovals?.keys()));
+	useEffect(() => {
+		if (!approvalOwner || !ownedApprovals) return void 0;
+		const sync = () => {
+			for (const id of ownedApprovalIdsRef.current) toolApprovalResponsesRef.current.delete(id);
+			for (const [id_0, entry] of ownedApprovals) toolApprovalResponsesRef.current.set(id_0, entry.response);
+			ownedApprovalIdsRef.current = new Set(ownedApprovals.keys());
+			hostApprovalIdsRef.current = new Set(toolApprovalResponsesRef.current.keys());
+			setToolApprovalResponses((prev) => {
+				const next = new Map(toolApprovalResponsesRef.current);
+				return prev.size === next.size && [...next].every(([id_1, response]) => prev.get(id_1) === response) ? prev : next;
+			});
+		};
+		const unsubscribe = subscribeToHostApprovals(approvalOwner, sync);
+		sync();
+		return unsubscribe;
+	}, [approvalOwner, ownedApprovals]);
+	const lastApprovalOwnerRef = useRef(approvalOwner);
+	if (lastApprovalOwnerRef.current !== approvalOwner) {
+		lastApprovalOwnerRef.current = approvalOwner;
+		hostApprovalIdsRef.current = new Set(ownedApprovals?.keys());
+		toolApprovalResponsesRef.current = new Map(toApprovalResponses(ownedApprovals));
+		ownedApprovalIdsRef.current = new Set(ownedApprovals?.keys());
+		setToolApprovalResponses(new Map(toolApprovalResponsesRef.current));
+	}
 	const toolArgsKeyOrderCacheRef = useRef(/* @__PURE__ */ new Map());
 	const toolLastInputCacheRef = useRef(/* @__PURE__ */ new Map());
 	const toolArgsTextCacheRef = useRef(/* @__PURE__ */ new WeakMap());
@@ -182,6 +300,19 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 	const optimisticMessageId = isRunning && lastMessage?.role === "assistant" ? lastMessage.id : void 0;
 	const cancelledMessageIds = cancelledMessages?.chatId === chatHelpers.id ? cancelledMessages.ids : NO_CANCELLED_MESSAGE_IDS;
 	const supportsRichToolApprovalResponses = customOnRespondToToolApproval != null;
+	const supersededApprovalProjectionRef = useRef(NO_SUPERSEDED_APPROVAL_PROJECTION);
+	const supersededApprovalProjection = useMemo(() => {
+		const projection = getSupersededApprovalProjection(chatHelpers.messages, hostApprovalIdsRef.current, joinStrategy, isRunning);
+		const previous = supersededApprovalProjectionRef.current;
+		if (projection.approvalIds.size === previous.approvalIds.size && projection.statusMessageIds.size === previous.statusMessageIds.size && [...projection.approvalIds].every((id_3) => previous.approvalIds.has(id_3)) && [...projection.statusMessageIds].every((id_2) => previous.statusMessageIds.has(id_2))) return previous;
+		supersededApprovalProjectionRef.current = projection;
+		return projection;
+	}, [
+		chatHelpers.messages,
+		joinStrategy,
+		isRunning,
+		toolApprovalResponses
+	]);
 	const toThreadMessages = useCallback((sourceMessages) => {
 		const metadata = {
 			supportsRichToolApprovalResponses,
@@ -192,9 +323,9 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 		return AISDKMessageConverter.toThreadMessages(sourceMessages, false, metadata);
 	}, [supportsRichToolApprovalResponses]);
 	const retractCancellation = useCallback((chatId, messageId) => {
-		setCancelledMessages((prev) => {
-			if (prev?.chatId !== chatId || !prev.ids.has(messageId)) return prev;
-			const ids = new Set(prev.ids);
+		setCancelledMessages((prev_0) => {
+			if (prev_0?.chatId !== chatId || !prev_0.ids.has(messageId)) return prev_0;
+			const ids = new Set(prev_0.ids);
 			ids.delete(messageId);
 			return {
 				chatId,
@@ -228,6 +359,8 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 			toolArtifacts: toolArtifactsRef.current,
 			toolInteractions: toolInteractionsRef.current,
 			supportsRichToolApprovalResponses,
+			cancelledToolApprovalIds: supersededApprovalProjection.approvalIds,
+			cancelledStatusMessageIds: supersededApprovalProjection.statusMessageIds,
 			...optimisticMessageId && { optimisticMessageId },
 			...chatHelpers.error && { error: toChatError(chatHelpers.error) },
 			...cancelledMessageIds.size > 0 && { cancelledMessageIds },
@@ -240,6 +373,7 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 			cancelledMessageIds,
 			toolApprovalResponses,
 			supportsRichToolApprovalResponses,
+			supersededApprovalProjection,
 			toolArtifactEpoch,
 			toolInteractionEpoch
 		])
@@ -256,6 +390,7 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 	const { isLoading, deleteMessage: deleteHistoryMessage, persistToolInteractions, persistToolApprovalResponses } = useExternalHistory(runtimeRef, adapters?.history ?? contextAdapters?.history, toThreadMessages, aiSDKV6FormatAdapter, (messages_0) => {
 		chatHelpers.setMessages(messages_0);
 	}, toolArtifactsRef.current, markToolArtifactsChanged, toolInteractionsRef.current, markToolInteractionsChanged, toolApprovalResponsesRef.current, () => {
+		for (const [id_4, entry_0] of ownedApprovals ?? []) toolApprovalResponsesRef.current.set(id_4, entry_0.response);
 		hostApprovalIdsRef.current = new Set(toolApprovalResponsesRef.current.keys());
 		setToolApprovalResponses(new Map(toolApprovalResponsesRef.current));
 	});
@@ -273,7 +408,7 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 		if (!cancelPendingToolCallsOnSend) return;
 		chatHelpers.setMessages((messages_1) => {
 			let hasChanges = false;
-			const next = messages_1.map((message) => {
+			const next_0 = messages_1.map((message) => {
 				if (message.role !== "assistant") return message;
 				let messageChanged = false;
 				const parts = message.parts?.map((part) => {
@@ -295,7 +430,7 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 				};
 			});
 			if (!hasChanges) return messages_1;
-			return next;
+			return next_0;
 		});
 	};
 	const respondViaAISDK = ({ approvalId, approved, reason }) => Promise.resolve(chatHelpers.addToolApprovalResponse({
@@ -304,38 +439,51 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 		...reason != null && { reason },
 		options: { metadata: lastRunConfigRef.current }
 	}));
-	const respondViaHost = async (onRespond, response) => {
-		const { approvalId: approvalId_0 } = response;
+	const respondViaHost = async (onRespond, response_0) => {
+		const { approvalId: approvalId_0 } = response_0;
 		const requested = chatHelpers.messages.flatMap((message_0) => message_0.parts.flatMap((part_0) => isToolUIPart(part_0) ? [{
 			messageId: message_0.id,
 			part: part_0
-		}] : [])).find(({ part: part_1 }) => part_1.state === "approval-requested" && part_1.approval.id === approvalId_0);
+		}] : [])).find(({ part: part_1 }) => part_1.state === "approval-requested" && part_1.approval?.id === approvalId_0);
 		if (!requested || hostApprovalIdsRef.current.has(approvalId_0)) throw new Error(`Tool approval ${approvalId_0} is not waiting for a response.`);
+		const startedWith = ownedApprovals;
+		const startedOwner = approvalOwner;
+		const startedResponses = toolApprovalResponsesRef.current;
+		let isApplied = false;
 		const applyResponse = (applied) => {
+			isApplied = applied;
+			if (applied) startedWith?.set(approvalId_0, { response: response_0 });
+			else startedWith?.delete(approvalId_0);
+			if (startedOwner) {
+				notifyHostApprovals(startedOwner);
+				return;
+			}
+			if (lastApprovalOwnerRef.current !== startedOwner) return;
 			if (applied) hostApprovalIdsRef.current.add(approvalId_0);
 			else hostApprovalIdsRef.current.delete(approvalId_0);
-			if (applied) toolApprovalResponsesRef.current.set(approvalId_0, response);
+			if (applied) toolApprovalResponsesRef.current.set(approvalId_0, response_0);
 			else toolApprovalResponsesRef.current.delete(approvalId_0);
 			setToolApprovalResponses(new Map(toolApprovalResponsesRef.current));
 		};
 		applyResponse(true);
 		try {
-			await onRespond(response, {
+			await onRespond(response_0, {
 				toolCallId: requested.part.toolCallId,
 				toolName: getToolName(requested.part),
 				respondViaAISDK: async () => {
 					try {
-						await respondViaAISDK(response);
+						await respondViaAISDK(response_0);
 					} finally {
 						applyResponse(false);
 					}
 				}
 			});
 		} catch (error_0) {
-			if (hostApprovalIdsRef.current.has(approvalId_0)) applyResponse(false);
+			if (isApplied) applyResponse(false);
 			throw error_0;
 		}
-		if (hostApprovalIdsRef.current.has(approvalId_0)) await persistToolApprovalResponses(requested.messageId);
+		const remainsApplied = startedWith?.has(approvalId_0) ?? startedResponses.has(approvalId_0);
+		if (isApplied && remainsApplied) await persistToolApprovalResponses(requested.messageId);
 	};
 	const hasSeededRepositoryRef = useRef(false);
 	const shouldFeedRepository = exportedMessageRepository != null && !hasSeededRepositoryRef.current && messages.length === 0;
@@ -377,8 +525,8 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 			const cancelledId = isRunning && message_2?.role === "assistant" ? message_2.id : void 0;
 			if (cancelledId) {
 				const liveIds = new Set(chatHelpers.messages.map((m) => m.id));
-				setCancelledMessages((prev_0) => {
-					const kept = prev_0?.chatId === chatHelpers.id ? [...prev_0.ids].filter((id) => liveIds.has(id)) : [];
+				setCancelledMessages((prev_1) => {
+					const kept = prev_1?.chatId === chatHelpers.id ? [...prev_1.ids].filter((id_5) => liveIds.has(id_5)) : [];
 					return {
 						chatId: chatHelpers.id,
 						ids: /* @__PURE__ */ new Set([...kept, cancelledId])
@@ -427,13 +575,17 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 				removedToolArtifact = toolArtifactsRef.current.delete(part_2.toolCallId) || removedToolArtifact;
 				removedToolInteractions = toolInteractionsRef.current.delete(part_2.toolCallId) || removedToolInteractions;
 				if (part_2.approval) {
+					ownedApprovals?.delete(part_2.approval.id);
 					removedToolApprovalResponse = toolApprovalResponsesRef.current.delete(part_2.approval.id) || removedToolApprovalResponse;
 					removedHostApprovalId = hostApprovalIdsRef.current.delete(part_2.approval.id) || removedHostApprovalId;
 				}
 			}
 			if (removedToolArtifact) markToolArtifactsChanged();
 			if (removedToolInteractions) markToolInteractionsChanged();
-			if (removedToolApprovalResponse || removedHostApprovalId) setToolApprovalResponses(new Map(toolApprovalResponsesRef.current));
+			if (removedToolApprovalResponse || removedHostApprovalId) {
+				if (approvalOwner) notifyHostApprovals(approvalOwner);
+				setToolApprovalResponses(new Map(toolApprovalResponsesRef.current));
+			}
 			const deleteIds = new Set(getExternalStoreMessages(threadMessages[messageIndex]).map((message_6) => message_6.id));
 			chatHelpers.setMessages((current_3) => current_3.filter((message_7) => !deleteIds.has(message_7.id)));
 		},
@@ -443,34 +595,60 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 			chatHelpers.setMessages(newMessages);
 			await chatHelpers.regenerate({ metadata: config.runConfig });
 		},
-		onAddToolResult: ({ toolCallId, toolName, result: result_0, isError, artifact, modelContent }) => {
+		onAddToolResult: ({ messageId: messageId_1, toolCallId, toolName, result: result_0, isError, artifact, modelContent }) => {
 			if (artifact !== void 0) {
 				toolArtifactsRef.current.set(toolCallId, artifact);
 				markToolArtifactsChanged();
 			}
+			const targetIndex = findRawToolMessageIndex(chatHelpers.messages, messageId_1, toolCallId, joinStrategy);
+			const errorText = typeof result_0 === "string" ? result_0 : JSON.stringify(result_0);
+			const output = !isError && modelContent !== void 0 ? wrapModelContentEnvelope(result_0, modelContent) : result_0;
+			if (targetIndex >= 0 && targetIndex !== chatHelpers.messages.length - 1) {
+				const target = chatHelpers.messages[targetIndex];
+				const targetPart = target.parts.find((part_3) => isToolUIPart(part_3) && part_3.toolCallId === toolCallId);
+				if (targetPart?.state === "output-error" || targetPart?.state === "output-denied" || targetPart?.state === "output-available" && !targetPart.preliminary) return Promise.resolve();
+				const targetMessageId = target.id;
+				chatHelpers.setMessages((current_4) => current_4.map((message_8) => {
+					if (message_8.id !== targetMessageId) return message_8;
+					return {
+						...message_8,
+						parts: message_8.parts.map((part_4) => {
+							if (!isToolUIPart(part_4) || part_4.toolCallId !== toolCallId) return part_4;
+							const { preliminary: _preliminary, ...finalPart } = part_4;
+							return {
+								...finalPart,
+								state: isError ? "output-error" : "output-available",
+								output: isError ? void 0 : output,
+								errorText: isError ? errorText : void 0
+							};
+						})
+					};
+				}));
+				return Promise.resolve();
+			}
 			const options = { metadata: lastRunConfigRef.current };
 			if (isError) return Promise.resolve(chatHelpers.addToolOutput({
 				state: "output-error",
 				tool: toolName ?? toolCallId,
 				toolCallId,
-				errorText: typeof result_0 === "string" ? result_0 : JSON.stringify(result_0),
+				errorText,
 				options
 			}));
-			else {
-				const output = modelContent !== void 0 ? wrapModelContentEnvelope(result_0, modelContent) : result_0;
-				return Promise.resolve(chatHelpers.addToolOutput({
-					tool: toolName,
-					toolCallId,
-					output,
-					options
-				}));
-			}
+			return Promise.resolve(chatHelpers.addToolOutput({
+				tool: toolName,
+				toolCallId,
+				output,
+				options
+			}));
+		},
+		onRespondToToolApproval: (response_1) => {
+			if (supersededApprovalProjection.approvalIds.has(response_1.approvalId)) return Promise.reject(/* @__PURE__ */ new Error(`Tool approval ${response_1.approvalId} is not waiting for a response.`));
+			return customOnRespondToToolApproval ? respondViaHost(customOnRespondToToolApproval, response_1) : respondViaAISDK(response_1);
 		},
-		onRespondToToolApproval: customOnRespondToToolApproval ? (response_0) => respondViaHost(customOnRespondToToolApproval, response_0) : respondViaAISDK,
-		unstable_onRecordToolInteraction: ({ messageId: messageId_1, toolCallId: toolCallId_0, interaction }) => {
+		unstable_onRecordToolInteraction: ({ messageId: messageId_2, toolCallId: toolCallId_0, interaction }) => {
 			toolInteractionsRef.current.set(toolCallId_0, appendToolInteraction(toolInteractionsRef.current.get(toolCallId_0), interaction));
 			markToolInteractionsChanged();
-			return persistToolInteractions(messageId_1);
+			return persistToolInteractions(messageId_2);
 		},
 		...pickExternalStoreSharedOptions(adapter),
 		...adapter.unstable_messageRepositoryInstance && { unstable_messageRepositoryInstance: adapter.unstable_messageRepositoryInstance },