@assistant-ui/react-ag-ui 0.0.64 → 0.0.66

package/dist/runtime/adapter/run-aggregator.js

@@ -1,11 +1,11 @@
 "use client";
 import { readMcpAppResourceUri } from "../mcp-tool-result.js";
+import { AG_UI_METADATA_NAMESPACE } from "../types.js";
 import { projectAgUiToolApprovals } from "./tool-approval.js";
 import { isMcpAppUri } from "@assistant-ui/core";
 import jsonpatch from "fast-json-patch";
-import { applyA2uiOperations, convertSurfaceToUISpec, surfaceToOperations } from "@assistant-ui/react-generative-ui/a2ui";
+import { applyA2uiOperations, surfaceToPresentToolCall } from "@assistant-ui/generative-ui/a2ui";
 //#region src/runtime/adapter/run-aggregator.ts
-const AG_UI_METADATA_NAMESPACE = "agui";
 const ROOT_SCOPE = "";
 const MAX_SUBAGENT_DEPTH = 16;
 const createJSONContainerScanner = () => ({
@@ -62,8 +62,6 @@ const scanJSONContainerDelta = (scanner, delta) => {
 	}
 	return completed && scanner.state === "complete";
 };
-const MCP_APPS_ACTIVITY_TYPE = "mcp-apps";
-const A2UI_SURFACE_ACTIVITY_TYPE = "a2ui-surface";
 const isPlainObject = (value) => !!value && typeof value === "object" && !Array.isArray(value);
 /**
 * Collects AG-UI events into assistant-ui run snapshots that can be yielded from a ChatModelAdapter.
@@ -108,6 +106,7 @@ var RunAggregator = class {
 	activityParts = /* @__PURE__ */ new Map();
 	lastResolvedToolCallIdByScope = /* @__PURE__ */ new Map();
 	partOrder = [];
+	lastReasoningBeforeToolCall;
 	textPartCounter = 0;
 	serverMessageIdReported = false;
 	reportedServerMessageId;
@@ -228,7 +227,7 @@ var RunAggregator = class {
 			case "THINKING_TEXT_MESSAGE_START":
 			case "REASONING_START":
 			case "REASONING_MESSAGE_START":
-				this.handleReasoningStart(this.scopeOf("subagentRunId" in event ? event : {}), "messageId" in event ? event.messageId : void 0, event.type === "REASONING_MESSAGE_START");
+				this.handleReasoningStart(this.scopeOf("subagentRunId" in event ? event : {}), "messageId" in event ? event.messageId : void 0, event.type === "REASONING_MESSAGE_START", event.timestamp);
 				break;
 			case "REASONING_ENCRYPTED_VALUE":
 				if (event.subtype === "message") {
@@ -259,11 +258,11 @@ var RunAggregator = class {
 				break;
 			case "THINKING_TEXT_MESSAGE_END":
 			case "THINKING_END":
-				this.handleReasoningEnd(ROOT_SCOPE);
+				this.handleReasoningEnd(ROOT_SCOPE, event.timestamp);
 				break;
 			case "REASONING_MESSAGE_END":
 			case "REASONING_END":
-				this.handleReasoningEnd(this.scopeOf(event));
+				this.handleReasoningEnd(this.scopeOf(event), event.timestamp);
 				break;
 			case "TOOL_CALL_START": {
 				const scope = this.scopeOf(event);
@@ -441,19 +440,19 @@ var RunAggregator = class {
 		for (const bucket of this.a2uiBuckets.values()) for (const [surfaceId, surface] of bucket) surfaces.set(surfaceId, surface);
 		const activeToolCallIds = /* @__PURE__ */ new Set();
 		for (const [surfaceId, surface] of surfaces) {
-			const toolCallId = `a2ui:${surfaceId}`;
-			const { spec, warnings } = convertSurfaceToUISpec(surface);
+			const { toolCall, warnings } = surfaceToPresentToolCall(surfaceId, surface);
 			for (const warning of warnings) this.logger.debug("[agui] a2ui surface conversion warning", warning);
-			if (!spec) continue;
+			if (!toolCall) continue;
+			const { toolCallId } = toolCall;
 			activeToolCallIds.add(toolCallId);
 			const entry = {
-				toolCallId,
-				toolCallName: "present",
-				argsText: JSON.stringify(spec),
-				parsedArgs: spec,
-				result: {},
+				toolCallId: toolCall.toolCallId,
+				toolCallName: toolCall.toolName,
+				argsText: toolCall.argsText,
+				parsedArgs: toolCall.args,
+				result: toolCall.result,
 				isError: void 0,
-				artifact: { a2ui: surfaceToOperations(surface) },
+				artifact: toolCall.artifact,
 				snapshotResultApplied: false
 			};
 			if (!this.toolCalls.has(toolCallId)) this.partOrder.push({
@@ -508,6 +507,7 @@ var RunAggregator = class {
 		this.activityParts.clear();
 		this.lastResolvedToolCallIdByScope.clear();
 		this.partOrder.length = 0;
+		this.lastReasoningBeforeToolCall = void 0;
 		this.textPartCounter = 0;
 		this.activeTextMessageIdByScope.clear();
 		this.reportedServerMessageId = void 0;
@@ -523,8 +523,70 @@ var RunAggregator = class {
 	}
 	beginDistinctTextMessage(messageId) {
 		if (!messageId || this.lastTextMessageId === void 0 || this.lastTextMessageId === messageId || !this.onTextMessageStart) return;
+		const carried = this.takeTrailingReasoning();
+		if (carried.parts.length > 0 || carried.hiddenSignatures.length > 0) this.emit();
+		const reasoningPartCounter = this.reasoningPartCounter;
 		this.resetMessageParts();
 		this.onTextMessageStart(messageId);
+		this.reasoningPartCounter = reasoningPartCounter;
+		this.restoreReasoning(carried);
+	}
+	takeTrailingReasoning() {
+		const parts = [];
+		let last = this.partOrder.at(-1);
+		while (last?.kind === "reasoning" && last.subagentRunId === void 0 && last !== this.lastReasoningBeforeToolCall) {
+			this.partOrder.pop();
+			parts.unshift({
+				key: last.key,
+				entry: this.reasoningParts.get(last.key) ?? { buffer: "" },
+				signature: this.reasoningSignatures.get(last.key),
+				signatureId: this.reasoningSignatureIds.get(last.key),
+				reasoningId: this.reasoningMessageIds.get(last.key),
+				anonymous: this.anonymousReasoningKeys.has(last.key)
+			});
+			last = this.partOrder.at(-1);
+		}
+		const cut = this.partOrder.length;
+		const hiddenSignatures = Array.from(this.hiddenSignatures).filter(([id]) => this.hiddenSignatureAnchors.get(id) >= cut);
+		for (const [id] of hiddenSignatures) {
+			this.hiddenSignatures.delete(id);
+			this.hiddenSignatureAnchors.delete(id);
+		}
+		const hiddenBlockIds = Array.from(this.hiddenBlockAnchors).filter(([, anchor]) => anchor >= cut).map(([id]) => id);
+		const hiddenAnonymous = this.hiddenAnonymousAnchor !== void 0 && this.hiddenAnonymousAnchor >= cut;
+		const activeCarried = this.hiddenActiveReasoning === "identified" && hiddenBlockIds.length > 0 || this.hiddenActiveReasoning === "anonymous" && hiddenAnonymous;
+		return {
+			parts,
+			activeKey: this.activeReasoningKeyByScope.get(ROOT_SCOPE),
+			hiddenSignatures,
+			hiddenBlockIds,
+			hiddenAnonymous,
+			hiddenActiveReasoning: activeCarried ? this.hiddenActiveReasoning : "none"
+		};
+	}
+	restoreReasoning(carried) {
+		for (const part of carried.parts) {
+			this.partOrder.push({
+				kind: "reasoning",
+				key: part.key
+			});
+			this.reasoningParts.set(part.key, part.entry);
+			if (part.signature !== void 0) this.reasoningSignatures.set(part.key, part.signature);
+			if (part.signatureId !== void 0) this.reasoningSignatureIds.set(part.key, part.signatureId);
+			if (part.reasoningId !== void 0) this.reasoningMessageIds.set(part.key, part.reasoningId);
+			if (part.anonymous) this.anonymousReasoningKeys.add(part.key);
+		}
+		if (carried.activeKey !== void 0 && this.reasoningParts.has(carried.activeKey)) this.activeReasoningKeyByScope.set(ROOT_SCOPE, carried.activeKey);
+		for (const [id, encryptedValue] of carried.hiddenSignatures) {
+			this.hiddenSignatures.set(id, encryptedValue);
+			this.hiddenSignatureAnchors.set(id, 0);
+		}
+		for (const id of carried.hiddenBlockIds) {
+			this.hiddenReasoningIds.add(id);
+			this.hiddenBlockAnchors.set(id, 0);
+		}
+		if (carried.hiddenAnonymous) this.hiddenAnonymousAnchor = 0;
+		this.hiddenActiveReasoning = carried.hiddenActiveReasoning;
 	}
 	generateTextKey() {
 		this.textPartCounter += 1;
@@ -571,7 +633,10 @@ var RunAggregator = class {
 	startToolCall(scope, id, name, parentMessageId) {
 		if (!id) return;
 		this.activeTextMessageIdByScope.delete(scope);
-		if (!this.partOrder.some((part) => part.kind === "tool-call" && part.toolCallId === id)) this.insertToolPart(scope, id, parentMessageId);
+		if (!this.partOrder.some((part) => part.kind === "tool-call" && part.toolCallId === id)) {
+			if (scope === ROOT_SCOPE) this.lastReasoningBeforeToolCall = this.partOrder.findLast((part) => part.kind === "reasoning" && part.subagentRunId === void 0);
+			this.insertToolPart(scope, id, parentMessageId);
+		}
 		const state = {
 			toolCallId: id,
 			toolCallName: name ?? "tool",
@@ -773,7 +838,8 @@ var RunAggregator = class {
 			const rawScope = "subagentRunId" in part ? part.subagentRunId ?? ROOT_SCOPE : ROOT_SCOPE;
 			if (part.kind === "reasoning") {
 				if (this.showThinking) {
-					const buffer = this.reasoningParts.get(part.key) ?? "";
+					const entry = this.reasoningParts.get(part.key);
+					const buffer = entry?.buffer ?? "";
 					const isActive = this.activeReasoningKeyByScope.get(rawScope) === part.key;
 					if (buffer.length > 0 || isActive) {
 						const encryptedValue = this.reasoningSignatures.get(part.key);
@@ -785,6 +851,7 @@ var RunAggregator = class {
 						snapshot.push({
 							type: "reasoning",
 							text: buffer,
+							...entry?.timing ? { timing: entry.timing } : {},
 							...Object.keys(meta).length > 0 ? { providerMetadata: { [AG_UI_METADATA_NAMESPACE]: meta } } : {}
 						});
 						materialized();
@@ -905,7 +972,7 @@ var RunAggregator = class {
 			toolCallCount: this.toolCalls.size
 		};
 	}
-	handleReasoningStart(scope, messageId, isMessageId = false) {
+	handleReasoningStart(scope, messageId, isMessageId = false, timestamp) {
 		if (!this.showThinking) {
 			if (scope !== ROOT_SCOPE) return;
 			const anchor = this.partOrder.length;
@@ -924,7 +991,7 @@ var RunAggregator = class {
 		if (messageId === void 0) this.anonymousReasoningKeys.add(key);
 		else if (isMessageId) this.reasoningMessageIds.set(key, messageId);
 		if (!this.reasoningParts.has(key)) {
-			this.reasoningParts.set(key, "");
+			this.reasoningParts.set(key, { buffer: "" });
 			this.partOrder.push(scope === ROOT_SCOPE ? {
 				kind: "reasoning",
 				key
@@ -934,6 +1001,9 @@ var RunAggregator = class {
 				subagentRunId: scope
 			});
 		}
+		const entry = this.reasoningParts.get(key);
+		if (entry?.timing) entry.timing = { startedAt: entry.timing.startedAt };
+		else if (entry && typeof timestamp === "number" && Number.isFinite(timestamp)) entry.timing = { startedAt: timestamp };
 		this.activeReasoningKeyByScope.set(scope, key);
 		this.emit();
 	}
@@ -947,14 +1017,21 @@ var RunAggregator = class {
 		if (!this.activeReasoningKeyByScope.has(scope)) this.handleReasoningStart(scope, messageId, isMessageId);
 		const key = this.activeReasoningKeyByScope.get(scope);
 		if (!key) return;
-		this.reasoningParts.set(key, (this.reasoningParts.get(key) ?? "") + delta);
+		const entry = this.reasoningParts.get(key);
+		if (entry) entry.buffer += delta;
 		this.emit();
 	}
-	handleReasoningEnd(scope) {
+	handleReasoningEnd(scope, timestamp) {
 		if (!this.showThinking) {
 			if (scope === ROOT_SCOPE) this.hiddenActiveReasoning = "none";
 			return;
 		}
+		const key = this.activeReasoningKeyByScope.get(scope);
+		const entry = key ? this.reasoningParts.get(key) : void 0;
+		if (entry?.timing && typeof timestamp === "number" && Number.isFinite(timestamp)) entry.timing = {
+			...entry.timing,
+			completedAt: timestamp
+		};
 		this.activeReasoningKeyByScope.delete(scope);
 		this.emit();
 	}
@@ -968,6 +1045,6 @@ function tryParseJSON(value) {
 	}
 }
 //#endregion
-export { A2UI_SURFACE_ACTIVITY_TYPE, AG_UI_METADATA_NAMESPACE, MCP_APPS_ACTIVITY_TYPE, RunAggregator, isPlainObject, tryParseJSON };
+export { RunAggregator, isPlainObject, tryParseJSON };
 
 //# sourceMappingURL=run-aggregator.js.map