assistant-stream/dist/core/serialization/ui-message-stream/UIMessageStream.js

Streaming utilities for AI assistants

import { generateId } from "../../utils/generateId.js";
import { PipeableTransformStream } from "../../utils/stream/PipeableTransformStream.js";
import { createSSEJsonDecoder } from "../../utils/stream/SSEJson.js";
import { AssistantTransformStream } from "../../utils/stream/AssistantTransformStream.js";
import { createToolCallPartRegistry } from "../tool-call-part-registry.js";
import { createChunkNormalizer } from "./chunk-normalizer.js";
import sjson from "secure-json-parse";
// src/core/serialization/ui-message-stream/UIMessageStream.ts
const isDataChunk = (chunk) => chunk.type.startsWith("data-");
/**
* Decodes AI SDK v6 UI Message Stream format into AssistantStreamChunks.
*/
var UIMessageStreamDecoder = class extends PipeableTransformStream {
	constructor(options = {}) {
		super((readable) => {
			const toolCallPartRegistry = createToolCallPartRegistry();
			const normalizer = createChunkNormalizer();
			const textParts = {
				text: /* @__PURE__ */ new Map(),
				reasoning: /* @__PURE__ */ new Map()
			};
			const activeTextIds = {};
			const closeTextParts = () => {
				for (const kind of ["text", "reasoning"]) {
					for (const part of textParts[kind].values()) part.close();
					textParts[kind].clear();
					delete activeTextIds[kind];
				}
			};
			let activeToolCallId;
			let currentMessageId;
			let stepOpenedByStart = false;
			const approvalsByToolCallId = /* @__PURE__ */ new Map();
			const transformer = {
				transform(chunk, controller) {
					const type = chunk.type;
					if (isDataChunk(chunk)) {
						const name = chunk.type.slice(5);
						if (options.onData) options.onData({
							type: chunk.type,
							name,
							data: chunk.data,
							...chunk.transient !== void 0 && { transient: chunk.transient }
						});
						if (!chunk.transient) controller.enqueue({
							type: "data",
							path: [],
							data: [{
								name,
								data: chunk.data
							}]
						});
						return;
					}
					switch (type) {
						case "start":
							currentMessageId = chunk.messageId;
							stepOpenedByStart = true;
							controller.enqueue({
								type: "step-start",
								path: [],
								messageId: chunk.messageId
							});
							break;
						case "text-start":
						case "reasoning-start": {
							const kind = type === "text-start" ? "text" : "reasoning";
							textParts[kind].get(chunk.id)?.close();
							textParts[kind].set(chunk.id, kind === "text" ? controller.addTextPart(chunk.providerMetadata != null ? { providerMetadata: chunk.providerMetadata } : void 0) : controller.addReasoningPart(chunk.providerMetadata != null ? { providerMetadata: chunk.providerMetadata } : void 0));
							activeTextIds[kind] = chunk.id;
							break;
						}
						case "text-delta":
						case "reasoning-delta": {
							const kind = type === "text-delta" ? "text" : "reasoning";
							const id = chunk.id ?? activeTextIds[kind];
							const part = id === void 0 ? void 0 : textParts[kind].get(id);
							const delta = type === "text-delta" ? chunk.textDelta : chunk.delta;
							if (part) part.append(delta, chunk.providerMetadata);
							else if (kind === "text") controller.appendText(delta);
							else controller.appendReasoning(delta);
							break;
						}
						case "text-end":
						case "reasoning-end": {
							const kind = type === "text-end" ? "text" : "reasoning";
							const id = chunk.id ?? activeTextIds[kind];
							if (id === void 0) break;
							textParts[kind].get(id)?.close(chunk.providerMetadata);
							textParts[kind].delete(id);
							if (activeTextIds[kind] === id) delete activeTextIds[kind];
							break;
						}
						case "source":
							controller.appendSource({
								type: "source",
								sourceType: chunk.source.sourceType,
								id: chunk.source.id,
								url: chunk.source.url,
								...chunk.source.title && { title: chunk.source.title }
							});
							break;
						case "source-document":
							controller.appendSource({
								type: "source",
								sourceType: "document",
								id: chunk.sourceId,
								title: chunk.title,
								mediaType: chunk.mediaType,
								...chunk.filename && { filename: chunk.filename }
							});
							break;
						case "file":
							controller.appendFile({
								type: "file",
								mimeType: chunk.file.mimeType,
								data: chunk.file.data
							});
							break;
						case "tool-call-start":
							if (activeToolCallId !== void 0) {
								toolCallPartRegistry.closeArgsText(toolCallPartRegistry.get(activeToolCallId));
								activeToolCallId = void 0;
							}
							toolCallPartRegistry.start(chunk.toolCallId, () => controller.addToolCallPart({
								toolCallId: chunk.toolCallId,
								toolName: chunk.toolName
							}));
							activeToolCallId = chunk.toolCallId;
							break;
						case "tool-call-delta":
							if (activeToolCallId !== void 0) toolCallPartRegistry.appendArgsText(toolCallPartRegistry.get(activeToolCallId), chunk.argsText);
							break;
						case "tool-call-end":
							if (activeToolCallId !== void 0) {
								toolCallPartRegistry.closeArgsText(toolCallPartRegistry.get(activeToolCallId));
								activeToolCallId = void 0;
							}
							break;
						case "tool-result": {
							if (chunk.toolCallId === activeToolCallId) activeToolCallId = void 0;
							const toolCallController = toolCallPartRegistry.tryGet(chunk.toolCallId);
							if (!toolCallController) throw new Error(`Encountered tool result with unknown id: ${chunk.toolCallId}`);
							toolCallPartRegistry.setResponse(toolCallController, {
								result: chunk.result,
								isError: chunk.isError ?? false,
								...chunk.isPreliminary ? { isPreliminary: true } : {},
								...chunk.messages !== void 0 ? { messages: chunk.messages } : {}
							});
							break;
						}
						case "tool-approval-request": {
							const toolCallController = toolCallPartRegistry.tryGet(chunk.toolCallId);
							if (!toolCallController) break;
							approvalsByToolCallId.set(chunk.toolCallId, chunk.approval);
							toolCallController.setApproval(chunk.approval);
							break;
						}
						case "tool-approval-response": {
							let entry;
							for (const candidate of approvalsByToolCallId) {
								if (candidate[1].id !== chunk.approvalId) continue;
								entry = candidate;
								break;
							}
							if (!entry) break;
							const [toolCallId, approval] = entry;
							const decided = {
								...approval,
								approved: chunk.approved,
								...chunk.reason !== void 0 && { reason: chunk.reason }
							};
							approvalsByToolCallId.set(toolCallId, decided);
							toolCallPartRegistry.tryGet(toolCallId)?.setApproval(decided);
							break;
						}
						case "tool-output-denied": {
							const toolCallController = toolCallPartRegistry.tryGet(chunk.toolCallId);
							if (!toolCallController) break;
							const reason = approvalsByToolCallId.get(chunk.toolCallId)?.reason;
							toolCallPartRegistry.setResponse(toolCallController, {
								result: { error: reason || "Tool approval denied" },
								isError: true
							});
							break;
						}
						case "start-step": {
							const messageId = chunk.messageId ?? currentMessageId ?? generateId();
							const isOpenStep = stepOpenedByStart && messageId === currentMessageId;
							stepOpenedByStart = false;
							if (isOpenStep) break;
							controller.enqueue({
								type: "step-start",
								path: [],
								messageId
							});
							break;
						}
						case "finish-step":
							stepOpenedByStart = false;
							closeTextParts();
							controller.enqueue({
								type: "step-finish",
								path: [],
								finishReason: chunk.finishReason,
								usage: chunk.usage,
								isContinued: chunk.isContinued
							});
							break;
						case "finish":
							closeTextParts();
							controller.enqueue({
								type: "message-finish",
								path: [],
								finishReason: chunk.finishReason,
								usage: chunk.usage
							});
							break;
						case "error": controller.enqueue({
							type: "error",
							path: [],
							error: chunk.errorText
						});
					}
				},
				flush(controller) {
					normalizer.flush({ enqueue(chunk) {
						transformer.transform(chunk, controller);
					} });
					closeTextParts();
					if (activeToolCallId !== void 0) toolCallPartRegistry.closeArgsText(toolCallPartRegistry.get(activeToolCallId));
					toolCallPartRegistry.closeAll();
				}
			};
			const transform = new AssistantTransformStream(transformer);
			return createSSEJsonDecoder({
				strict: true,
				maxLineLength: options.maxLineLength,
				maxEventLength: options.maxEventLength,
				parse(data, controller) {
					let chunk;
					try {
						chunk = sjson.parse(data);
					} catch {
						chunk = void 0;
					}
					if (typeof chunk !== "object" || chunk === null || Array.isArray(chunk) || typeof chunk.type !== "string") {
						console.warn(`Dropped invalid UIMessageStream chunk: ${data.slice(0, 200)}`);
						return;
					}
					normalizer.normalize(chunk, controller);
				},
				done: {
					marker: "[DONE]",
					onMissing() {
						throw new Error("Stream ended abruptly without receiving [DONE] marker");
					}
				}
			})(readable).pipeThrough(transform);
		});
	}
};
export { UIMessageStreamDecoder };

//# sourceMappingURL=UIMessageStream.js.map