@assistant-ui/react-ag-ui 0.0.44 → 0.0.45

package/dist/runtime/AgUiThreadRuntimeCore.js

@@ -2,12 +2,14 @@
 import { AG_UI_METADATA_NAMESPACE, RunAggregator, tryParseJSON } from "./adapter/run-aggregator.js";
 import { fromAgUiMessages, toAgUiMessages, toAgUiTools } from "./adapter/conversions.js";
 import { createAgUiSubscriber } from "./adapter/subscriber.js";
+import { MessageRepository } from "@assistant-ui/core/internal";
 import { fromThreadMessageLike, generateId } from "@assistant-ui/core";
 import jsonpatch from "fast-json-patch";
 //#region src/runtime/AgUiThreadRuntimeCore.ts
 const optimisticPrefix = "__optimistic__";
 const generateOptimisticId = () => `${optimisticPrefix}${generateId()}`;
 const isOptimisticId = (id) => id.startsWith(optimisticPrefix);
+const isResolvedToolCall = (part) => part.type === "tool-call" && "result" in part && part.result !== void 0;
 const symbolResumeShim = Symbol("agui-resume-shim");
 const FALLBACK_USER_STATUS = {
 	type: "complete",
@@ -17,11 +19,13 @@ var AgUiThreadRuntimeCore = class {
 	agent;
 	logger;
 	showThinking;
+	autoCancelPendingToolCalls;
 	onError;
 	onCancel;
 	notifyUpdate;
 	runtime;
-	messages = [];
+	repository = new MessageRepository();
+	exportedRepository;
 	isRunningFlag = false;
 	abortController = null;
 	stateSnapshot;
@@ -37,6 +41,7 @@ var AgUiThreadRuntimeCore = class {
 		this.agent = options.agent;
 		this.logger = options.logger;
 		this.showThinking = options.showThinking;
+		this.autoCancelPendingToolCalls = options.autoCancelPendingToolCalls;
 		this.onError = options.onError;
 		this.onCancel = options.onCancel;
 		this.history = options.history;
@@ -47,6 +52,7 @@ var AgUiThreadRuntimeCore = class {
 		this.agent = options.agent;
 		this.logger = options.logger;
 		this.showThinking = options.showThinking;
+		this.autoCancelPendingToolCalls = options.autoCancelPendingToolCalls;
 		this.onError = options.onError;
 		this.onCancel = options.onCancel;
 		this.history = options.history;
@@ -59,7 +65,74 @@ var AgUiThreadRuntimeCore = class {
 		this.runtime = void 0;
 	}
 	getMessages() {
-		return this.messages;
+		return this.repository.getMessages();
+	}
+	getMessageRepository() {
+		if (this.exportedRepository) return this.exportedRepository;
+		const exportedRepository = this.repository.export();
+		let parentId = null;
+		for (const message of this.repository.getMessages()) {
+			if (message.metadata.isOptimistic) exportedRepository.messages.push({
+				parentId,
+				message
+			});
+			parentId = message.id;
+		}
+		this.exportedRepository = {
+			...exportedRepository,
+			headId: this.repository.headId
+		};
+		return this.exportedRepository;
+	}
+	tryGetMessage(messageId) {
+		try {
+			return this.repository.getMessage(messageId);
+		} catch {
+			return;
+		}
+	}
+	tryGetMessages(messageId) {
+		try {
+			return this.repository.getMessages(messageId);
+		} catch {
+			return;
+		}
+	}
+	hasMessage(messageId) {
+		return this.tryGetMessage(messageId) !== void 0;
+	}
+	addOrUpdateMessage(parentId, message) {
+		this.repository.addOrUpdateMessage(parentId, message);
+		this.exportedRepository = void 0;
+	}
+	deleteMessage(messageId, replacementId) {
+		this.repository.deleteMessage(messageId, replacementId);
+		this.exportedRepository = void 0;
+	}
+	tryDeleteMessage(messageId) {
+		if (!this.hasMessage(messageId)) return false;
+		this.deleteMessage(messageId);
+		return true;
+	}
+	switchToBranch(messageId) {
+		this.repository.switchToBranch(messageId);
+		this.exportedRepository = void 0;
+	}
+	resetRepositoryHead(messageId) {
+		this.repository.resetHead(messageId);
+		this.exportedRepository = void 0;
+	}
+	clearRepository() {
+		this.repository.clear();
+		this.exportedRepository = void 0;
+	}
+	updateMessage(messageId, updater) {
+		const item = this.tryGetMessage(messageId);
+		if (!item) return false;
+		const message = updater(item.message);
+		if (message === item.message) return false;
+		this.addOrUpdateMessage(item.parentId, message);
+		return true;
 	}
 	getState() {
 		return this.stateSnapshot;
@@ -76,11 +149,10 @@ var AgUiThreadRuntimeCore = class {
 		this._isLoading = true;
 		this._loadPromise = promise.then(async (repo) => {
 			if (!repo) return;
-			const messages = repo.messages.map((item) => item.message);
-			this.applyExternalMessages(messages);
+			this.applyExternalMessageRepository(repo);
 			if (repo.state !== void 0) this.loadExternalState(repo.state);
 			if (repo.unstable_resume) {
-				const parentId = repo.headId ?? messages.at(-1)?.id ?? null;
+				const parentId = repo.headId ?? this.repository.headId;
 				const resumeStream = this.history?.resume?.bind(this.history);
 				await this.startRun(parentId, this.lastRunConfig, void 0, resumeStream);
 			}
@@ -96,18 +168,29 @@ var AgUiThreadRuntimeCore = class {
 	}
 	async append(message) {
 		const startRun = message.startRun ?? message.role === "user";
-		if (startRun) this.assertNoPendingInterrupts();
+		if (startRun) {
+			this.assertNoPendingInterrupts();
+			this.maybeAutoCancelPendingToolCalls();
+		}
 		const threadMessageId = this.appendEntry(message);
 		if (!startRun) return;
 		await this.startRun(threadMessageId, message.runConfig);
 	}
+	maybeAutoCancelPendingToolCalls() {
+		if (this.autoCancelPendingToolCalls === false) return;
+		const pending = this.getPendingToolCalls();
+		if (!pending) return;
+		this.cancelUnresolvedToolCalls(pending.messageId);
+		this.maybeCompleteAfterToolResults(pending.messageId);
+	}
 	appendEntry(message) {
-		if (message.sourceId) this.messages = this.messages.filter((entry) => entry.id !== message.sourceId);
-		this.resetHead(message.parentId);
+		if (message.sourceId) this.tryDeleteMessage(message.sourceId);
 		const threadMessage = this.toThreadMessage(message);
-		this.messages = [...this.messages, threadMessage];
+		const parentId = message.parentId === null ? null : message.parentId && this.hasMessage(message.parentId) ? message.parentId : this.repository.headId;
+		this.addOrUpdateMessage(parentId, threadMessage);
+		this.switchToBranch(threadMessage.id);
 		this.notifyUpdate();
-		this.recordHistoryEntry(message.parentId ?? null, threadMessage);
+		this.recordHistoryEntry(parentId, threadMessage);
 		return threadMessage.id;
 	}
 	async edit(message) {
@@ -115,8 +198,7 @@ var AgUiThreadRuntimeCore = class {
 	}
 	async reload(parentId, config = {}) {
 		this.assertNoPendingInterrupts();
-		this.resetHead(parentId);
-		this.notifyUpdate();
+		this.maybeAutoCancelPendingToolCalls();
 		await this.startRun(parentId, config.runConfig);
 	}
 	async cancel() {
@@ -144,9 +226,14 @@ var AgUiThreadRuntimeCore = class {
 		if (!this.getPendingInterrupts()) return;
 		throw new Error("[agui] cannot start a new run while interrupts are pending; resolve them with submitInterruptResponses()");
 	}
+	findRequiresActionAssistant(reason) {
+		const assistant = this.getMessages().findLast((message) => message.role === "assistant");
+		if (!assistant || assistant.status?.type !== "requires-action" || assistant.status.reason !== reason) return null;
+		return assistant;
+	}
 	getPendingInterrupts() {
-		const assistant = this.messages.findLast((m) => m.role === "assistant");
-		if (!assistant || assistant.status?.type !== "requires-action" || assistant.status.reason !== "interrupt") return null;
+		const assistant = this.findRequiresActionAssistant("interrupt");
+		if (!assistant) return null;
 		const stored = assistant.metadata.custom[AG_UI_METADATA_NAMESPACE]?.interrupts;
 		if (!stored?.length) return null;
 		return {
@@ -154,6 +241,21 @@ var AgUiThreadRuntimeCore = class {
 			interrupts: stored
 		};
 	}
+	getPendingToolCalls() {
+		const assistant = this.findRequiresActionAssistant("tool-calls");
+		if (!assistant) return null;
+		const toolCallIds = [];
+		for (const part of assistant.content) {
+			if (part.type !== "tool-call") continue;
+			if (isResolvedToolCall(part)) continue;
+			toolCallIds.push(part.toolCallId);
+		}
+		if (toolCallIds.length === 0) return null;
+		return {
+			messageId: assistant.id,
+			toolCallIds
+		};
+	}
 	async submitInterruptResponses(responses) {
 		const pending = this.getPendingInterrupts();
 		if (!pending) throw new Error("[agui] submitInterruptResponses: no pending interrupts on this thread");
@@ -185,6 +287,17 @@ var AgUiThreadRuntimeCore = class {
 	async steerAway(message, responses) {
 		const pending = this.getPendingInterrupts();
 		if (!pending) {
+			const pendingTools = this.getPendingToolCalls();
+			if (pendingTools) {
+				if (responses?.length) throw new Error("[agui] steerAway: responses are only valid for pending interrupts");
+				if (this.isRunningFlag) throw new Error("[agui] steerAway: a run is already in progress");
+				this.cancelUnresolvedToolCalls(pendingTools.messageId);
+				this.maybeCompleteAfterToolResults(pendingTools.messageId);
+				const normalized = this.toAppendMessage(message);
+				const threadMessageId = this.appendEntry(normalized);
+				await this.startRun(threadMessageId, normalized.runConfig);
+				return;
+			}
 			if (responses?.length) throw new Error("[agui] steerAway: no pending interrupts on this thread");
 			await this.append(this.toAppendMessage(message));
 			return;
@@ -215,7 +328,7 @@ var AgUiThreadRuntimeCore = class {
 	toAppendMessage(message) {
 		if (typeof message === "string") return {
 			createdAt: /* @__PURE__ */ new Date(),
-			parentId: this.messages.at(-1)?.id ?? null,
+			parentId: this.repository.headId,
 			sourceId: null,
 			runConfig: {},
 			role: "user",
@@ -228,7 +341,7 @@ var AgUiThreadRuntimeCore = class {
 		};
 		return {
 			createdAt: message.createdAt ?? /* @__PURE__ */ new Date(),
-			parentId: message.parentId ?? this.messages.at(-1)?.id ?? null,
+			parentId: message.parentId ?? this.repository.headId,
 			sourceId: message.sourceId ?? null,
 			role: message.role ?? "user",
 			content: message.content,
@@ -239,16 +352,14 @@ var AgUiThreadRuntimeCore = class {
 		};
 	}
 	clearPendingInterrupts(messageId) {
-		let touched = false;
-		this.messages = this.messages.map((message) => {
-			if (message.id !== messageId || message.role !== "assistant") return message;
+		if (this.updateMessage(messageId, (message) => {
+			if (message.role !== "assistant") return message;
 			const assistant = message;
 			if (assistant.status?.type !== "requires-action" || assistant.status.reason !== "interrupt") return assistant;
-			touched = true;
 			const { interrupts: _drop, ...restAgui } = assistant.metadata.custom["agui"] ?? {};
 			const newCustom = { ...assistant.metadata.custom };
-			if (Object.keys(restAgui).length > 0) newCustom[AG_UI_METADATA_NAMESPACE] = restAgui;
-			else delete newCustom[AG_UI_METADATA_NAMESPACE];
+			if (Object.keys(restAgui).length > 0) newCustom["agui"] = restAgui;
+			else delete newCustom["agui"];
 			return {
 				...assistant,
 				status: {
@@ -260,26 +371,43 @@ var AgUiThreadRuntimeCore = class {
 					custom: newCustom
 				}
 			};
-		});
-		if (touched) this.notifyUpdate();
+		})) this.notifyUpdate();
 	}
 	findMessageIdForToolCall(toolCallId) {
 		let fallbackMessageId;
-		for (let index = this.messages.length - 1; index >= 0; index--) {
-			const message = this.messages[index];
+		const messages = this.getMessages();
+		for (let index = messages.length - 1; index >= 0; index--) {
+			const message = messages[index];
 			if (!message || message.role !== "assistant") continue;
 			for (const part of message.content) {
 				if (part.type !== "tool-call" || part.toolCallId !== toolCallId) continue;
-				if (!("result" in part) || part.result === void 0) return message.id;
+				if (!isResolvedToolCall(part)) return message.id;
 				fallbackMessageId ??= message.id;
 			}
 		}
 		return fallbackMessageId;
 	}
+	cancelUnresolvedToolCalls(messageId) {
+		if (this.updateMessage(messageId, (message) => {
+			if (message.role !== "assistant") return message;
+			const assistant = message;
+			const content = assistant.content.map((part) => {
+				if (part.type !== "tool-call" || isResolvedToolCall(part)) return part;
+				return {
+					...part,
+					result: { error: "Tool call cancelled by user" },
+					isError: true
+				};
+			});
+			return {
+				...assistant,
+				content
+			};
+		})) this.notifyUpdate();
+	}
 	addToolResult(options) {
-		let updated = false;
-		this.messages = this.messages.map((message) => {
-			if (message.id !== options.messageId || message.role !== "assistant") return message;
+		if (!this.updateMessage(options.messageId, (message) => {
+			if (message.role !== "assistant") return message;
 			const assistant = message;
 			let matchedToolCall = false;
 			const content = assistant.content.map((part) => {
@@ -293,13 +421,11 @@ var AgUiThreadRuntimeCore = class {
 				};
 			});
 			if (!matchedToolCall) return message;
-			updated = true;
 			return {
 				...assistant,
 				content
 			};
-		});
-		if (!updated) return;
+		})) return;
 		this.notifyUpdate();
 		this.maybeResumeAfterToolResults(options.messageId);
 	}
@@ -312,18 +438,18 @@ var AgUiThreadRuntimeCore = class {
 		this.startResumeRun(messageId);
 	}
 	maybeCompleteAfterToolResults(messageId) {
-		const message = this.messages.find((m) => m.id === messageId);
+		const message = this.tryGetMessage(messageId)?.message;
 		if (!message || message.role !== "assistant") return false;
 		const assistant = message;
 		if (assistant.status?.type !== "requires-action" || assistant.status.reason !== "tool-calls") return false;
-		if (!assistant.content.every((part) => part.type !== "tool-call" || "result" in part && part.result !== void 0)) return false;
-		this.messages = this.messages.map((m) => m.id === messageId && m.role === "assistant" ? {
-			...m,
+		if (!assistant.content.every((part) => part.type !== "tool-call" || isResolvedToolCall(part))) return false;
+		if (!this.updateMessage(messageId, (current) => current.role === "assistant" ? {
+			...current,
 			status: {
 				type: "complete",
 				reason: "unknown"
 			}
-		} : m);
+		} : current)) return false;
 		this.notifyUpdate();
 		this.persistAssistantHistory(messageId);
 		return true;
@@ -335,30 +461,122 @@ var AgUiThreadRuntimeCore = class {
 	}
 	applyExternalMessages(messages) {
 		this.assistantHistoryParents.clear();
-		this.messages = [...messages];
+		if (messages.length === 0) this.clearRepository();
+		else {
+			let expectedParentId = null;
+			let lastAppliedId = null;
+			let hardReplace = false;
+			const seen = /* @__PURE__ */ new Set();
+			for (const message of messages) {
+				if (seen.has(message.id)) continue;
+				seen.add(message.id);
+				const existing = this.tryGetMessage(message.id);
+				if (existing && existing.parentId !== expectedParentId) {
+					hardReplace = true;
+					break;
+				}
+				this.addOrUpdateMessage(expectedParentId, message);
+				expectedParentId = message.id;
+				lastAppliedId = message.id;
+			}
+			if (hardReplace) {
+				this.clearRepository();
+				expectedParentId = null;
+				lastAppliedId = null;
+				seen.clear();
+				for (const message of messages) {
+					if (seen.has(message.id)) continue;
+					seen.add(message.id);
+					this.addOrUpdateMessage(expectedParentId, message);
+					expectedParentId = message.id;
+					lastAppliedId = message.id;
+				}
+			}
+			this.resetRepositoryHead(lastAppliedId);
+		}
+		this.recordedHistoryIds.clear();
+		for (const { message } of this.getMessageRepository().messages) this.recordedHistoryIds.add(message.id);
+		this.notifyUpdate();
+	}
+	applyExternalMessageRepository(loaded) {
+		const headId = loaded.headId ?? loaded.messages.at(-1)?.message.id ?? null;
+		const ids = /* @__PURE__ */ new Set();
+		let degenerate = false;
+		for (const { message } of loaded.messages) {
+			if (ids.has(message.id)) {
+				degenerate = true;
+				break;
+			}
+			ids.add(message.id);
+		}
+		if (headId !== null && !ids.has(headId)) degenerate = true;
+		if (!degenerate) {
+			this.clearRepository();
+			let pending = [...loaded.messages];
+			const importedIds = /* @__PURE__ */ new Set();
+			while (pending.length > 0) {
+				const unresolved = [];
+				let progressed = false;
+				for (const item of pending) {
+					if (item.parentId !== null && !importedIds.has(item.parentId)) {
+						unresolved.push(item);
+						continue;
+					}
+					this.addOrUpdateMessage(item.parentId, item.message);
+					importedIds.add(item.message.id);
+					progressed = true;
+				}
+				if (!progressed) {
+					degenerate = true;
+					break;
+				}
+				pending = unresolved;
+			}
+		}
+		if (degenerate) {
+			this.clearRepository();
+			let previousId = null;
+			for (const { message } of loaded.messages) {
+				const existing = this.tryGetMessage(message.id);
+				this.addOrUpdateMessage(existing ? existing.parentId : previousId, message);
+				previousId = message.id;
+			}
+			this.resetRepositoryHead(previousId);
+		} else this.resetRepositoryHead(headId);
+		this.assistantHistoryParents.clear();
 		this.recordedHistoryIds.clear();
-		for (const message of this.messages) this.recordedHistoryIds.add(message.id);
+		for (const { message } of loaded.messages) this.recordedHistoryIds.add(message.id);
 		this.notifyUpdate();
 	}
 	loadExternalState(state) {
 		this.stateSnapshot = state;
 		this.notifyUpdate();
 	}
+	setState(next) {
+		this.stateSnapshot = typeof next === "function" ? next(this.stateSnapshot) : next;
+		this.notifyUpdate();
+	}
+	resetState() {
+		this.stateSnapshot = void 0;
+		this.notifyUpdate();
+	}
 	async startRun(parentId, runConfig, resume, resumeStream) {
 		const normalizedRunConfig = runConfig ?? {};
 		this.lastRunConfig = normalizedRunConfig;
-		this.resetHead(parentId);
-		const historicalMessages = [...this.messages];
+		const parent = parentId === null ? void 0 : this.tryGetMessage(parentId);
+		const shouldEagerlyInsertAssistant = parentId !== null && parent !== void 0 && parentId !== this.repository.headId;
+		const historicalMessages = [...shouldEagerlyInsertAssistant ? this.tryGetMessages(parentId) ?? this.repository.getMessages() : this.repository.getMessages()];
 		this.pendingError = null;
-		const assistantParentId = parentId ?? this.messages.at(-1)?.id ?? null;
+		const assistantParentId = parent ? parentId : this.repository.headId;
 		let assistantMessageId;
 		const ensureAssistant = () => {
 			if (assistantMessageId) return assistantMessageId;
-			const created = this.insertAssistantPlaceholder();
+			const created = this.insertAssistantPlaceholder(assistantParentId ?? null);
 			assistantMessageId = created;
 			this.markPendingAssistantHistory(created, assistantParentId ?? null);
 			return created;
 		};
+		if (shouldEagerlyInsertAssistant) ensureAssistant();
 		const applyUpdate = (update) => {
 			const resolved = this.updateAssistantMessage(ensureAssistant(), update);
 			if (resolved !== assistantMessageId) assistantMessageId = resolved;
@@ -370,8 +588,7 @@ var AgUiThreadRuntimeCore = class {
 			onServerMessageId: (serverId) => {
 				const placeholder = ensureAssistant();
 				if (placeholder === serverId) return;
-				this.reassignAssistantId(placeholder, serverId);
-				assistantMessageId = serverId;
+				if (this.reassignAssistantId(placeholder, serverId)) assistantMessageId = serverId;
 			}
 		});
 		const dispatch = (event) => this.handleEvent(aggregator, event);
@@ -461,7 +678,7 @@ var AgUiThreadRuntimeCore = class {
 			unstable_threadId: ctx.threadId,
 			unstable_parentId: ctx.parentId,
 			unstable_getMessage: () => {
-				const message = this.messages.find((m) => m.id === currentId());
+				const message = this.tryGetMessage(currentId())?.message;
 				if (!message) throw new Error("[agui] resume stream requested the assistant message before it existed");
 				return message;
 			}
@@ -484,7 +701,7 @@ var AgUiThreadRuntimeCore = class {
 			return;
 		}
 		if (ctx.abortSignal.aborted) return;
-		const current = this.messages.find((m) => m.id === currentId());
+		const current = this.tryGetMessage(currentId())?.message;
 		if (!current || current.status?.type === "running") ctx.applyUpdate({ status: {
 			type: "complete",
 			reason: "unknown"
@@ -492,7 +709,7 @@ var AgUiThreadRuntimeCore = class {
 	}
 	buildRunInput(runId, runConfig, historyMessages, resume) {
 		const threadId = this.agent.threadId || "main";
-		const messages = toAgUiMessages(historyMessages ?? this.messages);
+		const messages = toAgUiMessages(historyMessages ?? this.repository.getMessages());
 		const context = this.runtime?.thread.getModelContext();
 		return {
 			threadId,
@@ -535,7 +752,7 @@ var AgUiThreadRuntimeCore = class {
 		if (this.abortController === controller) this.abortController = null;
 		this.setRunning(false);
 	}
-	insertAssistantPlaceholder() {
+	insertAssistantPlaceholder(parentId) {
 		const id = generateOptimisticId();
 		const assistant = {
 			id,
@@ -552,44 +769,48 @@ var AgUiThreadRuntimeCore = class {
 				custom: {}
 			}
 		};
-		this.messages = [...this.messages, assistant];
+		this.addOrUpdateMessage(parentId ?? this.repository.headId, assistant);
+		this.switchToBranch(id);
 		this.notifyUpdate();
 		return id;
 	}
 	reassignAssistantId(oldId, newId) {
-		if (oldId === newId) return;
-		if (this.messages.some((m) => m.id === newId)) {
-			this.logger.debug?.("[agui] reassignAssistantId: server id already present in messages, dropping placeholder", {
+		if (oldId === newId) return true;
+		const oldItem = this.tryGetMessage(oldId);
+		if (!oldItem) return false;
+		const collidesWithExisting = this.hasMessage(newId);
+		if (collidesWithExisting) {
+			this.logger.debug?.("[agui] reassignAssistantId: server id already present in messages or repository, dropping placeholder", {
 				oldId,
 				newId
 			});
-			this.messages = this.messages.filter((m) => m.id !== oldId);
-		} else this.messages = this.messages.map((m) => {
-			if (m.id !== oldId) return m;
-			const { isOptimistic: _, ...metadata } = m.metadata;
-			return {
-				...m,
+			this.deleteMessage(oldId, oldItem.parentId ?? null);
+		} else {
+			const { isOptimistic: _, ...metadata } = oldItem.message.metadata;
+			this.addOrUpdateMessage(oldItem.parentId, {
+				...oldItem.message,
 				id: newId,
 				metadata
-			};
-		});
+			});
+			this.switchToBranch(newId);
+			this.tryDeleteMessage(oldId);
+		}
 		const pendingParent = this.assistantHistoryParents.get(oldId);
 		if (pendingParent !== void 0) {
 			this.assistantHistoryParents.delete(oldId);
-			if (!this.assistantHistoryParents.has(newId)) this.assistantHistoryParents.set(newId, pendingParent);
+			if (!collidesWithExisting && !this.assistantHistoryParents.has(newId)) this.assistantHistoryParents.set(newId, pendingParent);
 		}
 		if (this.recordedHistoryIds.has(oldId)) {
 			this.recordedHistoryIds.delete(oldId);
-			this.recordedHistoryIds.add(newId);
+			if (!collidesWithExisting) this.recordedHistoryIds.add(newId);
 		}
 		this.notifyUpdate();
+		return !collidesWithExisting;
 	}
 	updateAssistantMessage(messageId, update) {
-		let touched = false;
 		let latestStatus;
-		this.messages = this.messages.map((message) => {
-			if (message.id !== messageId || message.role !== "assistant") return message;
-			touched = true;
+		if (!this.updateMessage(messageId, (message) => {
+			if (message.role !== "assistant") return message;
 			const assistant = message;
 			const metadata = update.metadata ? this.mergeAssistantMetadata(assistant.metadata, update.metadata) : assistant.metadata;
 			latestStatus = update.status ?? assistant.status;
@@ -600,8 +821,7 @@ var AgUiThreadRuntimeCore = class {
 				status: latestStatus,
 				metadata
 			};
-		});
-		if (!touched) return messageId;
+		})) return messageId;
 		let resolvedMessageId = messageId;
 		if (latestStatus !== void 0 && latestStatus.type !== "running" && isOptimisticId(messageId)) {
 			const stableId = generateId();
@@ -614,12 +834,11 @@ var AgUiThreadRuntimeCore = class {
 	}
 	preserveToolResults(previous, next) {
 		const resolved = /* @__PURE__ */ new Map();
-		for (const part of previous) if (part.type === "tool-call" && "result" in part && part.result !== void 0) resolved.set(part.toolCallId, part);
+		for (const part of previous) if (part.type === "tool-call" && isResolvedToolCall(part)) resolved.set(part.toolCallId, part);
 		if (resolved.size === 0) return next;
 		let changed = false;
 		const merged = next.map((part) => {
-			if (part.type !== "tool-call") return part;
-			if ("result" in part && part.result !== void 0) return part;
+			if (part.type !== "tool-call" || isResolvedToolCall(part)) return part;
 			const prior = resolved.get(part.toolCallId);
 			if (!prior) return part;
 			changed = true;
@@ -683,9 +902,8 @@ var AgUiThreadRuntimeCore = class {
 		}
 	}
 	applyCrossRunToolResult(messageId, event) {
-		let updated = false;
-		this.messages = this.messages.map((message) => {
-			if (message.id !== messageId || message.role !== "assistant") return message;
+		if (!this.updateMessage(messageId, (message) => {
+			if (message.role !== "assistant") return message;
 			const assistant = message;
 			let matchedToolCall = false;
 			const content = assistant.content.map((part) => {
@@ -699,13 +917,11 @@ var AgUiThreadRuntimeCore = class {
 				};
 			});
 			if (!matchedToolCall) return message;
-			updated = true;
 			return {
 				...assistant,
 				content
 			};
-		});
-		if (!updated) return;
+		})) return;
 		this.notifyUpdate();
 		this.maybeCompleteAfterToolResults(messageId);
 	}
@@ -726,15 +942,6 @@ var AgUiThreadRuntimeCore = class {
 	toThreadMessage(message) {
 		return fromThreadMessageLike(message, generateId(), FALLBACK_USER_STATUS);
 	}
-	resetHead(parentId) {
-		if (!parentId) {
-			if (this.messages.length) this.messages = [];
-			return;
-		}
-		const idx = this.messages.findIndex((message) => message.id === parentId);
-		if (idx === -1) return;
-		this.messages = this.messages.slice(0, idx + 1);
-	}
 	isTerminalStatus(status) {
 		return status?.type === "complete" || status?.type === "incomplete";
 	}
@@ -753,7 +960,7 @@ var AgUiThreadRuntimeCore = class {
 		if (!this.history) return;
 		const parentId = this.assistantHistoryParents.get(messageId);
 		if (parentId === void 0) return;
-		const message = this.messages.find((m) => m.id === messageId);
+		const message = this.tryGetMessage(messageId)?.message;
 		if (!message || message.role !== "assistant") return;
 		if (!this.isPersistableStatus(message.status)) return;
 		this.assistantHistoryParents.delete(messageId);