@assistant-ui/react-ag-ui 0.0.54 → 0.0.57

package/dist/runtime/AgUiThreadRuntimeCore.js

@@ -1,10 +1,11 @@
 "use client";
 import { readMcpAppResourceUri } from "./mcp-tool-result.js";
+import { buildToolApprovalResume, projectAgUiToolApprovals, withSettledToolApprovals, withToolApprovalDecision } from "./adapter/tool-approval.js";
 import { AG_UI_METADATA_NAMESPACE, RunAggregator, isPlainObject, tryParseJSON } from "./adapter/run-aggregator.js";
 import { fromAgUiMessages, toAgUiMessages, toAgUiTools } from "./adapter/conversions.js";
 import { createAgUiSubscriber } from "./adapter/subscriber.js";
 import { fromThreadMessageLike, generateId, isMcpAppUri } from "@assistant-ui/core";
-import { MessageRepository } from "@assistant-ui/core/internal";
+import { createMessageRepositorySession, invokeUserCallback } from "@assistant-ui/core/internal";
 import jsonpatch from "fast-json-patch";
 //#region src/runtime/AgUiThreadRuntimeCore.ts
 const optimisticPrefix = "__optimistic__";
@@ -16,19 +17,8 @@ const FALLBACK_USER_STATUS = {
 	type: "complete",
 	reason: "unknown"
 };
-const reportCallbackError = (name, error) => {
-	console.error(`[react-ag-ui] ${name} callback threw an error`, error);
-};
 const invokeRuntimeCallback = (name, callback, ...args) => {
-	if (!callback) return;
-	try {
-		const result = callback(...args);
-		if (result !== null && (typeof result === "object" || typeof result === "function") && "then" in result && typeof result.then === "function") Promise.resolve(result).catch((error) => {
-			reportCallbackError(name, error);
-		});
-	} catch (error) {
-		reportCallbackError(name, error);
-	}
+	invokeUserCallback("react-ag-ui", name, callback, ...args);
 };
 const mergeAgUiNamespace = (current, incoming) => {
 	const merged = {
@@ -51,9 +41,22 @@ var AgUiThreadRuntimeCore = class {
 	onError;
 	onCancel;
 	notifyUpdate;
+	reportedErrors = /* @__PURE__ */ new WeakSet();
 	runtime;
-	repository = new MessageRepository();
-	exportedRepository;
+	session = createMessageRepositorySession({ decorateExport: (exported, repository) => {
+		let parentId = null;
+		for (const message of repository.getMessages()) {
+			if (message.metadata.isOptimistic) exported.messages.push({
+				parentId,
+				message
+			});
+			parentId = message.id;
+		}
+		return {
+			...exported,
+			headId: repository.headId
+		};
+	} });
 	isRunningFlag = false;
 	abortController = null;
 	activeRunAgent = null;
@@ -62,9 +65,12 @@ var AgUiThreadRuntimeCore = class {
 	history;
 	lastRunConfig;
 	assistantHistoryParents = /* @__PURE__ */ new Map();
-	recordedHistoryIds = /* @__PURE__ */ new Set();
+	snapshotHistoryIds = /* @__PURE__ */ new Set();
+	persistedHistoryIds = /* @__PURE__ */ new Set();
+	historyWrites = /* @__PURE__ */ new Map();
 	_isLoading = false;
 	_loadPromise;
+	_loadRequested = false;
 	pendingResumeMessageId = null;
 	pendingA2uiResume = false;
 	pendingA2uiAction;
@@ -86,84 +92,25 @@ var AgUiThreadRuntimeCore = class {
 		this.autoCancelPendingToolCalls = options.autoCancelPendingToolCalls;
 		this.onError = options.onError;
 		this.onCancel = options.onCancel;
+		const previousHistory = this.history;
 		this.history = options.history;
 		this.installResumeShim();
+		if (this._loadRequested && !this._loadPromise && !previousHistory && options.history && this.session.getMessages().length === 0) this.__internal_load();
 	}
 	attachRuntime(runtime) {
 		this.runtime = runtime;
 	}
 	detachRuntime() {
 		this.runtime = void 0;
+		this.cancel().catch((error) => {
+			this.logger.error("[agui] failed to cancel run during teardown", error);
+		});
 	}
 	getMessages() {
-		return this.repository.getMessages();
+		return this.session.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;
+		return this.session.export();
 	}
 	getState() {
 		return this.stateSnapshot;
@@ -175,15 +122,22 @@ var AgUiThreadRuntimeCore = class {
 		return this._isLoading;
 	}
 	__internal_load() {
+		this._loadRequested = true;
 		if (this._loadPromise) return this._loadPromise;
-		const promise = this.history?.load() ?? Promise.resolve(null);
+		if (!this.history) return Promise.resolve();
+		const promise = this.history.load();
 		this._isLoading = true;
 		this._loadPromise = promise.then(async (repo) => {
 			if (!repo) return;
-			this.applyExternalMessageRepository(repo);
+			this.session.applyExternalMessageRepository(repo);
+			this.assistantHistoryParents.clear();
+			this.snapshotHistoryIds.clear();
+			this.persistedHistoryIds.clear();
+			for (const { message } of repo.messages) this.persistedHistoryIds.add(message.id);
+			this.notifyUpdate();
 			if (repo.state !== void 0) this.loadExternalState(repo.state);
 			if (repo.unstable_resume) {
-				const parentId = repo.headId ?? this.repository.headId;
+				const parentId = repo.headId ?? this.session.headId;
 				const resumeStream = this.history?.resume?.bind(this.history);
 				await this.startRun(parentId, this.lastRunConfig, void 0, resumeStream);
 			}
@@ -215,11 +169,11 @@ var AgUiThreadRuntimeCore = class {
 		this.maybeCompleteAfterToolResults(pending.messageId);
 	}
 	appendEntry(message) {
-		if (message.sourceId) this.tryDeleteMessage(message.sourceId);
+		if (message.sourceId) this.session.tryDeleteMessage(message.sourceId);
 		const threadMessage = this.toThreadMessage(message);
-		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);
+		const parentId = message.parentId === null ? null : message.parentId && this.session.hasMessage(message.parentId) ? message.parentId : this.session.headId;
+		this.session.addOrUpdateMessage(parentId, threadMessage);
+		this.session.switchToBranch(threadMessage.id);
 		this.notifyUpdate();
 		this.recordHistoryEntry(parentId, threadMessage);
 		return threadMessage.id;
@@ -298,17 +252,65 @@ var AgUiThreadRuntimeCore = class {
 		const openIds = pending.interrupts.map((i) => i.id);
 		const missing = openIds.filter((id) => !responsesById.has(id));
 		if (missing.length > 0) throw new Error(`[agui] submitInterruptResponses: missing responses for open interrupts: ${missing.join(", ")}`);
+		this.assertInterruptsAnswerable("submitInterruptResponses", pending.interrupts);
+		await this.resumeWithResponses(pending.messageId, openIds.map((id) => responsesById.get(id)));
+	}
+	/**
+	* Resumes with an already validated resume array. Answerability is checked
+	* once per submission: re-checking here would let a clock crossing reject a
+	* decision this runtime has already recorded, stranding the gate decided and
+	* unretryable.
+	*/
+	async resumeWithResponses(messageId, resume) {
+		this.clearPendingInterrupts(messageId, resume);
+		await this.startRun(messageId, this.lastRunConfig, resume);
+	}
+	/**
+	* The core seams that reach `respondToToolApproval` discard the promise they
+	* receive, so a rejected decision would otherwise surface only as an
+	* unhandled rejection. A failure raised by the resumed run itself was already
+	* reported by `startRun` before it rethrew, so reporting it here again would
+	* give the consumer two notifications for one failure.
+	*/
+	reportError(error) {
+		if (this.reportedErrors.has(error)) {
+			this.reportedErrors.delete(error);
+			return;
+		}
+		invokeRuntimeCallback("onError", this.onError, error instanceof Error ? error : new Error(String(error)));
+	}
+	assertInterruptsAnswerable(method, interrupts) {
 		const now = Date.now();
-		for (const interrupt of pending.interrupts) {
+		for (const interrupt of interrupts) {
 			if (!interrupt.expiresAt) continue;
 			const expiry = new Date(interrupt.expiresAt).getTime();
-			if (Number.isNaN(expiry)) throw new Error(`[agui] submitInterruptResponses: interrupt ${interrupt.id} has malformed expiresAt "${interrupt.expiresAt}"`);
-			if (expiry <= now) throw new Error(`[agui] submitInterruptResponses: interrupt ${interrupt.id} expired at ${interrupt.expiresAt}`);
+			if (Number.isNaN(expiry)) throw new Error(`[agui] ${method}: interrupt ${interrupt.id} has malformed expiresAt "${interrupt.expiresAt}"`);
+			if (expiry <= now) throw new Error(`[agui] ${method}: interrupt ${interrupt.id} expired at ${interrupt.expiresAt}`);
 		}
-		const resume = openIds.map((id) => responsesById.get(id));
-		if (this.isRunningFlag) throw new Error("[agui] submitInterruptResponses: a run is already in progress");
-		this.clearPendingInterrupts(pending.messageId);
-		await this.startRun(pending.messageId, this.lastRunConfig, resume);
+		if (this.isRunningFlag) throw new Error(`[agui] ${method}: a run is already in progress`);
+	}
+	async respondToToolApproval(options) {
+		const pending = this.getPendingInterrupts();
+		if (!pending) throw new Error("[agui] respondToToolApproval: no pending interrupts on this thread");
+		const gatedMessage = this.session.tryGetMessage(pending.messageId)?.message;
+		if (![...projectAgUiToolApprovals(pending.interrupts, new Set((gatedMessage?.content ?? []).filter((part) => part.type === "tool-call").map((part) => part.toolCallId))).values()].some((approval) => approval.id === options.approvalId)) throw new Error(`[agui] respondToToolApproval: no pending tool-call interrupt for approval id "${options.approvalId}"`);
+		this.assertInterruptsAnswerable("respondToToolApproval", pending.interrupts);
+		if (!this.session.updateMessage(pending.messageId, (message) => {
+			if (message.role !== "assistant") return message;
+			const assistant = message;
+			const content = withToolApprovalDecision(assistant.content, options);
+			if (content === assistant.content) return assistant;
+			return {
+				...assistant,
+				content
+			};
+		})) throw new Error(`[agui] respondToToolApproval: approval "${options.approvalId}" is already decided`);
+		this.notifyUpdate();
+		const assistant = this.session.tryGetMessage(pending.messageId)?.message;
+		if (!assistant) return;
+		const resume = buildToolApprovalResume(assistant.content, pending.interrupts);
+		if (!resume) return;
+		await this.resumeWithResponses(pending.messageId, resume);
 	}
 	async steerAway(message, responses) {
 		const pending = this.getPendingInterrupts();
@@ -331,7 +333,7 @@ var AgUiThreadRuntimeCore = class {
 		const resume = this.resolveSteerAwayResume(pending.interrupts, responses);
 		if (this.isRunningFlag) throw new Error("[agui] steerAway: a run is already in progress");
 		const normalized = this.toAppendMessage(message);
-		this.clearPendingInterrupts(pending.messageId);
+		this.clearPendingInterrupts(pending.messageId, resume);
 		const threadMessageId = this.appendEntry(normalized);
 		await this.startRun(threadMessageId, normalized.runConfig, resume);
 	}
@@ -357,7 +359,7 @@ var AgUiThreadRuntimeCore = class {
 	toAppendMessage(message) {
 		if (typeof message === "string") return {
 			createdAt: /* @__PURE__ */ new Date(),
-			parentId: this.repository.headId,
+			parentId: this.session.headId,
 			sourceId: null,
 			runConfig: {},
 			role: "user",
@@ -370,7 +372,7 @@ var AgUiThreadRuntimeCore = class {
 		};
 		return {
 			createdAt: message.createdAt ?? /* @__PURE__ */ new Date(),
-			parentId: message.parentId ?? this.repository.headId,
+			parentId: message.parentId ?? this.session.headId,
 			sourceId: message.sourceId ?? null,
 			role: message.role ?? "user",
 			content: message.content,
@@ -380,8 +382,8 @@ var AgUiThreadRuntimeCore = class {
 			startRun: message.startRun
 		};
 	}
-	clearPendingInterrupts(messageId) {
-		if (this.updateMessage(messageId, (message) => {
+	clearPendingInterrupts(messageId, resume) {
+		if (this.session.updateMessage(messageId, (message) => {
 			if (message.role !== "assistant") return message;
 			const assistant = message;
 			if (assistant.status?.type !== "requires-action" || assistant.status.reason !== "interrupt") return assistant;
@@ -391,6 +393,7 @@ var AgUiThreadRuntimeCore = class {
 			else delete newCustom["agui"];
 			return {
 				...assistant,
+				content: withSettledToolApprovals(assistant.content, resume),
 				status: {
 					type: "complete",
 					reason: "unknown"
@@ -417,7 +420,7 @@ var AgUiThreadRuntimeCore = class {
 		return fallbackMessageId;
 	}
 	cancelUnresolvedToolCalls(messageId) {
-		if (this.updateMessage(messageId, (message) => {
+		if (this.session.updateMessage(messageId, (message) => {
 			if (message.role !== "assistant") return message;
 			const assistant = message;
 			const content = assistant.content.map((part) => {
@@ -435,7 +438,7 @@ var AgUiThreadRuntimeCore = class {
 		})) this.notifyUpdate();
 	}
 	addToolResult(options) {
-		if (!this.updateMessage(options.messageId, (message) => {
+		if (!this.session.updateMessage(options.messageId, (message) => {
 			if (message.role !== "assistant") return message;
 			const assistant = message;
 			let matchedToolCall = false;
@@ -461,7 +464,7 @@ var AgUiThreadRuntimeCore = class {
 	sendA2uiAction(action) {
 		this.assertNoPendingInterrupts();
 		this.maybeAutoCancelPendingToolCalls();
-		const parentId = this.repository.headId;
+		const parentId = this.session.headId;
 		if (parentId === null) {
 			this.logger.debug("[agui] sendA2uiAction: no messages to resume, dropping action");
 			return;
@@ -485,12 +488,12 @@ var AgUiThreadRuntimeCore = class {
 		this.startResumeRun(messageId);
 	}
 	maybeCompleteAfterToolResults(messageId) {
-		const message = this.tryGetMessage(messageId)?.message;
+		const message = this.session.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" || isResolvedToolCall(part))) return false;
-		if (!this.updateMessage(messageId, (current) => current.role === "assistant" ? {
+		if (!this.session.updateMessage(messageId, (current) => current.role === "assistant" ? {
 			...current,
 			status: {
 				type: "complete",
@@ -510,7 +513,7 @@ var AgUiThreadRuntimeCore = class {
 		this.pendingA2uiResume = false;
 		this.pendingA2uiAction = void 0;
 		this.assistantHistoryParents.clear();
-		if (messages.length === 0) this.clearRepository();
+		if (messages.length === 0) this.session.clear();
 		else {
 			let expectedParentId = null;
 			let lastAppliedId = null;
@@ -519,82 +522,32 @@ var AgUiThreadRuntimeCore = class {
 			for (const message of messages) {
 				if (seen.has(message.id)) continue;
 				seen.add(message.id);
-				const existing = this.tryGetMessage(message.id);
+				const existing = this.session.tryGetMessage(message.id);
 				if (existing && existing.parentId !== expectedParentId) {
 					hardReplace = true;
 					break;
 				}
-				this.addOrUpdateMessage(expectedParentId, message);
+				this.session.addOrUpdateMessage(expectedParentId, message);
 				expectedParentId = message.id;
 				lastAppliedId = message.id;
 			}
 			if (hardReplace) {
-				this.clearRepository();
+				this.session.clear();
 				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);
+					this.session.addOrUpdateMessage(expectedParentId, message);
 					expectedParentId = message.id;
 					lastAppliedId = message.id;
 				}
 			}
-			this.resetRepositoryHead(lastAppliedId);
+			this.session.resetHead(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 loaded.messages) this.recordedHistoryIds.add(message.id);
+		this.snapshotHistoryIds.clear();
+		for (const { message } of this.getMessageRepository().messages) this.snapshotHistoryIds.add(message.id);
 		this.notifyUpdate();
 	}
 	loadExternalState(state) {
@@ -612,18 +565,20 @@ var AgUiThreadRuntimeCore = class {
 	async startRun(parentId, runConfig, resume, resumeStream) {
 		const normalizedRunConfig = runConfig ?? {};
 		this.lastRunConfig = normalizedRunConfig;
-		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()];
+		const parent = parentId === null ? void 0 : this.session.tryGetMessage(parentId);
+		const shouldEagerlyInsertAssistant = parentId !== null && parent !== void 0 && parentId !== this.session.headId;
+		const historicalMessages = [...shouldEagerlyInsertAssistant ? this.session.tryGetMessages(parentId) ?? this.session.getMessages() : this.session.getMessages()];
+		const runStartMessageIds = new Set(this.getMessageRepository().messages.map(({ message }) => message.id));
 		this.pendingError = null;
-		const assistantParentId = parent ? parentId : this.repository.headId;
+		const assistantParentId = parent ? parentId : this.session.headId;
 		let assistantMessageId;
 		let assistantCollided = false;
 		const ensureAssistant = (allowRecreate = false) => {
 			const cached = assistantMessageId;
-			if (cached !== void 0 && this.tryGetMessage(cached)) return cached;
+			if (cached !== void 0 && this.session.tryGetMessage(cached)) return cached;
 			if (cached !== void 0 && (assistantCollided || !allowRecreate)) return cached;
-			const parentId = cached === void 0 && assistantParentId && this.hasMessage(assistantParentId) ? assistantParentId : this.repository.headId;
+			const repositoryHeadId = this.session.headId;
+			const parentId = (cached === void 0 || shouldEagerlyInsertAssistant && (repositoryHeadId === null || runStartMessageIds.has(repositoryHeadId))) && assistantParentId && this.session.hasMessage(assistantParentId) ? assistantParentId : repositoryHeadId;
 			const created = this.insertAssistantPlaceholder(parentId, cached);
 			assistantMessageId = created;
 			this.markPendingAssistantHistory(created, parentId);
@@ -631,20 +586,44 @@ var AgUiThreadRuntimeCore = class {
 		};
 		if (shouldEagerlyInsertAssistant) ensureAssistant();
 		const applyUpdate = (update) => {
-			const hasContent = Array.isArray(update.content) && update.content.length > 0;
-			const resolved = this.updateAssistantMessage(ensureAssistant(hasContent), update);
+			const hasStreamContent = Array.isArray(update.content) && update.content.some((part) => part.type !== "data");
+			const resolved = this.updateAssistantMessage(ensureAssistant(hasStreamContent), update);
 			if (resolved !== assistantMessageId) assistantMessageId = resolved;
 		};
+		const adoptServerMessageId = (serverId, startNewMessage) => {
+			if (startNewMessage && assistantMessageId !== void 0) {
+				applyUpdate({ status: {
+					type: "complete",
+					reason: "unknown"
+				} });
+				const previousId = assistantMessageId;
+				assistantMessageId = void 0;
+				if (this.session.tryGetMessage(previousId)) {
+					const created = this.insertAssistantPlaceholder(previousId);
+					assistantMessageId = created;
+					this.markPendingAssistantHistory(created, previousId);
+				}
+			}
+			const placeholder = ensureAssistant(true);
+			if (placeholder === serverId) return;
+			const reassigned = this.reassignAssistantId(placeholder, serverId);
+			const adoptsVisibleCollision = !reassigned && !runStartMessageIds.has(serverId) && this.session.headId === serverId;
+			if (reassigned || adoptsVisibleCollision) {
+				assistantMessageId = serverId;
+				if (adoptsVisibleCollision) {
+					const parentId = this.session.tryGetMessage(serverId)?.parentId ?? null;
+					this.markPendingAssistantHistory(serverId, parentId);
+				}
+			} else assistantCollided = true;
+		};
 		const aggregator = new RunAggregator({
 			showThinking: this.showThinking,
 			logger: this.logger,
 			emit: applyUpdate,
 			onServerMessageId: (serverId) => {
-				const placeholder = ensureAssistant(true);
-				if (placeholder === serverId) return;
-				if (this.reassignAssistantId(placeholder, serverId)) assistantMessageId = serverId;
-				else assistantCollided = true;
-			}
+				adoptServerMessageId(serverId, false);
+			},
+			onTextMessageStart: (serverId) => adoptServerMessageId(serverId, true)
 		});
 		const dispatch = (event) => this.handleEvent(aggregator, event, assistantMessageId);
 		const abortController = new AbortController();
@@ -714,6 +693,7 @@ var AgUiThreadRuntimeCore = class {
 		}
 		if (this.pendingError) {
 			const err = this.pendingError;
+			this.reportedErrors.add(err);
 			this.pendingError = null;
 			this.pendingResumeMessageId = null;
 			this.pendingA2uiResume = false;
@@ -735,7 +715,7 @@ var AgUiThreadRuntimeCore = class {
 					return;
 				}
 				this.maybeAutoCancelPendingToolCalls();
-				const parentId = this.repository.headId;
+				const parentId = this.session.headId;
 				if (parentId !== null) this.startResumeRun(parentId);
 				else {
 					this.pendingA2uiAction = void 0;
@@ -758,7 +738,7 @@ var AgUiThreadRuntimeCore = class {
 			unstable_threadId: ctx.threadId,
 			unstable_parentId: ctx.parentId,
 			unstable_getMessage: () => {
-				const message = this.tryGetMessage(currentId())?.message;
+				const message = this.session.tryGetMessage(currentId())?.message;
 				if (!message) throw new Error("[agui] resume stream requested the assistant message before it existed");
 				return message;
 			}
@@ -781,7 +761,7 @@ var AgUiThreadRuntimeCore = class {
 			return;
 		}
 		if (ctx.abortSignal.aborted) return;
-		const current = this.tryGetMessage(currentId())?.message;
+		const current = this.session.tryGetMessage(currentId())?.message;
 		if (!current || current.status?.type === "running") ctx.applyUpdate({ status: {
 			type: "complete",
 			reason: "unknown"
@@ -789,7 +769,7 @@ var AgUiThreadRuntimeCore = class {
 	}
 	buildRunInput(runId, runConfig, historyMessages, resume) {
 		const threadId = this.agent.threadId || "main";
-		const messages = toAgUiMessages(historyMessages ?? this.repository.getMessages());
+		const messages = toAgUiMessages(historyMessages ?? this.session.getMessages());
 		const context = this.runtime?.thread.getModelContext();
 		const input = {
 			threadId,
@@ -854,47 +834,62 @@ var AgUiThreadRuntimeCore = class {
 				custom: {}
 			}
 		};
-		this.addOrUpdateMessage(parentId ?? this.repository.headId, assistant);
-		this.switchToBranch(id);
+		this.session.addOrUpdateMessage(parentId ?? this.session.headId, assistant);
+		this.session.switchToBranch(id);
 		this.notifyUpdate();
 		return id;
 	}
 	reassignAssistantId(oldId, newId) {
 		if (oldId === newId) return true;
-		const oldItem = this.tryGetMessage(oldId);
+		const oldItem = this.session.tryGetMessage(oldId);
 		if (!oldItem) return false;
-		const collidesWithExisting = this.hasMessage(newId);
+		const collidesWithExisting = this.session.hasMessage(newId);
 		if (collidesWithExisting) {
 			this.logger.debug?.("[agui] reassignAssistantId: server id already present in messages or repository, dropping placeholder", {
 				oldId,
 				newId
 			});
-			this.deleteMessage(oldId, oldItem.parentId ?? null);
+			this.session.deleteMessage(oldId, oldItem.parentId ?? null);
 		} else {
 			const { isOptimistic: _, ...metadata } = oldItem.message.metadata;
-			this.addOrUpdateMessage(oldItem.parentId, {
+			this.session.addOrUpdateMessage(oldItem.parentId, {
 				...oldItem.message,
 				id: newId,
 				metadata
 			});
-			this.switchToBranch(newId);
-			this.tryDeleteMessage(oldId);
+			this.session.switchToBranch(newId);
+			this.session.tryDeleteMessage(oldId);
 		}
 		const pendingParent = this.assistantHistoryParents.get(oldId);
 		if (pendingParent !== void 0) {
 			this.assistantHistoryParents.delete(oldId);
 			if (!collidesWithExisting && !this.assistantHistoryParents.has(newId)) this.assistantHistoryParents.set(newId, pendingParent);
 		}
-		if (this.recordedHistoryIds.has(oldId)) {
-			this.recordedHistoryIds.delete(oldId);
-			if (!collidesWithExisting) this.recordedHistoryIds.add(newId);
+		for (const ids of [this.snapshotHistoryIds, this.persistedHistoryIds]) if (ids.has(oldId)) {
+			ids.delete(oldId);
+			if (!collidesWithExisting) ids.add(newId);
+		}
+		const pendingWrite = this.historyWrites.get(oldId);
+		if (pendingWrite) {
+			this.historyWrites.delete(oldId);
+			if (!collidesWithExisting) {
+				this.historyWrites.set(newId, pendingWrite);
+				const settle = () => {
+					if (this.historyWrites.get(newId) === pendingWrite) this.historyWrites.delete(newId);
+				};
+				pendingWrite.then(() => {
+					this.persistedHistoryIds.delete(oldId);
+					this.persistedHistoryIds.add(newId);
+					settle();
+				}, settle);
+			}
 		}
 		this.notifyUpdate();
 		return !collidesWithExisting;
 	}
 	updateAssistantMessage(messageId, update) {
 		let latestStatus;
-		if (!this.updateMessage(messageId, (message) => {
+		if (!this.session.updateMessage(messageId, (message) => {
 			if (message.role !== "assistant") return message;
 			const assistant = message;
 			const metadata = update.metadata ? this.mergeAssistantMetadata(assistant.metadata, update.metadata) : assistant.metadata;
@@ -996,7 +991,7 @@ var AgUiThreadRuntimeCore = class {
 		}
 	}
 	applyCrossRunToolResult(messageId, event) {
-		if (!this.updateMessage(messageId, (message) => {
+		if (!this.session.updateMessage(messageId, (message) => {
 			if (message.role !== "assistant") return message;
 			const assistant = message;
 			const mcpAppUri = readMcpAppResourceUri(event.mcpResult?._meta);
@@ -1041,7 +1036,7 @@ var AgUiThreadRuntimeCore = class {
 		const serverHash = event.content["serverHash"];
 		const appServerId = typeof serverId === "string" && serverId.length > 0 ? serverId : typeof serverHash === "string" && serverHash.length > 0 ? serverHash : void 0;
 		const result = event.content["result"];
-		if (!this.updateMessage(messageId, (message) => {
+		if (!this.session.updateMessage(messageId, (message) => {
 			if (message.role !== "assistant") return message;
 			const assistant = message;
 			let matchedToolCall = false;
@@ -1075,7 +1070,7 @@ var AgUiThreadRuntimeCore = class {
 	}
 	importMessagesSnapshot(rawMessages, activeAssistantId) {
 		try {
-			const activeMessage = activeAssistantId ? this.tryGetMessage(activeAssistantId)?.message : void 0;
+			const activeMessage = activeAssistantId ? this.session.tryGetMessage(activeAssistantId)?.message : void 0;
 			const activeAssistant = activeMessage?.role === "assistant" ? activeMessage : void 0;
 			const normalized = fromAgUiMessages(rawMessages, { showThinking: this.showThinking });
 			const converted = [];
@@ -1084,14 +1079,16 @@ var AgUiThreadRuntimeCore = class {
 			} catch (error) {
 				this.logger.error?.("[agui] failed to import message from snapshot", error);
 			}
-			const snapshotHeadId = converted.at(-1)?.id ?? null;
 			const snapshotContainsActiveAssistant = converted.some((message) => message.id === activeAssistant?.id);
 			const preservesActiveAssistant = activeAssistant !== void 0 && !snapshotContainsActiveAssistant && (activeAssistant.metadata.isOptimistic !== true || converted.at(-1)?.role !== "assistant");
 			if (preservesActiveAssistant) converted.push(activeAssistant);
 			this.applyExternalMessages(converted);
-			if (preservesActiveAssistant) {
-				this.recordedHistoryIds.delete(activeAssistant.id);
-				this.markPendingAssistantHistory(activeAssistant.id, snapshotHeadId);
+			if (activeAssistant !== void 0) {
+				const activeItem = this.session.tryGetMessage(activeAssistant.id);
+				if (activeItem) {
+					if (preservesActiveAssistant) this.snapshotHistoryIds.delete(activeAssistant.id);
+					this.markPendingAssistantHistory(activeAssistant.id, activeItem.parentId);
+				}
 			}
 		} catch (error) {
 			this.logger.error?.("[agui] failed to import messages snapshot", error);
@@ -1108,32 +1105,82 @@ var AgUiThreadRuntimeCore = class {
 		return status?.type === "requires-action" && status.reason === "interrupt";
 	}
 	recordHistoryEntry(parentId, message) {
-		this.appendHistoryItem(parentId, message);
+		this.appendHistoryItem(parentId, message)?.catch((error) => {
+			this.logger.error?.("[agui] failed to append history entry", error);
+		});
 	}
 	markPendingAssistantHistory(messageId, parentId) {
 		if (!this.history) return;
 		this.assistantHistoryParents.set(messageId, parentId);
 	}
 	persistAssistantHistory(messageId) {
-		if (!this.history) return;
+		const history = this.history;
+		if (!history) return;
 		const parentId = this.assistantHistoryParents.get(messageId);
 		if (parentId === void 0) return;
-		const message = this.tryGetMessage(messageId)?.message;
+		const message = this.session.tryGetMessage(messageId)?.message;
 		if (!message || message.role !== "assistant") return;
 		if (!this.isPersistableStatus(message.status)) return;
+		const wasPersisted = this.persistedHistoryIds.has(messageId);
+		const update = history.update;
+		if (update !== void 0 && (wasPersisted || this.snapshotHistoryIds.has(messageId))) {
+			const write = this.chainHistoryWrite(messageId, () => update.call(history, {
+				parentId,
+				message
+			}));
+			this.assistantHistoryParents.delete(messageId);
+			write.then(() => {
+				this.persistedHistoryIds.add(messageId);
+			}, (error) => {
+				const pending = this.historyWrites.get(messageId);
+				if (pending === void 0 || pending === write) this.assistantHistoryParents.set(messageId, parentId);
+				this.logger.error?.("[agui] failed to update history entry", error);
+			});
+			return;
+		}
+		if (wasPersisted) {
+			this.assistantHistoryParents.delete(messageId);
+			return;
+		}
+		const write = this.appendHistoryItem(parentId, message);
+		if (!write) return;
 		this.assistantHistoryParents.delete(messageId);
-		this.appendHistoryItem(parentId, message);
+		write.then(() => {}, (error) => {
+			const pending = this.historyWrites.get(messageId);
+			if (pending === void 0 || pending === write) this.assistantHistoryParents.set(messageId, parentId);
+			this.logger.error?.("[agui] failed to append history entry", error);
+		});
 	}
 	appendHistoryItem(parentId, message) {
-		if (!this.history || this.recordedHistoryIds.has(message.id)) return;
-		this.recordedHistoryIds.add(message.id);
-		this.history.append({
+		if (!this.history || this.persistedHistoryIds.has(message.id)) return;
+		const pending = this.historyWrites.get(message.id);
+		if (pending) return pending;
+		const append = this.history.append.bind(this.history);
+		const write = this.chainHistoryWrite(message.id, () => append({
 			parentId,
 			message
-		}).catch((error) => {
-			this.recordedHistoryIds.delete(message.id);
-			this.logger.error?.("[agui] failed to append history entry", error);
+		}));
+		write.then(() => {
+			this.persistedHistoryIds.add(message.id);
+		}, () => {});
+		return write;
+	}
+	chainHistoryWrite(id, write) {
+		const pending = this.historyWrites.get(id);
+		let next;
+		if (pending) next = pending.then(write, write);
+		else try {
+			next = Promise.resolve(write());
+		} catch (error) {
+			next = Promise.reject(error);
+		}
+		this.historyWrites.set(id, next);
+		next.then(() => {
+			if (this.historyWrites.get(id) === next) this.historyWrites.delete(id);
+		}, () => {
+			if (this.historyWrites.get(id) === next) this.historyWrites.delete(id);
 		});
+		return next;
 	}
 };
 //#endregion