assistant-stream/dist/core/serialization/data-stream/DataStream.js

Streaming utilities for AI assistants

import { PipeableTransformStream } from "../../utils/stream/PipeableTransformStream.js";
import { AssistantTransformStream } from "../../utils/stream/AssistantTransformStream.js";
import { DataStreamStreamChunkType } from "./chunk-types.js";
import { LineDecoderStream } from "../../utils/stream/LineDecoderStream.js";
import { DataStreamChunkDecoder, DataStreamChunkEncoder } from "./serialization.js";
import { AssistantMetaTransformStream } from "../../utils/stream/AssistantMetaTransformStream.js";
import { createToolCallPartRegistry } from "../tool-call-part-registry.js";
// src/core/serialization/data-stream/DataStream.ts
const isString = (value) => typeof value === "string";
const isArray = (value) => Array.isArray(value);
const isObject = (value) => typeof value === "object" && value !== null && !Array.isArray(value);
const optional = (check) => (value) => value === void 0 || value === null || check(value);
const isBoolean = (value) => typeof value === "boolean";
const present = (value) => value !== void 0;
const objectWith = (fields) => {
	const entries = Object.entries(fields);
	return (value) => isObject(value) && entries.every(([key, check]) => check(value[key]));
};
const unchecked = () => true;
const isPath = (value) => Array.isArray(value) && value.every(isString);
const isStateOperation = (value) => isObject(value) && isPath(value["path"]) && (value["type"] === "set" ? present(value["value"]) : value["type"] === "append-text" && isString(value["value"]));
const FINISH_REASONS = /* @__PURE__ */ new Set([
	"stop",
	"length",
	"content-filter",
	"tool-calls",
	"error",
	"other",
	"unknown"
]);
const readFinishFields = (value) => ({
	finishReason: FINISH_REASONS.has(value.finishReason) ? value.finishReason : "other",
	usage: isObject(value.usage) ? value.usage : {
		inputTokens: 0,
		outputTokens: 0
	}
});
const VALUE_RULES = {
	[DataStreamStreamChunkType.TextDelta]: isString,
	[DataStreamStreamChunkType.Data]: isArray,
	[DataStreamStreamChunkType.Error]: unchecked,
	[DataStreamStreamChunkType.Annotation]: isArray,
	[DataStreamStreamChunkType.ToolCall]: objectWith({
		toolCallId: isString,
		toolName: isString,
		args: optional((value) => typeof value === "object")
	}),
	[DataStreamStreamChunkType.ToolCallResult]: objectWith({
		toolCallId: isString,
		isError: optional(isBoolean),
		isPreliminary: optional(isBoolean),
		modelContent: optional(isArray)
	}),
	[DataStreamStreamChunkType.StartToolCall]: objectWith({
		toolCallId: isString,
		toolName: isString,
		parentId: optional(isString)
	}),
	[DataStreamStreamChunkType.ToolCallArgsTextDelta]: objectWith({
		toolCallId: isString,
		argsTextDelta: isString,
		isFinal: optional(isBoolean)
	}),
	[DataStreamStreamChunkType.FinishMessage]: objectWith({
		finishReason: isString,
		usage: optional(isObject)
	}),
	[DataStreamStreamChunkType.FinishStep]: objectWith({
		finishReason: isString,
		usage: optional(isObject),
		isContinued: optional(isBoolean)
	}),
	[DataStreamStreamChunkType.StartStep]: objectWith({ messageId: isString }),
	[DataStreamStreamChunkType.ReasoningDelta]: isString,
	[DataStreamStreamChunkType.Source]: objectWith({ parentId: optional(isString) }),
	[DataStreamStreamChunkType.RedactedReasoning]: unchecked,
	[DataStreamStreamChunkType.ReasoningSignature]: unchecked,
	[DataStreamStreamChunkType.File]: objectWith({
		data: isString,
		mimeType: isString,
		parentId: optional(isString)
	}),
	[DataStreamStreamChunkType.AuiUpdateStateOperations]: (value) => Array.isArray(value) && value.every(isStateOperation),
	[DataStreamStreamChunkType.AuiTextDelta]: objectWith({
		textDelta: isString,
		parentId: isString
	}),
	[DataStreamStreamChunkType.AuiReasoningDelta]: objectWith({
		reasoningDelta: isString,
		parentId: isString
	}),
	[DataStreamStreamChunkType.AuiDataPart]: objectWith({
		name: isString,
		data: present,
		parentId: optional(isString)
	}),
	[DataStreamStreamChunkType.AuiReasoningPartStart]: objectWith({
		unstable_summary: optional(isString),
		parentId: optional(isString)
	}),
	[DataStreamStreamChunkType.AuiToolCallApproval]: objectWith({
		toolCallId: isString,
		approval: objectWith({
			id: isString,
			approved: (value) => value === void 0 || isBoolean(value)
		})
	})
};
const ERROR_SEVERITIES = /* @__PURE__ */ new Set([
	"critical",
	"warning",
	"info"
]);
const readErrorValue = (value) => {
	if (typeof value === "string") return {
		type: "error",
		path: [],
		error: value
	};
	if (typeof value === "object" && value !== null && "error" in value && typeof value.error === "string") {
		const code = "code" in value ? value.code : void 0;
		const severity = "severity" in value ? value.severity : void 0;
		return {
			type: "error",
			path: [],
			error: value.error,
			...typeof code === "string" ? { code } : {},
			...ERROR_SEVERITIES.has(severity) ? { severity } : {}
		};
	}
	return {
		type: "error",
		path: [],
		error: JSON.stringify(value) ?? ""
	};
};
var DataStreamEncoder = class extends PipeableTransformStream {
	headers = new Headers({
		"Content-Type": "text/plain; charset=utf-8",
		"x-vercel-ai-data-stream": "v1"
	});
	constructor() {
		super((readable) => {
			const openToolCallArgs = /* @__PURE__ */ new Map();
			const finishToolCallArgs = (controller, toolCallId) => {
				const hasArgsText = openToolCallArgs.get(toolCallId);
				if (hasArgsText === void 0) return;
				openToolCallArgs.delete(toolCallId);
				controller.enqueue({
					type: DataStreamStreamChunkType.ToolCallArgsTextDelta,
					value: {
						toolCallId,
						argsTextDelta: hasArgsText ? "" : "{}",
						isFinal: true
					}
				});
			};
			const finishOpenToolCallArgs = (controller) => {
				for (const toolCallId of openToolCallArgs.keys()) finishToolCallArgs(controller, toolCallId);
			};
			const transform = new TransformStream({
				transform(chunk, controller) {
					const type = chunk.type;
					switch (type) {
						case "part-start": {
							const part = chunk.part;
							if (part.type === "tool-call") {
								const { type, ...value } = part;
								controller.enqueue({
									type: DataStreamStreamChunkType.StartToolCall,
									value
								});
								openToolCallArgs.set(part.toolCallId, false);
							}
							if (part.type === "source") {
								const { type, ...value } = part;
								controller.enqueue({
									type: DataStreamStreamChunkType.Source,
									value
								});
							}
							if (part.type === "file") {
								const { type, ...value } = part;
								controller.enqueue({
									type: DataStreamStreamChunkType.File,
									value
								});
							}
							if (part.type === "data") {
								const { type, ...value } = part;
								controller.enqueue({
									type: DataStreamStreamChunkType.AuiDataPart,
									value
								});
							}
							if (part.type === "reasoning" && part.unstable_summary !== void 0) controller.enqueue({
								type: DataStreamStreamChunkType.AuiReasoningPartStart,
								value: {
									unstable_summary: part.unstable_summary,
									...part.parentId !== void 0 ? { parentId: part.parentId } : {}
								}
							});
							break;
						}
						case "text-delta": {
							const part = chunk.meta;
							switch (part.type) {
								case "text":
									if (part.parentId) controller.enqueue({
										type: DataStreamStreamChunkType.AuiTextDelta,
										value: {
											textDelta: chunk.textDelta,
											parentId: part.parentId
										}
									});
									else controller.enqueue({
										type: DataStreamStreamChunkType.TextDelta,
										value: chunk.textDelta
									});
									break;
								case "reasoning":
									if (part.parentId) controller.enqueue({
										type: DataStreamStreamChunkType.AuiReasoningDelta,
										value: {
											reasoningDelta: chunk.textDelta,
											parentId: part.parentId
										}
									});
									else controller.enqueue({
										type: DataStreamStreamChunkType.ReasoningDelta,
										value: chunk.textDelta
									});
									break;
								case "tool-call":
									if (!openToolCallArgs.has(part.toolCallId)) break;
									openToolCallArgs.set(part.toolCallId, true);
									controller.enqueue({
										type: DataStreamStreamChunkType.ToolCallArgsTextDelta,
										value: {
											toolCallId: part.toolCallId,
											argsTextDelta: chunk.textDelta
										}
									});
									break;
								default: throw new Error(`Unsupported part type for text-delta: ${part.type}`);
							}
							break;
						}
						case "result": {
							const part = chunk.meta;
							if (part.type !== "tool-call") throw new Error(`Result chunk on non-tool-call part not supported: ${part.type}`);
							openToolCallArgs.delete(part.toolCallId);
							controller.enqueue({
								type: DataStreamStreamChunkType.ToolCallResult,
								value: {
									toolCallId: part.toolCallId,
									result: chunk.result,
									artifact: chunk.artifact,
									...chunk.isError ? { isError: chunk.isError } : {},
									...chunk.isPreliminary ? { isPreliminary: true } : {},
									...chunk.modelContent !== void 0 ? { modelContent: chunk.modelContent } : {},
									...chunk.messages !== void 0 ? { messages: chunk.messages } : {}
								}
							});
							break;
						}
						case "step-start": {
							const { type, ...value } = chunk;
							controller.enqueue({
								type: DataStreamStreamChunkType.StartStep,
								value
							});
							break;
						}
						case "step-finish": {
							finishOpenToolCallArgs(controller);
							const { type, ...value } = chunk;
							controller.enqueue({
								type: DataStreamStreamChunkType.FinishStep,
								value
							});
							break;
						}
						case "message-finish": {
							finishOpenToolCallArgs(controller);
							const { type, ...value } = chunk;
							controller.enqueue({
								type: DataStreamStreamChunkType.FinishMessage,
								value
							});
							break;
						}
						case "error":
							if (chunk.severity !== "warning" && chunk.severity !== "info") finishOpenToolCallArgs(controller);
							controller.enqueue({
								type: DataStreamStreamChunkType.Error,
								value: chunk.code === void 0 && chunk.severity === void 0 ? chunk.error : {
									error: chunk.error,
									...chunk.code !== void 0 ? { code: chunk.code } : {},
									...chunk.severity !== void 0 ? { severity: chunk.severity } : {}
								}
							});
							break;
						case "annotations":
							controller.enqueue({
								type: DataStreamStreamChunkType.Annotation,
								value: chunk.annotations
							});
							break;
						case "data":
							controller.enqueue({
								type: DataStreamStreamChunkType.Data,
								value: chunk.data
							});
							break;
						case "update-state":
							controller.enqueue({
								type: DataStreamStreamChunkType.AuiUpdateStateOperations,
								value: chunk.operations
							});
							break;
						case "tool-call-args-text-finish":
							finishToolCallArgs(controller, chunk.meta.toolCallId);
							break;
						case "tool-call-approval":
							if (chunk.meta.type !== "tool-call") break;
							controller.enqueue({
								type: DataStreamStreamChunkType.AuiToolCallApproval,
								value: {
									toolCallId: chunk.meta.toolCallId,
									approval: chunk.approval
								}
							});
							break;
						case "part-finish":
							if (chunk.meta.type === "tool-call") finishToolCallArgs(controller, chunk.meta.toolCallId);
							break;
						default: throw new Error(`Unsupported chunk type: ${type}`);
					}
				},
				flush(controller) {
					finishOpenToolCallArgs(controller);
				}
			});
			return readable.pipeThrough(new AssistantMetaTransformStream()).pipeThrough(transform).pipeThrough(new DataStreamChunkEncoder()).pipeThrough(new TextEncoderStream());
		});
	}
};
var DataStreamDecoder = class extends PipeableTransformStream {
	constructor(options = {}) {
		const strict = options.strict ?? true;
		super((readable) => {
			const toolCallPartRegistry = createToolCallPartRegistry();
			const warnedDroppedArgs = /* @__PURE__ */ new Set();
			const loggedDrops = /* @__PURE__ */ new Set();
			const logDropped = (key, message) => {
				if (loggedDrops.has(key) || loggedDrops.size >= 20) return;
				loggedDrops.add(key);
				console.error(message);
			};
			const closeOpenToolCallArgs = () => {
				toolCallPartRegistry.closeOpenArgsText();
			};
			const transform = new AssistantTransformStream({
				strict,
				transform(chunk, controller) {
					const { type, value } = chunk;
					const rule = Object.prototype.hasOwnProperty.call(VALUE_RULES, type) ? VALUE_RULES[type] : void 0;
					if (rule && !rule(value)) {
						const preview = JSON.stringify(value)?.slice(0, 200);
						if (strict) throw new Error(`Invalid value for data-stream chunk type "${type}": ${preview}`);
						logDropped(`value:${type}`, `Dropped data-stream chunk with invalid value for type "${type}": ${preview}`);
						return;
					}
					switch (type) {
						case DataStreamStreamChunkType.ReasoningDelta:
							controller.appendReasoning(value);
							break;
						case DataStreamStreamChunkType.TextDelta:
							controller.appendText(value);
							break;
						case DataStreamStreamChunkType.AuiTextDelta:
							controller.withParentId(value.parentId).appendText(value.textDelta);
							break;
						case DataStreamStreamChunkType.AuiReasoningPartStart: {
							const target = value.parentId ? controller.withParentId(value.parentId) : controller;
							const unstable_summary = value.unstable_summary ?? void 0;
							target.appendReasoning("", { ...unstable_summary !== void 0 ? { unstable_summary } : {} });
							break;
						}
						case DataStreamStreamChunkType.AuiReasoningDelta:
							controller.withParentId(value.parentId).appendReasoning(value.reasoningDelta);
							break;
						case DataStreamStreamChunkType.StartToolCall: {
							const { toolCallId, toolName, parentId } = value;
							const ctrl = parentId ? controller.withParentId(parentId) : controller;
							if (toolCallPartRegistry.tryGet(toolCallId)) {
								if (strict) throw new Error(`Encountered duplicate tool call id: ${toolCallId}`);
								logDropped(`duplicate:${toolCallId}`, `Dropped duplicate tool call start: ${toolCallId}`);
								break;
							}
							toolCallPartRegistry.start(toolCallId, () => ctrl.addToolCallPart({
								toolCallId,
								toolName
							}));
							break;
						}
						case DataStreamStreamChunkType.ToolCallArgsTextDelta: {
							const { toolCallId, argsTextDelta, isFinal } = value;
							const toolCallController = toolCallPartRegistry.tryGet(toolCallId);
							if (!toolCallController) {
								if (strict) throw new Error(`Encountered tool call with unknown id: ${toolCallId}`);
								logDropped(`args:${toolCallId}`, `Dropped args delta for unknown tool call: ${toolCallId}`);
								break;
							}
							if (toolCallPartRegistry.isArgsTextClosed(toolCallController)) {
								if (!warnedDroppedArgs.has(toolCallId)) {
									warnedDroppedArgs.add(toolCallId);
									console.warn(`Dropped tool-call args delta for closed args stream: ${toolCallId}`);
								}
								break;
							}
							if (argsTextDelta.length > 0) toolCallPartRegistry.appendArgsText(toolCallController, argsTextDelta);
							if (isFinal === true) toolCallPartRegistry.closeArgsText(toolCallController);
							break;
						}
						case DataStreamStreamChunkType.ToolCallResult: {
							const { toolCallId, artifact, result, isError, isPreliminary, modelContent, messages } = value;
							const toolCallController = toolCallPartRegistry.tryGet(toolCallId);
							if (!toolCallController) {
								if (strict) throw new Error(`Encountered tool call result with unknown id: ${toolCallId}`);
								logDropped(`result:${toolCallId}`, `Dropped result for unknown tool call: ${toolCallId}`);
								break;
							}
							toolCallPartRegistry.setResponse(toolCallController, {
								artifact,
								result,
								isError,
								...isPreliminary ? { isPreliminary: true } : {},
								...modelContent != null ? { modelContent } : {},
								...messages !== void 0 ? { messages } : {}
							});
							break;
						}
						case DataStreamStreamChunkType.AuiToolCallApproval: {
							const { toolCallId, approval } = value;
							const toolCallController = toolCallPartRegistry.tryGet(toolCallId);
							if (!toolCallController) {
								if (strict) throw new Error(`Encountered tool call approval with unknown id: ${toolCallId}`);
								logDropped(`approval:${toolCallId}`, `Dropped approval for unknown tool call: ${toolCallId}`);
								break;
							}
							toolCallController.setApproval(approval);
							break;
						}
						case DataStreamStreamChunkType.ToolCall: {
							const { toolCallId, toolName } = value;
							const args = value.args ?? void 0;
							const toolCallController = toolCallPartRegistry.tryGet(toolCallId) ?? toolCallPartRegistry.start(toolCallId, () => controller.addToolCallPart({
								toolCallId,
								toolName
							}));
							if (args !== void 0 && !toolCallPartRegistry.hasArgsText(toolCallController) && !toolCallPartRegistry.isArgsTextClosed(toolCallController)) toolCallPartRegistry.appendArgsText(toolCallController, JSON.stringify(args));
							toolCallPartRegistry.closeArgsText(toolCallController);
							break;
						}
						case DataStreamStreamChunkType.FinishMessage:
							closeOpenToolCallArgs();
							controller.enqueue({
								...value,
								...readFinishFields(value),
								type: "message-finish",
								path: []
							});
							break;
						case DataStreamStreamChunkType.StartStep:
							controller.enqueue({
								...value,
								type: "step-start",
								path: []
							});
							break;
						case DataStreamStreamChunkType.FinishStep:
							closeOpenToolCallArgs();
							controller.enqueue({
								...value,
								...readFinishFields(value),
								isContinued: value.isContinued ?? false,
								type: "step-finish",
								path: []
							});
							break;
						case DataStreamStreamChunkType.Data:
							controller.enqueue({
								type: "data",
								path: [],
								data: value
							});
							break;
						case DataStreamStreamChunkType.Annotation:
							controller.enqueue({
								type: "annotations",
								path: [],
								annotations: value
							});
							break;
						case DataStreamStreamChunkType.Source: {
							const { parentId, ...sourceData } = value;
							(parentId ? controller.withParentId(parentId) : controller).appendSource({
								...sourceData,
								type: "source"
							});
							break;
						}
						case DataStreamStreamChunkType.Error:
							controller.enqueue(readErrorValue(value));
							break;
						case DataStreamStreamChunkType.File: {
							const { parentId, ...fileData } = value;
							(parentId ? controller.withParentId(parentId) : controller).appendFile({
								...fileData,
								type: "file"
							});
							break;
						}
						case DataStreamStreamChunkType.AuiDataPart:
							controller.appendData({
								...value,
								type: "data"
							});
							break;
						case DataStreamStreamChunkType.AuiUpdateStateOperations:
							controller.enqueue({
								type: "update-state",
								path: [],
								operations: value
							});
							break;
						case DataStreamStreamChunkType.ReasoningSignature:
						case DataStreamStreamChunkType.RedactedReasoning: break;
						default: {
							const exhaustiveCheck = type;
							if (strict) throw new Error(`unsupported chunk type: ${exhaustiveCheck}`);
							logDropped(`type:${exhaustiveCheck}`, `Dropped unsupported chunk type: ${exhaustiveCheck}`);
						}
					}
				},
				flush() {
					closeOpenToolCallArgs();
					toolCallPartRegistry.closeAll();
				}
			});
			return readable.pipeThrough(new TextDecoderStream()).pipeThrough(new LineDecoderStream({ maxLineLength: options.maxLineLength })).pipeThrough(new DataStreamChunkDecoder()).pipeThrough(transform);
		});
	}
};
export { DataStreamDecoder, DataStreamEncoder };

//# sourceMappingURL=DataStream.js.map