@assistant-ui/react-ag-ui 0.0.56 → 0.0.58

package/dist/runtime/AgUiThreadRuntimeCore.js

@@ -5,31 +5,19 @@ import { AG_UI_METADATA_NAMESPACE, RunAggregator, isPlainObject, tryParseJSON }
 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, iterateToolCallParts, mapToolCallPartsDeep } from "@assistant-ui/core/internal";
 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",
 	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 = {
@@ -54,8 +42,20 @@ var AgUiThreadRuntimeCore = class {
 	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;
@@ -69,6 +69,7 @@ var AgUiThreadRuntimeCore = class {
 	historyWrites = /* @__PURE__ */ new Map();
 	_isLoading = false;
 	_loadPromise;
+	_loadRequested = false;
 	pendingResumeMessageId = null;
 	pendingA2uiResume = false;
 	pendingA2uiAction;
@@ -81,7 +82,6 @@ var AgUiThreadRuntimeCore = class {
 		this.onCancel = options.onCancel;
 		this.history = options.history;
 		this.notifyUpdate = options.notifyUpdate;
-		this.installResumeShim();
 	}
 	updateOptions(options) {
 		this.agent = options.agent;
@@ -90,8 +90,9 @@ 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;
@@ -103,74 +104,10 @@ var AgUiThreadRuntimeCore = class {
 		});
 	}
 	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;
@@ -182,15 +119,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);
 			}
@@ -222,11 +166,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;
@@ -287,8 +231,7 @@ var AgUiThreadRuntimeCore = class {
 		const assistant = this.findRequiresActionAssistant("tool-calls");
 		if (!assistant) return null;
 		const toolCallIds = [];
-		for (const part of assistant.content) {
-			if (part.type !== "tool-call") continue;
+		for (const part of iterateToolCallParts(assistant.content)) {
 			if (isResolvedToolCall(part)) continue;
 			toolCallIds.push(part.toolCallId);
 		}
@@ -319,9 +262,9 @@ var AgUiThreadRuntimeCore = class {
 		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
+	* A rejected decision reaches the caller through the approval seam, and
+	* `onError` in addition, because a consumer watching only the runtime would
+	* otherwise miss it. 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.
 	*/
@@ -345,10 +288,10 @@ var AgUiThreadRuntimeCore = class {
 	async respondToToolApproval(options) {
 		const pending = this.getPendingInterrupts();
 		if (!pending) throw new Error("[agui] respondToToolApproval: no pending interrupts on this thread");
-		const gatedMessage = this.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}"`);
+		const gatedMessage = this.session.tryGetMessage(pending.messageId)?.message;
+		if (![...projectAgUiToolApprovals(pending.interrupts, new Set(Array.from(iterateToolCallParts(gatedMessage?.content ?? [])).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.updateMessage(pending.messageId, (message) => {
+		if (!this.session.updateMessage(pending.messageId, (message) => {
 			if (message.role !== "assistant") return message;
 			const assistant = message;
 			const content = withToolApprovalDecision(assistant.content, options);
@@ -359,7 +302,7 @@ var AgUiThreadRuntimeCore = class {
 			};
 		})) throw new Error(`[agui] respondToToolApproval: approval "${options.approvalId}" is already decided`);
 		this.notifyUpdate();
-		const assistant = this.tryGetMessage(pending.messageId)?.message;
+		const assistant = this.session.tryGetMessage(pending.messageId)?.message;
 		if (!assistant) return;
 		const resume = buildToolApprovalResume(assistant.content, pending.interrupts);
 		if (!resume) return;
@@ -412,7 +355,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",
@@ -425,7 +368,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,
@@ -436,7 +379,7 @@ var AgUiThreadRuntimeCore = class {
 		};
 	}
 	clearPendingInterrupts(messageId, resume) {
-		if (this.updateMessage(messageId, (message) => {
+		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;
@@ -464,8 +407,8 @@ var AgUiThreadRuntimeCore = class {
 		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;
+			for (const part of iterateToolCallParts(message.content)) {
+				if (part.toolCallId !== toolCallId) continue;
 				if (!isResolvedToolCall(part)) return message.id;
 				fallbackMessageId ??= message.id;
 			}
@@ -473,11 +416,11 @@ 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) => {
-				if (part.type !== "tool-call" || isResolvedToolCall(part)) return part;
+			const { content } = mapToolCallPartsDeep(assistant.content, (part) => {
+				if (isResolvedToolCall(part)) return part;
 				return {
 					...part,
 					result: { error: "Tool call cancelled by user" },
@@ -491,12 +434,14 @@ var AgUiThreadRuntimeCore = class {
 		})) this.notifyUpdate();
 	}
 	addToolResult(options) {
-		if (!this.updateMessage(options.messageId, (message) => {
+		const sessionMessageId = this.session.tryGetMessage(options.messageId) ? options.messageId : this.findMessageIdForToolCall(options.toolCallId);
+		if (sessionMessageId === void 0) return;
+		if (!this.session.updateMessage(sessionMessageId, (message) => {
 			if (message.role !== "assistant") return message;
 			const assistant = message;
 			let matchedToolCall = false;
-			const content = assistant.content.map((part) => {
-				if (part.type !== "tool-call" || part.toolCallId !== options.toolCallId) return part;
+			const { content } = mapToolCallPartsDeep(assistant.content, (part) => {
+				if (part.toolCallId !== options.toolCallId) return part;
 				matchedToolCall = true;
 				return {
 					...part,
@@ -512,12 +457,12 @@ var AgUiThreadRuntimeCore = class {
 			};
 		})) return;
 		this.notifyUpdate();
-		this.maybeResumeAfterToolResults(options.messageId);
+		this.maybeResumeAfterToolResults(sessionMessageId);
 	}
 	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;
@@ -541,12 +486,17 @@ 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" ? {
+		let allResolved = true;
+		for (const part of iterateToolCallParts(assistant.content)) if (!isResolvedToolCall(part)) {
+			allResolved = false;
+			break;
+		}
+		if (!allResolved) return false;
+		if (!this.session.updateMessage(messageId, (current) => current.role === "assistant" ? {
 			...current,
 			status: {
 				type: "complete",
@@ -566,7 +516,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;
@@ -575,85 +525,34 @@ 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.snapshotHistoryIds.clear();
 		for (const { message } of this.getMessageRepository().messages) this.snapshotHistoryIds.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.snapshotHistoryIds.clear();
-		this.persistedHistoryIds.clear();
-		for (const { message } of loaded.messages) this.persistedHistoryIds.add(message.id);
-		this.notifyUpdate();
-	}
 	loadExternalState(state) {
 		this.stateSnapshot = state;
 		this.notifyUpdate();
@@ -669,20 +568,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 repositoryHeadId = this.repository.headId;
-			const parentId = (cached === void 0 || shouldEagerlyInsertAssistant && (repositoryHeadId === null || runStartMessageIds.has(repositoryHeadId))) && assistantParentId && this.hasMessage(assistantParentId) ? assistantParentId : repositoryHeadId;
+			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);
@@ -694,23 +593,40 @@ var AgUiThreadRuntimeCore = class {
 			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;
-				const reassigned = this.reassignAssistantId(placeholder, serverId);
-				const adoptsVisibleCollision = !reassigned && !runStartMessageIds.has(serverId) && this.repository.headId === serverId;
-				if (reassigned || adoptsVisibleCollision) {
-					assistantMessageId = serverId;
-					if (adoptsVisibleCollision) {
-						const parentId = this.tryGetMessage(serverId)?.parentId ?? null;
-						this.markPendingAssistantHistory(serverId, parentId);
-					}
-				} else assistantCollided = true;
-			}
+				adoptServerMessageId(serverId, false);
+			},
+			onTextMessageStart: (serverId) => adoptServerMessageId(serverId, true)
 		});
 		const dispatch = (event) => this.handleEvent(aggregator, event, assistantMessageId);
 		const abortController = new AbortController();
@@ -802,7 +718,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;
@@ -825,7 +741,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;
 			}
@@ -848,7 +764,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"
@@ -856,7 +772,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,
@@ -879,21 +795,6 @@ var AgUiThreadRuntimeCore = class {
 		this.pendingA2uiAction = void 0;
 		return input;
 	}
-	installResumeShim() {
-		const agent = this.agent;
-		if (agent[symbolResumeShim]) return;
-		agent[symbolResumeShim] = true;
-		const original = Object.hasOwn(agent, "prepareRunAgentInput") ? agent.prepareRunAgentInput : Object.getPrototypeOf(agent)?.prepareRunAgentInput;
-		if (typeof original !== "function") return;
-		agent.prepareRunAgentInput = function(params) {
-			const input = original.call(this, params);
-			if (params?.resume !== void 0 && input && typeof input === "object") return {
-				...input,
-				resume: params.resume
-			};
-			return input;
-		};
-	}
 	setRunning(running) {
 		this.isRunningFlag = running;
 		this.notifyUpdate();
@@ -921,31 +822,31 @@ 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) {
@@ -976,7 +877,7 @@ var AgUiThreadRuntimeCore = class {
 	}
 	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;
@@ -1001,14 +902,12 @@ var AgUiThreadRuntimeCore = class {
 	}
 	preserveToolResults(previous, next) {
 		const resolved = /* @__PURE__ */ new Map();
-		for (const part of previous) if (part.type === "tool-call" && isResolvedToolCall(part)) resolved.set(part.toolCallId, part);
+		for (const part of iterateToolCallParts(previous)) if (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" || isResolvedToolCall(part)) return part;
+		const { content: merged, changed } = mapToolCallPartsDeep(next, (part) => {
+			if (isResolvedToolCall(part)) return part;
 			const prior = resolved.get(part.toolCallId);
 			if (!prior) return part;
-			changed = true;
 			return {
 				...part,
 				result: prior.result,
@@ -1078,13 +977,13 @@ 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);
 			let matchedToolCall = false;
-			const content = assistant.content.map((part) => {
-				if (part.type !== "tool-call" || part.toolCallId !== event.toolCallId) return part;
+			const { content } = mapToolCallPartsDeep(assistant.content, (part) => {
+				if (part.toolCallId !== event.toolCallId) return part;
 				matchedToolCall = true;
 				if (part.mcp?.app !== void 0 && isPlainObject(part.result) && Array.isArray(part.result["content"])) return {
 					...part,
@@ -1123,12 +1022,12 @@ 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;
-			const content = assistant.content.map((part) => {
-				if (part.type !== "tool-call" || part.toolCallId !== toolCallId) return part;
+			const { content } = mapToolCallPartsDeep(assistant.content, (part) => {
+				if (part.toolCallId !== toolCallId) return part;
 				matchedToolCall = true;
 				return {
 					...part,
@@ -1157,7 +1056,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 = [];
@@ -1171,7 +1070,7 @@ var AgUiThreadRuntimeCore = class {
 			if (preservesActiveAssistant) converted.push(activeAssistant);
 			this.applyExternalMessages(converted);
 			if (activeAssistant !== void 0) {
-				const activeItem = this.tryGetMessage(activeAssistant.id);
+				const activeItem = this.session.tryGetMessage(activeAssistant.id);
 				if (activeItem) {
 					if (preservesActiveAssistant) this.snapshotHistoryIds.delete(activeAssistant.id);
 					this.markPendingAssistantHistory(activeAssistant.id, activeItem.parentId);
@@ -1205,7 +1104,7 @@ var AgUiThreadRuntimeCore = class {
 		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);