@assistant-ui/ai-sdk 0.0.10 → 0.0.12

package/dist/runtime/useAISDKRuntime.js

@@ -2,7 +2,6 @@
 import { aiSDKV6FormatAdapter } from "../adapters/aiSDKFormatAdapter.js";
 import { vercelAttachmentAdapter } from "../adapters/vercelAttachmentAdapter.js";
 import { aiSDKExtras } from "../aiSDKExtras.js";
-import { wrapModelContentEnvelope } from "../converters/modelContentEnvelope.js";
 import { AISDKMessageConverter } from "../converters/convertMessage.js";
 import { toCreateMessage } from "../converters/toCreateMessage.js";
 import { sliceMessagesUntil } from "../utils/sliceMessagesUntil.js";
@@ -13,8 +12,8 @@ import { MessageRepository, appendToolInteraction, consumeSuggestionResult } fro
 import { generateId, getToolName, isToolUIPart } from "ai";
 import { useExternalStoreRuntime, useRuntimeAdapters } from "@assistant-ui/core/react";
 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 { wrapModelContentEnvelope } from "assistant-stream/internal";
+import { useCallback, useEffect, useInsertionEffect, useMemo, useRef, useState, useSyncExternalStore } from "@assistant-ui/tap/react-shim";
 import { useReplaySafeEffect } from "@assistant-ui/store/internal";
 //#region src/runtime/useAISDKRuntime.ts
 const toUIMessage = (createMessage, fallbackRole) => ({
@@ -36,109 +35,83 @@ const toVoiceTranscriptUIMessage = (message) => ({
 });
 const EMPTY_SUGGESTIONS = [];
 const useGeneratedSuggestions = (suggestionAdapter, messages, isRunning) => {
-	const $ = c(13);
 	const [suggestions, setSuggestions] = useState(EMPTY_SUGGESTIONS);
 	const controllerRef = useRef(null);
 	const wasRunningRef = useRef(false);
 	const messagesRef = useRef(messages);
-	let t0;
-	let t1;
-	if ($[0] !== messages) {
-		t0 = () => {
-			messagesRef.current = messages;
-		};
-		t1 = [messages];
-		$[0] = messages;
-		$[1] = t0;
-		$[2] = t1;
-	} else {
-		t0 = $[1];
-		t1 = $[2];
-	}
-	useInsertionEffect(t0, t1);
+	useInsertionEffect(() => {
+		messagesRef.current = messages;
+	}, [messages]);
+	const adapterKey = suggestionAdapter?.key;
 	const adapterRef = useRef(suggestionAdapter);
-	let t2;
-	let t3;
-	if ($[3] !== suggestionAdapter) {
-		t2 = () => {
-			adapterRef.current = suggestionAdapter;
-		};
-		t3 = [suggestionAdapter];
-		$[3] = suggestionAdapter;
-		$[4] = t2;
-		$[5] = t3;
-	} else {
-		t2 = $[4];
-		t3 = $[5];
-	}
-	useInsertionEffect(t2, t3);
-	const hasAdapter = suggestionAdapter != null;
-	let t4;
-	if ($[6] !== isRunning) {
-		t4 = () => {
-			const clearSuggestions = () => {
-				controllerRef.current?.abort();
-				controllerRef.current = null;
-				setSuggestions(_temp);
-			};
-			const adapter = adapterRef.current;
-			if (!adapter) {
-				clearSuggestions();
-				wasRunningRef.current = isRunning;
-				return;
-			}
-			if (isRunning) {
-				if (!wasRunningRef.current) clearSuggestions();
-				wasRunningRef.current = true;
-				return;
-			}
-			if (!wasRunningRef.current) return;
-			wasRunningRef.current = false;
-			const currentMessages = messagesRef.current;
-			const last = currentMessages.at(-1);
-			if (last?.role !== "assistant") return;
-			if (last.status?.type === "requires-action") return;
-			const controller = new AbortController();
-			controllerRef.current = controller;
-			const { signal } = controller;
-			(async () => {
-				try {
-					const promiseOrGenerator = adapter.generate({
-						messages: currentMessages,
-						signal
-					});
-					await consumeSuggestionResult(promiseOrGenerator, {
-						signal,
-						onUpdate: setSuggestions
-					});
-				} catch {}
-			})();
+	const lastAdapterKeyRef = useRef(adapterKey);
+	const adapterReplacedRef = useRef(false);
+	useInsertionEffect(() => {
+		const previousAdapter = adapterRef.current;
+		adapterRef.current = suggestionAdapter;
+		const adapterReplaced = previousAdapter !== void 0 && suggestionAdapter !== void 0 && lastAdapterKeyRef.current !== adapterKey;
+		adapterReplacedRef.current ||= adapterReplaced;
+		if (suggestionAdapter === void 0 || adapterReplaced) controllerRef.current?.abort();
+		lastAdapterKeyRef.current = adapterKey;
+	}, [adapterKey, suggestionAdapter]);
+	const hasAdapter = suggestionAdapter !== void 0;
+	const adapterPresentRef = useRef(hasAdapter);
+	useEffect(() => {
+		const clearSuggestions = () => {
+			controllerRef.current?.abort();
+			controllerRef.current = null;
+			setSuggestions((prev) => prev.length === 0 ? prev : EMPTY_SUGGESTIONS);
 		};
-		$[6] = isRunning;
-		$[7] = t4;
-	} else t4 = $[7];
-	let t5;
-	if ($[8] !== hasAdapter || $[9] !== isRunning) {
-		t5 = [hasAdapter, isRunning];
-		$[8] = hasAdapter;
-		$[9] = isRunning;
-		$[10] = t5;
-	} else t5 = $[10];
-	useEffect(t4, t5);
-	let t6;
-	let t7;
-	if ($[11] === Symbol.for("react.memo_cache_sentinel")) {
-		t6 = () => () => {
+		const adapter = adapterRef.current;
+		const adapterRestored = hasAdapter && !adapterPresentRef.current;
+		adapterPresentRef.current = hasAdapter;
+		const adapterReplaced_0 = adapterReplacedRef.current;
+		adapterReplacedRef.current = false;
+		if (!adapter) {
+			clearSuggestions();
+			wasRunningRef.current = isRunning;
+			return;
+		}
+		if (adapterReplaced_0) clearSuggestions();
+		if (isRunning) {
+			if (!wasRunningRef.current) clearSuggestions();
+			wasRunningRef.current = true;
+			return;
+		}
+		if (!wasRunningRef.current && !adapterReplaced_0 && !adapterRestored) return;
+		wasRunningRef.current = false;
+		const currentMessages = messagesRef.current;
+		const last = currentMessages.at(-1);
+		if (last?.role !== "assistant") return;
+		if (last.status?.type === "requires-action") return;
+		const controller = new AbortController();
+		controllerRef.current = controller;
+		const { signal } = controller;
+		(async () => {
+			try {
+				const promiseOrGenerator = adapter.generate({
+					messages: currentMessages,
+					signal
+				});
+				await consumeSuggestionResult(promiseOrGenerator, {
+					signal,
+					onUpdate: (nextSuggestions) => {
+						if (controllerRef.current !== controller) return;
+						setSuggestions(nextSuggestions);
+					}
+				});
+			} catch {}
+		})();
+	}, [
+		adapterKey,
+		hasAdapter,
+		isRunning
+	]);
+	useReplaySafeEffect(() => {
+		return () => {
 			controllerRef.current?.abort();
 		};
-		t7 = [];
-		$[11] = t6;
-		$[12] = t7;
-	} else {
-		t6 = $[11];
-		t7 = $[12];
-	}
-	useReplaySafeEffect(t6, t7);
+	}, []);
 	return suggestions;
 };
 const NO_CANCELLED_MESSAGE_IDS = /* @__PURE__ */ new Set();
@@ -147,7 +120,19 @@ const NO_SUPERSEDED_APPROVAL_PROJECTION = Object.freeze({
 	statusMessageIds: Object.freeze(/* @__PURE__ */ new Set())
 });
 const NO_TOOL_APPROVAL_RESPONSES = /* @__PURE__ */ new Map();
-const hostToolApprovalsByChat = /* @__PURE__ */ new WeakMap();
+const chatRuntimeStates = /* @__PURE__ */ new WeakMap();
+const getChatRuntimeState = (owner) => {
+	if (owner === void 0) return void 0;
+	let state = chatRuntimeStates.get(owner);
+	if (state === void 0) {
+		state = {
+			approvals: /* @__PURE__ */ new Map(),
+			cancelledMessageIds: NO_CANCELLED_MESSAGE_IDS
+		};
+		chatRuntimeStates.set(owner, state);
+	}
+	return state;
+};
 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
@@ -155,17 +140,17 @@ const toApprovalResponses = (owned) => owned && owned.size > 0 ? new Map([...own
 * 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);
+const chatRuntimeListeners = /* @__PURE__ */ new WeakMap();
+const subscribeToChatRuntimeState = (owner, listener) => {
+	const listeners = chatRuntimeListeners.get(owner) ?? /* @__PURE__ */ new Set();
+	chatRuntimeListeners.set(owner, listeners);
 	listeners.add(listener);
 	return () => {
 		listeners.delete(listener);
 	};
 };
-const notifyHostApprovals = (owner) => {
-	for (const listener of [...hostApprovalListenersByChat.get(owner) ?? []]) listener();
+const notifyChatRuntimeState = (owner) => {
+	for (const listener of [...chatRuntimeListeners.get(owner) ?? []]) listener();
 };
 const getSupersededApprovalProjection = (messages, hostApprovalIds, joinStrategy, isRunning) => {
 	const approvalIds = /* @__PURE__ */ new Set();
@@ -176,7 +161,7 @@ const getSupersededApprovalProjection = (messages, hostApprovalIds, joinStrategy
 	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;
+			if (typeof part?.type !== "string" || !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";
@@ -196,7 +181,7 @@ const getSupersededApprovalProjection = (messages, hostApprovalIds, joinStrategy
 			previousWasVoice = isVoice;
 		}
 		for (const part of message.parts ?? []) {
-			if (!isToolUIPart(part) || part.state !== "approval-requested") continue;
+			if (typeof part?.type !== "string" || !isToolUIPart(part) || part.state !== "approval-requested") continue;
 			const approval = part.approval;
 			if (!approval) continue;
 			const approvalId = approval.id;
@@ -217,7 +202,7 @@ const getSupersededApprovalProjection = (messages, hostApprovalIds, joinStrategy
 	};
 };
 const findRawToolMessageIndex = (messages, messageId, toolCallId, joinStrategy) => {
-	const containsToolCall = (message) => message.parts?.some((part) => isToolUIPart(part) && part.toolCallId === toolCallId) === true;
+	const containsToolCall = (message) => message.parts?.some((part) => typeof part?.type === "string" && 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;
@@ -243,11 +228,9 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 	const [toolStatuses, setToolStatuses] = useState({});
 	const [cancelledMessages, setCancelledMessages] = useState(null);
 	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 chatRuntimeState = getChatRuntimeState(approvalOwner);
+	const ownedApprovals = chatRuntimeState?.approvals;
+	const ownedCancelledMessageIds = useSyncExternalStore(useCallback((listener) => approvalOwner ? subscribeToChatRuntimeState(approvalOwner, listener) : () => {}, [approvalOwner]), useCallback(() => chatRuntimeState?.cancelledMessageIds ?? NO_CANCELLED_MESSAGE_IDS, [chatRuntimeState]), () => NO_CANCELLED_MESSAGE_IDS);
 	const [toolApprovalResponses, setToolApprovalResponses] = useState(() => toApprovalResponses(ownedApprovals));
 	const [toolArtifactEpoch, setToolArtifactEpoch] = useState(0);
 	const [toolInteractionEpoch, setToolInteractionEpoch] = useState(0);
@@ -266,7 +249,7 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 				return prev.size === next.size && [...next].every(([id_1, response]) => prev.get(id_1) === response) ? prev : next;
 			});
 		};
-		const unsubscribe = subscribeToHostApprovals(approvalOwner, sync);
+		const unsubscribe = subscribeToChatRuntimeState(approvalOwner, sync);
 		sync();
 		return unsubscribe;
 	}, [approvalOwner, ownedApprovals]);
@@ -298,7 +281,7 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 	const messageTiming = useStreamingTiming$1(chatHelpers.messages, isRunning);
 	const lastMessage = chatHelpers.messages.at(-1);
 	const optimisticMessageId = isRunning && lastMessage?.role === "assistant" ? lastMessage.id : void 0;
-	const cancelledMessageIds = cancelledMessages?.chatId === chatHelpers.id ? cancelledMessages.ids : NO_CANCELLED_MESSAGE_IDS;
+	const cancelledMessageIds = chatRuntimeState ? ownedCancelledMessageIds : cancelledMessages?.chatId === chatHelpers.id ? cancelledMessages.ids : NO_CANCELLED_MESSAGE_IDS;
 	const supportsRichToolApprovalResponses = customOnRespondToToolApproval != null;
 	const supersededApprovalProjectionRef = useRef(NO_SUPERSEDED_APPROVAL_PROJECTION);
 	const supersededApprovalProjection = useMemo(() => {
@@ -323,16 +306,24 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 		return AISDKMessageConverter.toThreadMessages(sourceMessages, false, metadata);
 	}, [supportsRichToolApprovalResponses]);
 	const retractCancellation = useCallback((chatId, messageId) => {
+		if (chatRuntimeState !== void 0 && approvalOwner !== void 0) {
+			if (!chatRuntimeState.cancelledMessageIds.has(messageId)) return;
+			const ids = new Set(chatRuntimeState.cancelledMessageIds);
+			ids.delete(messageId);
+			chatRuntimeState.cancelledMessageIds = ids;
+			notifyChatRuntimeState(approvalOwner);
+			return;
+		}
 		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);
+			const ids_0 = new Set(prev_0.ids);
+			ids_0.delete(messageId);
 			return {
 				chatId,
-				ids
+				ids: ids_0
 			};
 		});
-	}, []);
+	}, [approvalOwner, chatRuntimeState]);
 	const resumedMessageId = providerIsRunning && lastMessage?.role === "assistant" ? lastMessage.id : void 0;
 	useEffect(() => {
 		const wasProviderRunning = wasProviderRunningRef.current;
@@ -412,7 +403,7 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 				if (message.role !== "assistant") return message;
 				let messageChanged = false;
 				const parts = message.parts?.map((part) => {
-					if (!isToolUIPart(part)) return part;
+					if (typeof part?.type !== "string" || !isToolUIPart(part)) return part;
 					if (part.state === "output-available" || part.state === "output-error" || part.state === "output-denied") return part;
 					messageChanged = true;
 					const { approval: _approval, ...rest } = part;
@@ -441,7 +432,7 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 	}));
 	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) ? [{
+		const requested = chatHelpers.messages.flatMap((message_0) => message_0.parts.flatMap((part_0) => typeof part_0?.type === "string" && 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);
@@ -455,7 +446,7 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 			if (applied) startedWith?.set(approvalId_0, { response: response_0 });
 			else startedWith?.delete(approvalId_0);
 			if (startedOwner) {
-				notifyHostApprovals(startedOwner);
+				notifyChatRuntimeState(startedOwner);
 				return;
 			}
 			if (lastApprovalOwnerRef.current !== startedOwner) return;
@@ -485,6 +476,11 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 		const remainsApplied = startedWith?.has(approvalId_0) ?? startedResponses.has(approvalId_0);
 		if (isApplied && remainsApplied) await persistToolApprovalResponses(requested.messageId);
 	};
+	const lastAISDKMessageId = (threadMessageId) => {
+		if (threadMessageId == null) return null;
+		const threadMessage = runtimeRef.current.thread.getState().messages.find((message_1) => message_1.id === threadMessageId);
+		return (threadMessage && getVercelAIMessages(threadMessage).at(-1)?.id) ?? threadMessageId;
+	};
 	const hasSeededRepositoryRef = useRef(false);
 	const shouldFeedRepository = exportedMessageRepository != null && !hasSeededRepositoryRef.current && messages.length === 0;
 	const runtime = useExternalStoreRuntime({
@@ -495,7 +491,7 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 		setToolStatuses,
 		setMessages: (messages_2) => chatHelpers.setMessages(messages_2.map(getVercelAIMessages).filter(Boolean).flat()),
 		onImport: (messages_3) => chatHelpers.setMessages(messages_3.map(getVercelAIMessages).filter(Boolean).flat()),
-		onVoiceTranscript: (message_1) => chatHelpers.setMessages((current) => [...current, toVoiceTranscriptUIMessage(message_1)]),
+		onVoiceTranscript: (message_2) => chatHelpers.setMessages((current) => [...current, toVoiceTranscriptUIMessage(message_2)]),
 		onExportExternalState: () => {
 			const exported = runtimeRef.current.thread.export();
 			const expandedMessages = [];
@@ -521,12 +517,15 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 			runtimeRef.current.thread.import(exportedRepo);
 		},
 		onCancel: async () => {
-			const message_2 = chatHelpers.messages.at(-1);
-			const cancelledId = isRunning && message_2?.role === "assistant" ? message_2.id : void 0;
+			const message_3 = chatHelpers.messages.at(-1);
+			const cancelledId = isRunning && message_3?.role === "assistant" ? message_3.id : void 0;
 			if (cancelledId) {
 				const liveIds = new Set(chatHelpers.messages.map((m) => m.id));
-				setCancelledMessages((prev_1) => {
-					const kept = prev_1?.chatId === chatHelpers.id ? [...prev_1.ids].filter((id_5) => liveIds.has(id_5)) : [];
+				if (chatRuntimeState !== void 0 && approvalOwner !== void 0) {
+					chatRuntimeState.cancelledMessageIds = /* @__PURE__ */ new Set([...[...chatRuntimeState.cancelledMessageIds].filter((id_5) => liveIds.has(id_5)), cancelledId]);
+					notifyChatRuntimeState(approvalOwner);
+				} else setCancelledMessages((prev_1) => {
+					const kept = prev_1?.chatId === chatHelpers.id ? [...prev_1.ids].filter((id_6) => liveIds.has(id_6)) : [];
 					return {
 						chatId: chatHelpers.id,
 						ids: /* @__PURE__ */ new Set([...kept, cancelledId])
@@ -542,29 +541,30 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 				}
 			}
 		},
-		onNew: async (message_3) => {
-			const createMessage = (customToCreateMessage ?? toCreateMessage)(message_3);
-			if (!(message_3.startRun ?? message_3.role === "user")) {
-				chatHelpers.setMessages((current_0) => [...current_0, toUIMessage(createMessage, message_3.role)]);
+		onNew: async (message_4) => {
+			const createMessage = (customToCreateMessage ?? toCreateMessage)(message_4);
+			if (!(message_4.startRun ?? message_4.role === "user")) {
+				chatHelpers.setMessages((current_0) => [...current_0, toUIMessage(createMessage, message_4.role)]);
 				return;
 			}
-			lastRunConfigRef.current = message_3.runConfig;
+			lastRunConfigRef.current = message_4.runConfig;
 			await completePendingToolCalls();
-			await chatHelpers.sendMessage(createMessage, { metadata: message_3.runConfig });
+			await chatHelpers.sendMessage(createMessage, { metadata: message_4.runConfig });
 		},
-		onEdit: async (message_4) => {
-			const createMessage_0 = (customToCreateMessage ?? toCreateMessage)(message_4);
-			if (!(message_4.startRun ?? message_4.role === "user")) {
-				chatHelpers.setMessages((current_1) => [...sliceMessagesUntil(current_1, message_4.parentId), toUIMessage(createMessage_0, message_4.role)]);
+		onEdit: async (message_5) => {
+			const createMessage_0 = (customToCreateMessage ?? toCreateMessage)(message_5);
+			const parentId_0 = lastAISDKMessageId(message_5.parentId);
+			if (!(message_5.startRun ?? message_5.role === "user")) {
+				chatHelpers.setMessages((current_1) => [...sliceMessagesUntil(current_1, parentId_0), toUIMessage(createMessage_0, message_5.role)]);
 				return;
 			}
-			lastRunConfigRef.current = message_4.runConfig;
-			chatHelpers.setMessages((current_2) => sliceMessagesUntil(current_2, message_4.parentId));
-			await chatHelpers.sendMessage(createMessage_0, { metadata: message_4.runConfig });
+			lastRunConfigRef.current = message_5.runConfig;
+			chatHelpers.setMessages((current_2) => sliceMessagesUntil(current_2, parentId_0));
+			await chatHelpers.sendMessage(createMessage_0, { metadata: message_5.runConfig });
 		},
 		onDelete: async (messageId_0) => {
 			const threadMessages = runtimeRef.current.thread.getState().messages;
-			const messageIndex = threadMessages.findIndex((message_5) => message_5.id === messageId_0);
+			const messageIndex = threadMessages.findIndex((message_6) => message_6.id === messageId_0);
 			if (messageIndex === -1) return;
 			await deleteHistoryMessage(messageId_0);
 			let removedToolArtifact = false;
@@ -583,15 +583,17 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 			if (removedToolArtifact) markToolArtifactsChanged();
 			if (removedToolInteractions) markToolInteractionsChanged();
 			if (removedToolApprovalResponse || removedHostApprovalId) {
-				if (approvalOwner) notifyHostApprovals(approvalOwner);
+				if (approvalOwner) notifyChatRuntimeState(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)));
+			const deleteIds = new Set(getExternalStoreMessages(threadMessages[messageIndex]).map((message_7) => message_7.id));
+			chatHelpers.setMessages((current_3) => current_3.filter((message_8) => !deleteIds.has(message_8.id)));
 		},
-		onReload: async (parentId_0, config) => {
+		onReload: async (parentId_1, config) => {
 			lastRunConfigRef.current = config.runConfig;
-			const newMessages = sliceMessagesUntil(chatHelpers.messages, parentId_0);
+			const newMessages = sliceMessagesUntil(chatHelpers.messages, lastAISDKMessageId(parentId_1));
+			const reloaded = chatHelpers.messages[newMessages.length];
+			if (reloaded?.role === "assistant") newMessages.push(reloaded);
 			chatHelpers.setMessages(newMessages);
 			await chatHelpers.regenerate({ metadata: config.runConfig });
 		},
@@ -605,15 +607,15 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 			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);
+				const targetPart = target.parts.find((part_3) => typeof part_3?.type === "string" && 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;
+				chatHelpers.setMessages((current_4) => current_4.map((message_9) => {
+					if (message_9.id !== targetMessageId) return message_9;
 					return {
-						...message_8,
-						parts: message_8.parts.map((part_4) => {
-							if (!isToolUIPart(part_4) || part_4.toolCallId !== toolCallId) return part_4;
+						...message_9,
+						parts: message_9.parts.map((part_4) => {
+							if (typeof part_4?.type !== "string" || !isToolUIPart(part_4) || part_4.toolCallId !== toolCallId) return part_4;
 							const { preliminary: _preliminary, ...finalPart } = part_4;
 							return {
 								...finalPart,
@@ -682,9 +684,6 @@ const useAISDKRuntime = (chatHelpers, adapter = {}) => {
 	}, [exportedMessageRepository, chatHelpers.messages.length]);
 	return runtime;
 };
-function _temp(prev) {
-	return prev.length === 0 ? prev : EMPTY_SUGGESTIONS;
-}
 //#endregion
 export { useAISDKRuntime };