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

package/dist/runtime/AgUiThreadRuntimeCore.d.ts

@@ -32,8 +32,7 @@ declare class AgUiThreadRuntimeCore {
   private readonly notifyUpdate;
   private readonly reportedErrors;
   private runtime;
-  private readonly repository;
-  private exportedRepository;
+  private readonly session;
   private isRunningFlag;
   private abortController;
   private activeRunAgent;
@@ -47,6 +46,7 @@ declare class AgUiThreadRuntimeCore {
   private readonly historyWrites;
   private _isLoading;
   private _loadPromise;
+  private _loadRequested;
   private pendingResumeMessageId;
   private pendingA2uiResume;
   private pendingA2uiAction;
@@ -56,16 +56,6 @@ declare class AgUiThreadRuntimeCore {
   detachRuntime(): void;
   getMessages(): readonly ThreadMessage[];
   getMessageRepository(): ExportedMessageRepository;
-  private tryGetMessage;
-  private tryGetMessages;
-  private hasMessage;
-  private addOrUpdateMessage;
-  private deleteMessage;
-  private tryDeleteMessage;
-  private switchToBranch;
-  private resetRepositoryHead;
-  private clearRepository;
-  private updateMessage;
   getState(): ReadonlyJSONValue | undefined;
   isRunning(): boolean;
   get isLoading(): boolean;
@@ -121,7 +111,6 @@ declare class AgUiThreadRuntimeCore {
   private maybeCompleteAfterToolResults;
   private startResumeRun;
   applyExternalMessages(messages: readonly ThreadMessage[]): void;
-  private applyExternalMessageRepository;
   loadExternalState(state: ReadonlyJSONValue): void;
   setState(next: ReadonlyJSONValue | ((prev: ReadonlyJSONValue | undefined) => ReadonlyJSONValue)): void;
   resetState(): void;

package/dist/runtime/AgUiThreadRuntimeCore.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"AgUiThreadRuntimeCore.d.ts","names":[],"sources":["../../src/runtime/AgUiThreadRuntimeCore.ts"],"mappings":";;;;;;KA2EK,YAAY,YAAY;KACxB,gBACH,SAAS,wBACN,eAAe;KACf;EACH;EACA;EACA,WAAW;EACX,SAAS;;KAGN;EACH,OAAO;EACP,QAAQ;EACR;EACA;EACA,WAAW,OAAO;EAClB;EACA,UAAU;EACV;;cAuDW;UACH;UACA;UACA;UACA;UACA;UACA;mBACS;mBACA;UAET;mBACS;UACT;UACA;UACA;UAGA;UACA;UACA;UACA;UACA;mBACS;mBACA;mBACA;mBACA;UACT;UACA;UACA;UACA;UACA;EAEI,YAAA,SAAS;EAYrB,cAAc,SAAS,KAAK;EAW5B,cAAc,SAAS;EAIvB;EAOA,wBAAwB;EAIxB,wBAAwB;UAmBhB;UAQA;UAUA;UAIA;UAQA;UAQA;UAMA;UAKA;UAKA;UAKA;EAYR,YAAY;EAIZ;MAII;EAIJ,mBAAmB;EA6Cb,OAAO,SAAS,gBAAgB;UAW9B;UAQA;EAiBF,KAAK,SAAS,gBAAgB;EAI9B,OACJ,yBACA;IAAU,YAAY;MACrB;EAMG,UAAU;EAcV,OAAO,QAAQ,kBAAkB;EAUjC,kBAAkB,mBAAmB,kBAAkB;UAwBrD;UAOA;EAgBR;IACE;IACA,qBAAqB;;EAavB;IACE;IACA;;EAcI,yBACJ,oBAAoB,oBACnB;;;;;;;UAuCW;;;;;;;;EAed,YAAY;UAYJ;EAwBF,sBACJ,SAAS,+BACR;EAoEG,UACJ,SAAS,qBACT,qBAAqB,oBACpB;UAyCK;UAkBA;UAiCA;UA0BA;EAmCR,yBAAyB;UAkBjB;EAiBR,cAAc,SAAS;EAyBvB,eAAe,QAAQ;UA4Bf;UAYA;UA6BA;EAUR,sBAAsB,mBAAmB;UAkDjC;EAiER,kBAAkB,OAAO;EAKzB,SACE,MACI,sBACE,MAAM,kCAAkC;EAOhD;UAKc;UAoOA;UA2DN;UAkCA;UAqBA;UAKA;UAQA;UAyBA;UAkEA;UAiDA;UA4BA;UA6BA;UAgEA;UA2EA;UA8DA;UAuDA;UAQA;UAIA;UAKA;UAMA;UAQA;UAsDA;UAqBA"}
+{"version":3,"file":"AgUiThreadRuntimeCore.d.ts","names":[],"sources":["../../src/runtime/AgUiThreadRuntimeCore.ts"],"mappings":";;;;;;KA8EK,YAAY,YAAY;KACxB,gBACH,SAAS,wBACN,eAAe;KACf;EACH;EACA;EACA,WAAW;EACX,SAAS;;KAGN;EACH,OAAO;EACP,QAAQ;EACR;EACA;EACA,WAAW,OAAO;EAClB;EACA,UAAU;EACV;;cAoCW;UACH;UACA;UACA;UACA;UACA;UACA;mBACS;mBACA;UAET;mBACS;UAgBT;UACA;UAGA;UACA;UACA;UACA;UACA;mBACS;mBACA;mBACA;mBACA;UACT;UACA;UACA;UACA;UACA;UACA;EAEI,YAAA,SAAS;EAYrB,cAAc,SAAS,KAAK;EAsB5B,cAAc,SAAS;EAIvB;EAOA,wBAAwB;EAIxB,wBAAwB;EAIxB,YAAY;EAIZ;MAII;EAIJ,mBAAmB;EAsDb,OAAO,SAAS,gBAAgB;UAW9B;UAQA;EAiBF,KAAK,SAAS,gBAAgB;EAI9B,OACJ,yBACA;IAAU,YAAY;MACrB;EAMG,UAAU;EAcV,OAAO,QAAQ,kBAAkB;EAUjC,kBAAkB,mBAAmB,kBAAkB;UAwBrD;UAOA;EAgBR;IACE;IACA,qBAAqB;;EAavB;IACE;IACA;;EAcI,yBACJ,oBAAoB,oBACnB;;;;;;;UAuCW;;;;;;;;EAed,YAAY;UAYJ;EAwBF,sBACJ,SAAS,+BACR;EAsEG,UACJ,SAAS,qBACT,qBAAqB,oBACpB;UAyCK;UAkBA;UAiCA;UA0BA;EAmCR,yBAAyB;UAkBjB;EAiBR,cAAc,SAAS;EAyBvB,eAAe,QAAQ;UA4Bf;UAYA;UA6BA;EAUR,sBAAsB,mBAAmB;EAkDzC,kBAAkB,OAAO;EAKzB,SACE,MACI,sBACE,MAAM,kCAAkC;EAOhD;UAKc;UAyPA;UA2DN;UAkCA;UAqBA;UAKA;UAQA;UAyBA;UAkEA;UAiDA;UA4BA;UA6BA;UAgEA;UA2EA;UA8DA;UAuDA;UAQA;UAIA;UAKA;UAMA;UAQA;UAsDA;UAqBA"}

package/dist/runtime/AgUiThreadRuntimeCore.js

@@ -5,7 +5,7 @@ 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 } from "@assistant-ui/core/internal";
 import jsonpatch from "fast-json-patch";
 //#region src/runtime/AgUiThreadRuntimeCore.ts
 const optimisticPrefix = "__optimistic__";
@@ -17,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 = {
@@ -54,8 +43,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 +70,7 @@ var AgUiThreadRuntimeCore = class {
 	historyWrites = /* @__PURE__ */ new Map();
 	_isLoading = false;
 	_loadPromise;
+	_loadRequested = false;
 	pendingResumeMessageId = null;
 	pendingA2uiResume = false;
 	pendingA2uiAction;
@@ -90,8 +92,10 @@ 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 +107,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 +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);
 			}
@@ -222,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;
@@ -345,10 +292,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;
+		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.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 +306,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 +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",
@@ -425,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,
@@ -436,7 +383,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;
@@ -473,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) => {
@@ -491,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;
@@ -517,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;
@@ -541,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",
@@ -566,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;
@@ -575,85 +522,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 +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 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 +590,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 +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;
@@ -825,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;
 			}
@@ -848,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"
@@ -856,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,
@@ -921,31 +834,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 +889,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;
@@ -1078,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);
@@ -1123,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;
@@ -1157,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 = [];
@@ -1171,7 +1084,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 +1118,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);