@assistant-ui/react-ag-ui 0.0.60 → 0.0.61

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 { createMessageRepositorySession, invokeUserCallback, iterateToolCallParts, mapToolCallPartsDeep } from "@assistant-ui/core/internal";
+import { appendToolInteraction, createMessageRepositorySession, invokeUserCallback, iterateToolCallParts, mapToolCallPartsDeep } from "@assistant-ui/core/internal";
 import jsonpatch from "fast-json-patch";
 //#region src/runtime/AgUiThreadRuntimeCore.ts
 const optimisticPrefix = "__optimistic__";
@@ -162,6 +162,13 @@ var AgUiThreadRuntimeCore = class {
 		if (!startRun || !ownsThread) return;
 		await this.startRun(threadMessageId, message.runConfig);
 	}
+	appendVoiceTranscript(message) {
+		const parentId = this.session.headId;
+		this.session.addOrUpdateMessage(parentId, message);
+		this.session.switchToBranch(message.id);
+		this.notifyUpdate();
+		this.recordHistoryEntry(parentId, message);
+	}
 	maybeAutoCancelPendingToolCalls() {
 		if (this.autoCancelPendingToolCalls === false) return;
 		const pending = this.getPendingToolCalls();
@@ -484,6 +491,38 @@ var AgUiThreadRuntimeCore = class {
 		this.notifyUpdate();
 		this.maybeResumeAfterToolResults(sessionMessageId);
 	}
+	async recordToolInteraction(options) {
+		const sessionMessageId = this.session.tryGetMessage(options.messageId) ? options.messageId : this.findMessageIdForToolCall(options.toolCallId);
+		if (sessionMessageId === void 0) throw new Error(`[agui] recordToolInteraction: message "${options.messageId}" was not found`);
+		const item = this.session.tryGetMessage(sessionMessageId);
+		if (!item) throw new Error(`[agui] recordToolInteraction: message "${sessionMessageId}" was not found`);
+		let recorded = false;
+		const updated = this.session.updateMessage(sessionMessageId, (message) => {
+			if (message.role !== "assistant") return message;
+			const assistant = message;
+			const { content } = mapToolCallPartsDeep(assistant.content, (part) => {
+				if (part.toolCallId !== options.toolCallId) return part;
+				recorded = true;
+				return {
+					...part,
+					unstable_interactions: appendToolInteraction(part.unstable_interactions, options.interaction)
+				};
+			});
+			return recorded ? {
+				...assistant,
+				content
+			} : assistant;
+		});
+		if (!recorded || !updated) throw new Error(`[agui] recordToolInteraction: tool call "${options.toolCallId}" was not found on message "${options.messageId}"`);
+		this.notifyUpdate();
+		const message = this.session.tryGetMessage(sessionMessageId)?.message;
+		const history = this.history;
+		if (!message || !history?.update || !this.isPersistableStatus(message.status)) return;
+		await this.chainHistoryWrite(sessionMessageId, () => history.update({
+			parentId: item.parentId,
+			message
+		}));
+	}
 	sendA2uiAction(action) {
 		this.assertNoPendingInterrupts();
 		this.maybeAutoCancelPendingToolCalls();
@@ -543,6 +582,18 @@ var AgUiThreadRuntimeCore = class {
 		});
 	}
 	applyExternalMessages(messages) {
+		messages = messages.map((message) => {
+			if (message.role === "system" || message.metadata.modality !== void 0) return message;
+			const modality = this.session.tryGetMessage(message.id)?.message.metadata.modality;
+			if (modality === void 0) return message;
+			return {
+				...message,
+				metadata: {
+					...message.metadata,
+					modality
+				}
+			};
+		});
 		this.pendingA2uiResumeOwner = null;
 		this.pendingA2uiAction = void 0;
 		this.assistantHistoryParents.clear();
@@ -597,6 +648,26 @@ var AgUiThreadRuntimeCore = class {
 		this.stateSnapshot = void 0;
 		this.notifyUpdate();
 	}
+	resetThreadState() {
+		const controller = this.abortController;
+		const activeRunAgent = this.activeRunAgent;
+		this.stateSnapshot = void 0;
+		this.pendingResume = null;
+		this.pendingA2uiResumeOwner = null;
+		this.pendingA2uiAction = void 0;
+		if (controller) {
+			this.abortController = null;
+			this.activeRunAgent = null;
+			this.setRunning(false);
+			try {
+				(activeRunAgent ?? this.agent).abortRun();
+			} catch (error) {
+				this.logger.error?.("[agui] agent abortRun failed", error);
+			} finally {
+				controller.abort();
+			}
+		} else this.notifyUpdate();
+	}
 	async startRun(parentId, runConfig, resume, resumeStream) {
 		if (!this.supersedeActiveRun()) return;
 		const normalizedRunConfig = runConfig ?? {};
@@ -652,6 +723,8 @@ var AgUiThreadRuntimeCore = class {
 				}
 			} else assistantCollided = true;
 		};
+		const abortController = new AbortController();
+		const abortSignal = abortController.signal;
 		const aggregator = new RunAggregator({
 			showThinking: this.showThinking,
 			logger: this.logger,
@@ -661,9 +734,11 @@ var AgUiThreadRuntimeCore = class {
 			},
 			onTextMessageStart: (serverId) => adoptServerMessageId(serverId, true)
 		});
-		const dispatch = (event) => this.handleEvent(aggregator, event, assistantMessageId);
-		const abortController = new AbortController();
-		const abortSignal = abortController.signal;
+		const dispatch = (event) => {
+			if (this.abortController !== abortController) return;
+			const nextAssistantMessageId = this.handleEvent(aggregator, event, assistantMessageId);
+			if (event.type === "MESSAGES_SNAPSHOT") assistantMessageId = nextAssistantMessageId;
+		};
 		this.abortController = abortController;
 		const runAgentInstance = this.agent;
 		this.activeRunAgent = runAgentInstance;
@@ -948,22 +1023,36 @@ var AgUiThreadRuntimeCore = class {
 		return resolvedMessageId;
 	}
 	preserveToolResults(previous, next) {
-		const resolved = /* @__PURE__ */ new Map();
-		for (const part of iterateToolCallParts(previous)) if (isResolvedToolCall(part)) resolved.set(part.toolCallId, part);
-		if (resolved.size === 0) return next;
+		const preserved = /* @__PURE__ */ new Map();
+		for (const part of iterateToolCallParts(previous)) if (isResolvedToolCall(part)) preserved.set(part.toolCallId, part);
+		if (preserved.size === 0) return this.preserveToolInteractions(previous, next);
 		const { content: merged, changed } = mapToolCallPartsDeep(next, (part) => {
-			if (isResolvedToolCall(part)) return part;
-			const prior = resolved.get(part.toolCallId);
+			const prior = preserved.get(part.toolCallId);
 			if (!prior) return part;
-			return {
+			if (!isResolvedToolCall(part) && isResolvedToolCall(prior)) return {
 				...part,
 				result: prior.result,
 				...prior.artifact !== void 0 ? { artifact: prior.artifact } : {},
 				...prior.isError !== void 0 ? { isError: prior.isError } : {},
 				...prior.modelContent !== void 0 ? { modelContent: prior.modelContent } : {}
 			};
+			return part;
 		});
-		return changed ? merged : next;
+		return this.preserveToolInteractions(previous, changed ? merged : next);
+	}
+	preserveToolInteractions(previous, next) {
+		const interactions = /* @__PURE__ */ new Map();
+		for (const part of iterateToolCallParts(previous)) if (part.unstable_interactions !== void 0) interactions.set(part.toolCallId, part.unstable_interactions);
+		if (interactions.size === 0) return next;
+		const { content, changed } = mapToolCallPartsDeep(next, (part) => {
+			const interaction = interactions.get(part.toolCallId);
+			if (interaction === void 0 || part.unstable_interactions !== void 0) return part;
+			return {
+				...part,
+				unstable_interactions: interaction
+			};
+		});
+		return changed ? content : next;
 	}
 	mergeAssistantMetadata(current, incoming) {
 		const annotations = incoming.unstable_annotations ? [...current.unstable_annotations, ...incoming.unstable_annotations] : current.unstable_annotations;
@@ -983,46 +1072,69 @@ var AgUiThreadRuntimeCore = class {
 		switch (event.type) {
 			case "STATE_SNAPSHOT":
 				this.stateSnapshot = event.snapshot;
-				this.notifyUpdate();
-				return;
+				this.updateActiveAssistantState(activeAssistantId);
+				break;
 			case "STATE_DELTA":
-				if (event.delta.length === 0) return;
+				if (event.delta.length === 0) break;
 				try {
 					const state = this.stateSnapshot ?? {};
 					const result = jsonpatch.applyPatch(state, event.delta, true, false);
 					this.stateSnapshot = result.newDocument;
-					this.notifyUpdate();
+					this.updateActiveAssistantState(activeAssistantId);
 				} catch (error) {
 					this.logger.error?.("[agui] failed to apply state delta", error);
 				}
-				return;
-			case "MESSAGES_SNAPSHOT":
+				break;
+			case "MESSAGES_SNAPSHOT": {
+				const previousHeadId = this.session.headId;
+				const activeAssistant = activeAssistantId === void 0 ? void 0 : this.session.tryGetMessage(activeAssistantId)?.message;
+				const hasActiveText = activeAssistant?.role === "assistant" && activeAssistant.content.some((part) => part.type === "text");
 				this.importMessagesSnapshot(event.messages, activeAssistantId);
-				return;
+				const headId = this.session.headId;
+				if (hasActiveText && headId !== previousHeadId) activeAssistantId = (headId === null ? void 0 : this.session.tryGetMessage(headId)?.message)?.role === "assistant" ? headId ?? void 0 : void 0;
+				this.updateActiveAssistantState(activeAssistantId);
+				break;
+			}
 			case "TOOL_CALL_RESULT":
 				if (!aggregator.hasToolCall(event.toolCallId)) {
 					const messageId = this.findMessageIdForToolCall(event.toolCallId);
 					if (messageId !== void 0) {
 						this.applyCrossRunToolResult(messageId, event);
-						return;
+						break;
 					}
 				}
 				aggregator.handle(event);
-				return;
+				break;
 			case "ACTIVITY_SNAPSHOT": {
 				const toolCallId = event.content["toolCallId"];
 				if (event.activityType === "mcp-apps" && typeof toolCallId === "string" && !aggregator.hasToolCall(toolCallId)) {
 					const messageId = this.findMessageIdForToolCall(toolCallId);
 					if (messageId !== void 0) {
 						this.applyCrossRunActivitySnapshot(messageId, toolCallId, event);
-						return;
+						break;
 					}
 				}
 				aggregator.handle(event);
-				return;
+				break;
 			}
 			default: aggregator.handle(event);
 		}
+		return activeAssistantId;
+	}
+	updateActiveAssistantState(messageId) {
+		const activeAssistantId = (messageId === void 0 ? void 0 : this.session.tryGetMessage(messageId)?.message)?.role === "assistant" && this.session.headId === messageId ? messageId : this.session.headId;
+		if (activeAssistantId !== null && activeAssistantId !== void 0) this.session.updateMessage(activeAssistantId, (message) => {
+			if (message.role !== "assistant") return message;
+			const assistant = message;
+			return {
+				...assistant,
+				metadata: {
+					...assistant.metadata,
+					unstable_state: this.stateSnapshot ?? null
+				}
+			};
+		});
+		this.notifyUpdate();
 	}
 	applyCrossRunToolResult(messageId, event) {
 		if (!this.session.updateMessage(messageId, (message) => {
@@ -1109,7 +1221,13 @@ var AgUiThreadRuntimeCore = class {
 			const normalized = fromAgUiMessages(rawMessages, { showThinking: this.showThinking });
 			const converted = [];
 			for (const message of normalized) try {
-				converted.push(fromThreadMessageLike(message, generateId(), FALLBACK_USER_STATUS));
+				const convertedMessage = fromThreadMessageLike(message, generateId(), FALLBACK_USER_STATUS);
+				const existing = this.session.tryGetMessage(convertedMessage.id)?.message;
+				if (convertedMessage.role === "assistant" && existing?.role === "assistant") converted.push({
+					...convertedMessage,
+					content: this.preserveToolInteractions(existing.content, convertedMessage.content)
+				});
+				else converted.push(convertedMessage);
 			} catch (error) {
 				this.logger.error?.("[agui] failed to import message from snapshot", error);
 			}