@assistant-ui/react-ag-ui 0.0.30 → 0.0.32

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

@@ -1,367 +1,370 @@
 "use client";
-export const AG_UI_METADATA_NAMESPACE = "agui";
+//#region src/runtime/adapter/run-aggregator.ts
+const AG_UI_METADATA_NAMESPACE = "agui";
 /**
- * Collects AG-UI events into assistant-ui run snapshots that can be yielded from a ChatModelAdapter.
- *
- * The aggregator keeps a single assistant message worth of parts. Each incoming event updates the parts and
- * emits a fresh snapshot through the provided `emit` callback.
- */
-export class RunAggregator {
-    emitUpdate;
-    showThinking;
-    logger;
-    onServerMessageId;
-    status;
-    interrupts;
-    textParts = new Map();
-    activeTextMessageId;
-    reasoningBuffer = "";
-    reasoningActive = false;
-    toolCalls = new Map();
-    partOrder = [];
-    hasReasoningPart = false;
-    textPartCounter = 0;
-    serverMessageIdReported = false;
-    constructor(options) {
-        this.emitUpdate = options.emit;
-        this.showThinking = options.showThinking;
-        this.logger = options.logger;
-        this.onServerMessageId = options.onServerMessageId;
-    }
-    handle(event) {
-        switch (event.type) {
-            case "RUN_STARTED": {
-                this.clearTextParts();
-                this.reasoningBuffer = "";
-                this.reasoningActive = false;
-                this.toolCalls.clear();
-                this.partOrder.length = 0;
-                this.hasReasoningPart = false;
-                this.textPartCounter = 0;
-                this.activeTextMessageId = undefined;
-                this.interrupts = undefined;
-                this.serverMessageIdReported = false;
-                this.status = { type: "running" };
-                this.emit();
-                break;
-            }
-            case "RUN_FINISHED": {
-                if (event.outcome?.type === "interrupt") {
-                    this.interrupts = event.outcome.interrupts;
-                    this.status = { type: "requires-action", reason: "interrupt" };
-                    this.emit();
-                    break;
-                }
-                this.interrupts = undefined;
-                const hasUnresolvedToolCalls = Array.from(this.toolCalls.values()).some((tc) => tc.result === undefined);
-                this.status =
-                    event.outcome?.type === "success" || !hasUnresolvedToolCalls
-                        ? { type: "complete", reason: "unknown" }
-                        : { type: "requires-action", reason: "tool-calls" };
-                this.emit();
-                break;
-            }
-            case "RUN_ERROR": {
-                this.status = {
-                    type: "incomplete",
-                    reason: "error",
-                    ...(event.message !== undefined ? { error: event.message } : {}),
-                };
-                this.emit();
-                break;
-            }
-            case "RUN_CANCELLED": {
-                this.status = { type: "incomplete", reason: "cancelled" };
-                this.emit();
-                break;
-            }
-            case "TEXT_MESSAGE_START": {
-                this.reportServerMessageId(event.messageId);
-                const id = this.startTextMessage(event.messageId);
-                if (id) {
-                    this.markTextPartTouched(id);
-                }
-                this.emit();
-                break;
-            }
-            case "TEXT_MESSAGE_CONTENT":
-            case "TEXT_MESSAGE_CHUNK": {
-                const incomingId = "messageId" in event ? event.messageId : undefined;
-                this.reportServerMessageId(incomingId);
-                if (!event.delta)
-                    break;
-                const id = this.resolveTextMessageId(incomingId);
-                this.appendText(id, event.delta);
-                this.emit();
-                break;
-            }
-            case "TEXT_MESSAGE_END": {
-                this.reportServerMessageId(event.messageId);
-                if (event.messageId && this.activeTextMessageId === event.messageId) {
-                    this.activeTextMessageId = undefined;
-                }
-                this.emit();
-                break;
-            }
-            case "THINKING_START":
-            case "THINKING_TEXT_MESSAGE_START":
-            case "REASONING_START":
-            case "REASONING_MESSAGE_START":
-                this.handleReasoningStart();
-                break;
-            case "THINKING_TEXT_MESSAGE_CONTENT":
-            case "REASONING_MESSAGE_CONTENT":
-                this.handleReasoningContent(event.delta);
-                break;
-            case "THINKING_TEXT_MESSAGE_END":
-            case "THINKING_END":
-            case "REASONING_MESSAGE_END":
-            case "REASONING_END":
-                this.handleReasoningEnd();
-                break;
-            case "TOOL_CALL_START": {
-                this.reportServerMessageId(event.parentMessageId);
-                this.startToolCall(event.toolCallId, event.toolCallName, event.parentMessageId);
-                this.emit();
-                break;
-            }
-            case "TOOL_CALL_ARGS":
-            case "TOOL_CALL_CHUNK": {
-                if (event.type === "TOOL_CALL_CHUNK") {
-                    this.reportServerMessageId(event.parentMessageId);
-                }
-                if (!event.delta)
-                    break;
-                this.appendToolArgs(event.toolCallId, event.delta);
-                this.emit();
-                break;
-            }
-            case "TOOL_CALL_END": {
-                this.emit();
-                break;
-            }
-            case "TOOL_CALL_RESULT": {
-                this.finishToolCall(event.toolCallId, event.content ?? "", event.role === "tool" ? false : undefined, event.messageId);
-                this.emit();
-                break;
-            }
-            default: {
-                this.logger.debug?.("[agui] aggregator ignored event", event);
-            }
-        }
-    }
-    reportServerMessageId(messageId) {
-        if (this.serverMessageIdReported || !messageId)
-            return;
-        this.serverMessageIdReported = true;
-        this.onServerMessageId?.(messageId);
-    }
-    clearTextParts() {
-        this.textParts.clear();
-    }
-    generateTextKey() {
-        this.textPartCounter += 1;
-        return `text-${this.textPartCounter}`;
-    }
-    startTextMessage(messageId) {
-        const id = messageId ?? this.generateTextKey();
-        this.ensureTextPart(id);
-        this.activeTextMessageId = id;
-        return id;
-    }
-    resolveTextMessageId(messageId) {
-        if (messageId) {
-            this.ensureTextPart(messageId);
-            this.activeTextMessageId = messageId;
-            return messageId;
-        }
-        if (this.activeTextMessageId) {
-            return this.activeTextMessageId;
-        }
-        const generated = this.generateTextKey();
-        this.ensureTextPart(generated);
-        this.activeTextMessageId = generated;
-        return generated;
-    }
-    ensureTextPart(id) {
-        if (!this.textParts.has(id)) {
-            this.textParts.set(id, { buffer: "", touched: false });
-            if (!this.partOrder.some((part) => part.kind === "text" && part.key === id)) {
-                this.partOrder.push({ kind: "text", key: id });
-            }
-        }
-    }
-    markTextPartTouched(id) {
-        const entry = this.textParts.get(id);
-        if (!entry)
-            return;
-        entry.touched = true;
-    }
-    appendText(id, delta) {
-        this.ensureTextPart(id);
-        const entry = this.textParts.get(id);
-        if (!entry)
-            return;
-        entry.buffer += delta;
-        entry.touched = true;
-    }
-    startToolCall(id, name, parentMessageId) {
-        if (!id)
-            return;
-        if (!this.partOrder.some((part) => part.kind === "tool-call" && part.toolCallId === id)) {
-            this.partOrder.push({ kind: "tool-call", toolCallId: id });
-        }
-        const state = {
-            toolCallId: id,
-            toolCallName: name ?? "tool",
-            argsText: "",
-            parsedArgs: undefined,
-            result: undefined,
-            isError: undefined,
-        };
-        if (parentMessageId) {
-            state.parentMessageId = parentMessageId;
-        }
-        this.toolCalls.set(id, state);
-    }
-    appendToolArgs(id, delta) {
-        const entry = id ? this.toolCalls.get(id) : undefined;
-        if (!entry)
-            return;
-        entry.argsText += delta;
-        try {
-            const parsed = JSON.parse(entry.argsText);
-            if (parsed && typeof parsed === "object") {
-                entry.parsedArgs = parsed;
-            }
-            else {
-                entry.parsedArgs = undefined;
-            }
-        }
-        catch {
-            entry.parsedArgs = undefined;
-        }
-    }
-    finishToolCall(id, content, isError, toolMessageId) {
-        if (!id)
-            return;
-        let entry = this.toolCalls.get(id);
-        if (!entry) {
-            entry = {
-                toolCallId: id,
-                toolCallName: "tool",
-                argsText: "",
-                parsedArgs: undefined,
-                result: undefined,
-                isError: undefined,
-            };
-            this.toolCalls.set(id, entry);
-        }
-        if (!this.partOrder.some((part) => part.kind === "tool-call" && part.toolCallId === id)) {
-            this.partOrder.push({ kind: "tool-call", toolCallId: id });
-        }
-        entry.result = this.tryParseJSON(content);
-        entry.isError = isError;
-        if (toolMessageId) {
-            entry.toolMessageId = toolMessageId;
-        }
-    }
-    tryParseJSON(value) {
-        if (!value)
-            return value;
-        try {
-            return JSON.parse(value);
-        }
-        catch {
-            return value;
-        }
-    }
-    emit() {
-        const snapshot = [];
-        for (const part of this.partOrder) {
-            if (part.kind === "reasoning") {
-                if (this.showThinking &&
-                    (this.reasoningActive || this.reasoningBuffer.length > 0)) {
-                    snapshot.push({
-                        type: "reasoning",
-                        text: this.reasoningBuffer,
-                    });
-                }
-                continue;
-            }
-            if (part.kind === "text") {
-                const entry = this.textParts.get(part.key);
-                if (entry?.touched) {
-                    snapshot.push({ type: "text", text: entry.buffer });
-                }
-                continue;
-            }
-            const entry = this.toolCalls.get(part.toolCallId);
-            if (!entry)
-                continue;
-            const toolPart = {
-                type: "tool-call",
-                toolCallId: entry.toolCallId,
-                toolName: entry.toolCallName,
-                args: (entry.parsedArgs ?? {}),
-                argsText: entry.argsText,
-                ...(entry.result !== undefined ? { result: entry.result } : {}),
-                ...(entry.isError !== undefined ? { isError: entry.isError } : {}),
-                ...(entry.parentMessageId ? { parentId: entry.parentMessageId } : {}),
-                ...(entry.toolMessageId
-                    ? { unstable_toolMessageId: entry.toolMessageId }
-                    : {}),
-            };
-            snapshot.push(toolPart);
-        }
-        const result = {
-            content: snapshot,
-            ...(this.status ? { status: this.status } : undefined),
-            ...(this.interrupts
-                ? {
-                    metadata: {
-                        custom: {
-                            [AG_UI_METADATA_NAMESPACE]: {
-                                interrupts: this.interrupts,
-                            },
-                        },
-                    },
-                }
-                : undefined),
-        };
-        this.emitUpdate(result);
-    }
-    handleReasoningStart() {
-        if (!this.showThinking)
-            return;
-        this.reasoningActive = true;
-        this.ensureReasoningPart();
-        this.emit();
-    }
-    handleReasoningContent(delta) {
-        if (!this.showThinking || !delta)
-            return;
-        this.reasoningBuffer += delta;
-        this.ensureReasoningPart();
-        this.emit();
-    }
-    handleReasoningEnd() {
-        if (!this.showThinking)
-            return;
-        this.emit();
-    }
-    ensureReasoningPart() {
-        if (this.hasReasoningPart)
-            return;
-        // ensure reasoning appears before the first text segment if possible
-        const textIndex = this.partOrder.findIndex((part) => part.kind === "text");
-        if (textIndex === -1) {
-            this.partOrder.push({ kind: "reasoning" });
-        }
-        else {
-            this.partOrder.splice(textIndex, 0, { kind: "reasoning" });
-        }
-        this.hasReasoningPart = true;
-    }
-}
+* Collects AG-UI events into assistant-ui run snapshots that can be yielded from a ChatModelAdapter.
+*
+* The aggregator keeps a single assistant message worth of parts. Each incoming event updates the parts and
+* emits a fresh snapshot through the provided `emit` callback.
+*/
+var RunAggregator = class {
+	emitUpdate;
+	showThinking;
+	logger;
+	onServerMessageId;
+	status;
+	interrupts;
+	textParts = /* @__PURE__ */ new Map();
+	activeTextMessageId;
+	reasoningParts = /* @__PURE__ */ new Map();
+	activeReasoningKey;
+	reasoningPartCounter = 0;
+	toolCalls = /* @__PURE__ */ new Map();
+	partOrder = [];
+	textPartCounter = 0;
+	serverMessageIdReported = false;
+	streamStartTime;
+	firstTokenTime;
+	totalChunks = 0;
+	constructor(options) {
+		this.emitUpdate = options.emit;
+		this.showThinking = options.showThinking;
+		this.logger = options.logger;
+		this.onServerMessageId = options.onServerMessageId;
+	}
+	handle(event) {
+		switch (event.type) {
+			case "RUN_STARTED":
+				this.clearTextParts();
+				this.reasoningParts.clear();
+				this.activeReasoningKey = void 0;
+				this.reasoningPartCounter = 0;
+				this.toolCalls.clear();
+				this.partOrder.length = 0;
+				this.textPartCounter = 0;
+				this.activeTextMessageId = void 0;
+				this.interrupts = void 0;
+				this.serverMessageIdReported = false;
+				this.streamStartTime = Date.now();
+				this.firstTokenTime = void 0;
+				this.totalChunks = 0;
+				this.status = { type: "running" };
+				this.emit();
+				break;
+			case "RUN_FINISHED": {
+				if (event.outcome?.type === "interrupt") {
+					this.interrupts = event.outcome.interrupts;
+					this.status = {
+						type: "requires-action",
+						reason: "interrupt"
+					};
+					this.emit();
+					break;
+				}
+				this.interrupts = void 0;
+				const hasUnresolvedToolCalls = Array.from(this.toolCalls.values()).some((tc) => tc.result === void 0);
+				this.status = event.outcome?.type === "success" || !hasUnresolvedToolCalls ? {
+					type: "complete",
+					reason: "unknown"
+				} : {
+					type: "requires-action",
+					reason: "tool-calls"
+				};
+				this.emit();
+				break;
+			}
+			case "RUN_ERROR":
+				this.status = {
+					type: "incomplete",
+					reason: "error",
+					...event.message !== void 0 ? { error: event.message } : {}
+				};
+				this.emit();
+				break;
+			case "RUN_CANCELLED":
+				this.status = {
+					type: "incomplete",
+					reason: "cancelled"
+				};
+				this.emit();
+				break;
+			case "TEXT_MESSAGE_START": {
+				this.reportServerMessageId(event.messageId);
+				const id = this.startTextMessage(event.messageId);
+				if (id) this.markTextPartTouched(id);
+				this.emit();
+				break;
+			}
+			case "TEXT_MESSAGE_CONTENT":
+			case "TEXT_MESSAGE_CHUNK": {
+				const incomingId = "messageId" in event ? event.messageId : void 0;
+				this.reportServerMessageId(incomingId);
+				if (!event.delta) break;
+				this.recordFirstToken();
+				const id = this.resolveTextMessageId(incomingId);
+				this.appendText(id, event.delta);
+				this.totalChunks++;
+				this.emit();
+				break;
+			}
+			case "TEXT_MESSAGE_END":
+				this.reportServerMessageId(event.messageId);
+				if (event.messageId && this.activeTextMessageId === event.messageId) this.activeTextMessageId = void 0;
+				this.emit();
+				break;
+			case "THINKING_START":
+			case "THINKING_TEXT_MESSAGE_START":
+			case "REASONING_START":
+			case "REASONING_MESSAGE_START":
+				this.handleReasoningStart("messageId" in event ? event.messageId : void 0);
+				break;
+			case "THINKING_TEXT_MESSAGE_CONTENT":
+				this.handleReasoningContent(event.delta);
+				this.totalChunks++;
+				this.recordFirstToken();
+				break;
+			case "REASONING_MESSAGE_CONTENT":
+				this.handleReasoningContent(event.delta, "messageId" in event ? event.messageId : void 0);
+				this.totalChunks++;
+				this.recordFirstToken();
+				break;
+			case "THINKING_TEXT_MESSAGE_END":
+			case "THINKING_END":
+			case "REASONING_MESSAGE_END":
+			case "REASONING_END":
+				this.handleReasoningEnd();
+				break;
+			case "TOOL_CALL_START":
+				this.reportServerMessageId(event.parentMessageId);
+				this.startToolCall(event.toolCallId, event.toolCallName, event.parentMessageId);
+				this.emit();
+				break;
+			case "TOOL_CALL_ARGS":
+			case "TOOL_CALL_CHUNK":
+				if (event.type === "TOOL_CALL_CHUNK") this.reportServerMessageId(event.parentMessageId);
+				if (!event.delta) break;
+				this.appendToolArgs(event.toolCallId, event.delta);
+				this.emit();
+				break;
+			case "TOOL_CALL_END":
+				this.emit();
+				break;
+			case "TOOL_CALL_RESULT":
+				this.finishToolCall(event.toolCallId, event.content ?? "", event.role === "tool" ? false : void 0, event.messageId);
+				this.emit();
+				break;
+			default: this.logger.debug?.("[agui] aggregator ignored event", event);
+		}
+	}
+	reportServerMessageId(messageId) {
+		if (this.serverMessageIdReported || !messageId) return;
+		this.serverMessageIdReported = true;
+		this.onServerMessageId?.(messageId);
+	}
+	clearTextParts() {
+		this.textParts.clear();
+	}
+	generateTextKey() {
+		this.textPartCounter += 1;
+		return `text-${this.textPartCounter}`;
+	}
+	startTextMessage(messageId) {
+		const id = messageId ?? this.generateTextKey();
+		this.ensureTextPart(id);
+		this.activeTextMessageId = id;
+		return id;
+	}
+	resolveTextMessageId(messageId) {
+		if (messageId) {
+			this.ensureTextPart(messageId);
+			this.activeTextMessageId = messageId;
+			return messageId;
+		}
+		if (this.activeTextMessageId) return this.activeTextMessageId;
+		const generated = this.generateTextKey();
+		this.ensureTextPart(generated);
+		this.activeTextMessageId = generated;
+		return generated;
+	}
+	ensureTextPart(id) {
+		if (!this.textParts.has(id)) {
+			this.textParts.set(id, {
+				buffer: "",
+				touched: false
+			});
+			if (!this.partOrder.some((part) => part.kind === "text" && part.key === id)) this.partOrder.push({
+				kind: "text",
+				key: id
+			});
+		}
+	}
+	markTextPartTouched(id) {
+		const entry = this.textParts.get(id);
+		if (!entry) return;
+		entry.touched = true;
+	}
+	appendText(id, delta) {
+		this.ensureTextPart(id);
+		const entry = this.textParts.get(id);
+		if (!entry) return;
+		entry.buffer += delta;
+		entry.touched = true;
+	}
+	startToolCall(id, name, parentMessageId) {
+		if (!id) return;
+		this.activeTextMessageId = void 0;
+		if (!this.partOrder.some((part) => part.kind === "tool-call" && part.toolCallId === id)) this.partOrder.push({
+			kind: "tool-call",
+			toolCallId: id
+		});
+		const state = {
+			toolCallId: id,
+			toolCallName: name ?? "tool",
+			argsText: "",
+			parsedArgs: void 0,
+			result: void 0,
+			isError: void 0
+		};
+		if (parentMessageId) state.parentMessageId = parentMessageId;
+		this.toolCalls.set(id, state);
+	}
+	appendToolArgs(id, delta) {
+		const entry = id ? this.toolCalls.get(id) : void 0;
+		if (!entry) return;
+		entry.argsText += delta;
+		try {
+			const parsed = JSON.parse(entry.argsText);
+			if (parsed && typeof parsed === "object") entry.parsedArgs = parsed;
+			else entry.parsedArgs = void 0;
+		} catch {
+			entry.parsedArgs = void 0;
+		}
+	}
+	finishToolCall(id, content, isError, toolMessageId) {
+		if (!id) return;
+		let entry = this.toolCalls.get(id);
+		if (!entry) {
+			entry = {
+				toolCallId: id,
+				toolCallName: "tool",
+				argsText: "",
+				parsedArgs: void 0,
+				result: void 0,
+				isError: void 0
+			};
+			this.toolCalls.set(id, entry);
+		}
+		if (!this.partOrder.some((part) => part.kind === "tool-call" && part.toolCallId === id)) this.partOrder.push({
+			kind: "tool-call",
+			toolCallId: id
+		});
+		entry.result = this.tryParseJSON(content);
+		entry.isError = isError;
+		if (toolMessageId) entry.toolMessageId = toolMessageId;
+	}
+	tryParseJSON(value) {
+		if (!value) return value;
+		try {
+			return JSON.parse(value);
+		} catch {
+			return value;
+		}
+	}
+	emit() {
+		const snapshot = [];
+		for (const part of this.partOrder) {
+			if (part.kind === "reasoning") {
+				if (this.showThinking) {
+					const buffer = this.reasoningParts.get(part.key) ?? "";
+					if (buffer.length > 0 || this.activeReasoningKey === part.key) snapshot.push({
+						type: "reasoning",
+						text: buffer
+					});
+				}
+				continue;
+			}
+			if (part.kind === "text") {
+				const entry = this.textParts.get(part.key);
+				if (entry?.touched) snapshot.push({
+					type: "text",
+					text: entry.buffer
+				});
+				continue;
+			}
+			const entry = this.toolCalls.get(part.toolCallId);
+			if (!entry) continue;
+			const toolPart = {
+				type: "tool-call",
+				toolCallId: entry.toolCallId,
+				toolName: entry.toolCallName,
+				args: entry.parsedArgs ?? {},
+				argsText: entry.argsText,
+				...entry.result !== void 0 ? { result: entry.result } : {},
+				...entry.isError !== void 0 ? { isError: entry.isError } : {},
+				...entry.parentMessageId ? { parentId: entry.parentMessageId } : {},
+				...entry.toolMessageId ? { unstable_toolMessageId: entry.toolMessageId } : {}
+			};
+			snapshot.push(toolPart);
+		}
+		const timing = this.getTiming();
+		const metadata = {
+			...timing ? { timing } : {},
+			...this.interrupts ? { custom: { [AG_UI_METADATA_NAMESPACE]: { interrupts: this.interrupts } } } : {}
+		};
+		const result = {
+			content: snapshot,
+			...this.status ? { status: this.status } : void 0,
+			...Object.keys(metadata).length > 0 ? { metadata } : void 0
+		};
+		this.emitUpdate(result);
+	}
+	recordFirstToken() {
+		if (this.firstTokenTime === void 0 && this.streamStartTime !== void 0) this.firstTokenTime = Date.now() - this.streamStartTime;
+	}
+	getTiming() {
+		if (this.streamStartTime === void 0) return void 0;
+		const totalStreamTime = Date.now() - this.streamStartTime;
+		const tokenCount = this.totalChunks > 0 ? Math.ceil(Array.from(this.textParts.values()).reduce((sum, p) => sum + p.buffer.length, 0) / 4) : void 0;
+		const tokensPerSecond = tokenCount && totalStreamTime > 0 ? tokenCount / totalStreamTime * 1e3 : void 0;
+		return {
+			streamStartTime: this.streamStartTime,
+			...this.firstTokenTime !== void 0 ? { firstTokenTime: this.firstTokenTime } : {},
+			totalStreamTime,
+			...tokenCount !== void 0 ? { tokenCount } : {},
+			...tokensPerSecond !== void 0 ? { tokensPerSecond } : {},
+			totalChunks: this.totalChunks,
+			toolCallCount: this.toolCalls.size
+		};
+	}
+	handleReasoningStart(messageId) {
+		if (!this.showThinking) return;
+		this.activeTextMessageId = void 0;
+		const key = messageId ?? `__auto-reasoning-${++this.reasoningPartCounter}`;
+		if (!this.reasoningParts.has(key)) {
+			this.reasoningParts.set(key, "");
+			this.partOrder.push({
+				kind: "reasoning",
+				key
+			});
+		}
+		this.activeReasoningKey = key;
+		this.emit();
+	}
+	handleReasoningContent(delta, messageId) {
+		if (!this.showThinking || !delta) return;
+		if (!this.activeReasoningKey) this.handleReasoningStart(messageId);
+		const key = this.activeReasoningKey;
+		if (!key) return;
+		this.reasoningParts.set(key, (this.reasoningParts.get(key) ?? "") + delta);
+		this.emit();
+	}
+	handleReasoningEnd() {
+		if (!this.showThinking) return;
+		this.activeReasoningKey = void 0;
+		this.emit();
+	}
+};
+//#endregion
+export { AG_UI_METADATA_NAMESPACE, RunAggregator };
+
 //# sourceMappingURL=run-aggregator.js.map