assistant-stream/src/core/modules/assistant-stream.ts

Streaming utilities for AI assistants

import { AssistantStream } from "../AssistantStream";
import type { AssistantStreamChunk, PartInit } from "../AssistantStreamChunk";
import {
  createMergeStream,
  type MergeStreamFinishOrder,
} from "../utils/stream/merge";
import { TextStreamControllerImpl, type TextStreamController } from "./text";
import {
  ToolCallStreamControllerImpl,
  type ToolCallStreamController,
} from "./tool-call";
import { PathAppendEncoder } from "../utils/stream/path-utils";
import type {
  DataPart,
  FilePart,
  PartProviderMetadata,
  SourcePart,
} from "../utils/types";
import { generateId } from "../utils/generateId";
import type {
  ReadonlyJSONObject,
  ReadonlyJSONValue,
} from "../../utils/json/json-value";
import type { ToolResponseLike } from "../tool/ToolResponse";
import { promiseWithResolvers } from "../../utils/promiseWithResolvers";
import { createControllerStreamPair } from "../utils/stream/createControllerStream";

type ToolCallPartInit = {
  toolCallId?: string;
  toolName: string;
  argsText?: string;
  args?: ReadonlyJSONObject;
  response?: ToolResponseLike<ReadonlyJSONValue>;
};

type ReasoningPartInit = {
  /** @deprecated Experimental since 2026-08-07. Not scheduled for removal; the API may change in any release. */
  unstable_summary?: string;
  providerMetadata?: PartProviderMetadata;
};

type TextPartInit = {
  providerMetadata?: PartProviderMetadata;
};

/**
 * Imperative writer for constructing an {@link AssistantStream}.
 *
 * The controller handles part boundaries for common streaming operations. Use
 * `appendText` and `appendReasoning` for simple token streams, or open explicit
 * parts with `addTextPart` and `addToolCallPart` when you need direct control.
 */
export type AssistantStreamController = {
  /** Appends text to the current text part, opening one if needed. */
  appendText(textDelta: string): void;
  /**
   * Appends reasoning text to the current reasoning part, opening one if
   * needed. Supplying options opens a new part and applies them to it, so a
   * summary always lands on a part of its own.
   */
  appendReasoning(reasoningDelta: string, options?: ReasoningPartInit): void;
  /** Appends a source citation part to the stream. */
  appendSource(options: SourcePart): void;
  /** Appends a file part to the stream. */
  appendFile(options: FilePart): void;
  /** Appends a named data part to the stream. */
  appendData(options: DataPart): void;
  /**
   * Opens a text part and returns its writer.
   *
   * Close the returned controller when the text part is complete. Opening a new
   * part through this controller closes any implicit text or reasoning append
   * part first.
   */
  addTextPart(options?: TextPartInit): TextStreamController;
  /**
   * Opens a reasoning part and returns its writer.
   *
   * Use the options object to provide an app-authored summary for the part.
   * A summary makes the data stream emit a reasoning part-start frame, which
   * decoders older than that frame reject, so a stream only requires a current
   * client once it uses the field.
   */
  addReasoningPart(options?: ReasoningPartInit): TextStreamController;
  /**
   * Opens a tool-call part by tool name and returns its writer.
   *
   * A tool call id is generated automatically. Use the object overload when the
   * caller already has an id, initial args, args text, or response.
   */
  addToolCallPart(options: string): ToolCallStreamController;
  /**
   * Opens a tool-call part and returns its writer.
   *
   * Use this overload to provide a stable `toolCallId`, initial arguments,
   * streamed argument text, or an immediate {@link ToolResponseLike}.
   */
  addToolCallPart(options: ToolCallPartInit): ToolCallStreamController;
  /** Enqueues a raw protocol chunk. Prefer higher-level helpers when possible. */
  enqueue(chunk: AssistantStreamChunk): void;
  /**
   * Merges another assistant stream into this stream.
   *
   * Paths from the merged stream are remapped so its parts appear at the next
   * available position in this controller's output.
   */
  merge(stream: AssistantStream): void;
  /** Closes any active part and finishes the stream. */
  close(): void;
  /**
   * Returns a controller that writes child parts with `parentId` attached.
   *
   * Use this for nested or related parts that should be associated with an
   * existing message or part in downstream renderers.
   */
  withParentId(parentId: string): AssistantStreamController;
};

type AssistantStreamOptions = {
  strict?: boolean | undefined;
};

const createCountedPartStream = <TController>(
  makeController: (
    controller: ReadableStreamDefaultController<AssistantStreamChunk>,
  ) => TController,
) => {
  let written = 0;
  const [stream, controller] = createControllerStreamPair<
    AssistantStreamChunk,
    TController
  >((c) =>
    makeController({
      get desiredSize() {
        return c.desiredSize;
      },
      enqueue(chunk) {
        c.enqueue(chunk);
        written++;
      },
      close() {
        c.close();
      },
      error(reason) {
        c.error(reason);
      },
    }),
  );
  return [stream, controller, () => written] as const;
};

// Shared state between controller instances
type AssistantStreamControllerState = {
  strict: boolean;
  merger: ReturnType<typeof createMergeStream>;
  append?:
    | {
        controller: TextStreamController;
        kind: "text" | "reasoning";
        parentId: string | undefined;
      }
    | undefined;
  openInputs: Set<() => void>;
  lastToolCallFinish: Promise<void>;
  closeSubscriber?: () => void;
};

class AssistantStreamControllerImpl implements AssistantStreamController {
  private readonly _state: AssistantStreamControllerState;
  private _parentId?: string;

  constructor(
    state?: AssistantStreamControllerState,
    options: AssistantStreamOptions = {},
  ) {
    this._state = state || {
      strict: options.strict ?? true,
      merger: createMergeStream(),
      openInputs: new Set(),
      lastToolCallFinish: Promise.resolve(),
    };
  }

  get __internal_isClosed() {
    return (
      this._state.merger.isSealed() ||
      this._state.merger.isCancelled() ||
      this._state.merger.isErrored()
    );
  }

  get __internal_isCancelled() {
    return this._state.merger.isCancelled();
  }

  __internal_getReadable() {
    return this._state.merger.readable;
  }

  __internal_endOpenInputs() {
    this._state.append = undefined;
    for (const end of this._state.openInputs) end();
    this._state.openInputs.clear();
  }

  __internal_subscribeToClose(callback: () => void) {
    this._state.closeSubscriber = callback;
  }

  private _addTransformedStream(
    stream: AssistantStream,
    transformer: ReadableWritablePair<
      AssistantStreamChunk,
      AssistantStreamChunk
    >,
    options?: {
      orderedFinish?: MergeStreamFinishOrder | undefined;
      written?: () => number;
      pathTransform?: "merge" | undefined;
    },
  ) {
    if (stream.locked) {
      throw new TypeError(
        "Cannot merge a stream that is already locked to a reader.",
      );
    }

    const pipeTask = stream
      .pipeTo(transformer.writable)
      .catch(async (error) => {
        await transformer.writable.abort(error).catch(() => undefined);
        throw error;
      });
    this._state.merger.addStream(transformer.readable, pipeTask, options);
    return pipeTask;
  }

  private _trackOpenInput(pipeTask: Promise<void>, end: () => void) {
    const { openInputs } = this._state;
    openInputs.add(end);
    const forget = () => openInputs.delete(end);
    pipeTask.then(forget, forget);
  }

  private _addPart(
    part: PartInit,
    stream: AssistantStream,
    written: () => number,
    end?: () => void,
    orderedFinish?: MergeStreamFinishOrder,
  ) {
    if (this._state.append) {
      this._state.append.controller.close();
      this._state.append = undefined;
    }

    const partIndex = this._state.merger.enqueuePartStart({
      type: "part-start",
      part,
      path: [],
    });
    const pipeTask = this._addTransformedStream(
      stream,
      new PathAppendEncoder(partIndex),
      { orderedFinish, written },
    );
    if (end) this._trackOpenInput(pipeTask, end);
  }

  merge(stream: AssistantStream) {
    // The identity stage delays the merged stream's buffered chunks, so a raw
    // chunk enqueued before they are delivered lands ahead of the merged part,
    // never between its part-start and deltas.
    this._addTransformedStream(
      stream,
      new TransformStream<AssistantStreamChunk, AssistantStreamChunk>(),
      { pathTransform: "merge" },
    );
  }

  appendText(textDelta: string) {
    if (
      this._state.append?.kind !== "text" ||
      this._state.append.parentId !== this._parentId
    ) {
      this._state.append = {
        kind: "text",
        parentId: this._parentId,
        controller: this.addTextPart(),
      };
    }
    this._state.append.controller.append(textDelta);
  }

  appendReasoning(textDelta: string, options?: ReasoningPartInit) {
    // An init describes a part, so supplying one opens a part rather than
    // extending whichever one happens to be open.
    if (
      options !== undefined ||
      this._state.append?.kind !== "reasoning" ||
      this._state.append.parentId !== this._parentId
    ) {
      this._state.append = {
        kind: "reasoning",
        parentId: this._parentId,
        controller: this.addReasoningPart(options),
      };
    }
    // Opening a part from an init with no text is how a summary-only part is
    // created, and appending there would put a chunk on the stream that no
    // producer emitted. An empty delta from an ordinary caller still appends.
    if (options !== undefined && textDelta.length === 0) return;
    this._state.append.controller.append(textDelta);
  }

  addTextPart(options?: TextPartInit) {
    const [stream, controller, written] = createCountedPartStream(
      (c) => new TextStreamControllerImpl(c, { strict: this._state.strict }),
    );
    this._addPart(
      this._withParentIdOption({ type: "text", ...options }),
      stream,
      written,
      () => controller.__internal_close(),
    );
    return controller;
  }

  addReasoningPart(options?: ReasoningPartInit) {
    const [stream, controller, written] = createCountedPartStream(
      (c) => new TextStreamControllerImpl(c, { strict: this._state.strict }),
    );
    this._addPart(
      this._withParentIdOption({ type: "reasoning", ...options }),
      stream,
      written,
      () => controller.__internal_close(),
    );
    return controller;
  }

  addToolCallPart(
    options: string | ToolCallPartInit,
  ): ToolCallStreamController {
    const opt = typeof options === "string" ? { toolName: options } : options;
    const toolName = opt.toolName;
    const toolCallId = opt.toolCallId ?? generateId();

    const delivered = promiseWithResolvers<void>();
    let previous = Promise.resolve();
    const [stream, controller, written] = createCountedPartStream(
      (c) =>
        new ToolCallStreamControllerImpl(c, {
          strict: this._state.strict,
          onClose: () => {
            previous = this._state.lastToolCallFinish;
            this._state.lastToolCallFinish = delivered.promise;
          },
        }),
    );
    this._addPart(
      {
        type: "tool-call",
        toolName,
        toolCallId,
        ...(this._parentId && { parentId: this._parentId }),
      },
      stream,
      written,
      () => controller.__internal_truncate(),
      { previous: () => previous, delivered },
    );

    if (opt.argsText !== undefined) {
      controller.argsText.append(opt.argsText);
      controller.argsText.close();
    }
    if (opt.args !== undefined) {
      controller.argsText.append(JSON.stringify(opt.args));
      controller.argsText.close();
    }
    if (opt.response !== undefined) {
      controller.setResponse(opt.response);
    }

    return controller;
  }

  private _finishedPartStream(): AssistantStream {
    return new ReadableStream({
      start(controller) {
        controller.enqueue({ type: "part-finish", path: [] });
        controller.close();
      },
    });
  }

  private _withParentIdOption<T>(options: T): T {
    if (!this._parentId) return options;
    return { ...options, parentId: this._parentId };
  }

  appendSource(options: SourcePart) {
    this._addPart(
      this._withParentIdOption(options),
      this._finishedPartStream(),
      () => 1,
    );
  }

  appendFile(options: FilePart) {
    this._addPart(
      this._withParentIdOption(options),
      this._finishedPartStream(),
      () => 1,
    );
  }

  appendData(options: DataPart) {
    this._addPart(
      this._withParentIdOption(options),
      this._finishedPartStream(),
      () => 1,
    );
  }

  enqueue(chunk: AssistantStreamChunk) {
    this._state.merger.enqueue(chunk);
  }

  withParentId(parentId: string): AssistantStreamController {
    const controller = new AssistantStreamControllerImpl(this._state);
    controller._parentId = parentId;
    return controller;
  }

  close() {
    this._state.append?.controller?.close();
    this._state.merger.seal();

    this._state.closeSubscriber?.();
  }
}

/**
 * Creates an {@link AssistantStream} and writes to it with an
 * {@link AssistantStreamController}.
 *
 * The callback may write synchronously or asynchronously. If it throws while
 * the stream is open, an `error` chunk is emitted. Failures after an explicit
 * close are logged, while failures after consumer cancellation are discarded.
 * When the callback settles, the stream is closed automatically unless the
 * controller was already closed.
 */
export function createAssistantStream(
  callback: (controller: AssistantStreamController) => PromiseLike<void> | void,
  options: AssistantStreamOptions = {},
): AssistantStream {
  const controller = new AssistantStreamControllerImpl(undefined, options);

  const runTask = async () => {
    try {
      await callback(controller);
    } catch (e) {
      if (!controller.__internal_isClosed) {
        controller.enqueue({
          type: "error",
          path: [],
          error: String(e),
        });
        controller.__internal_endOpenInputs();
      } else if (!controller.__internal_isCancelled) {
        console.error(e);
      }
    } finally {
      if (!controller.__internal_isClosed) {
        controller.close();
      }
    }
  };
  runTask();

  return controller.__internal_getReadable();
}

/**
 * Creates an {@link AssistantStream} together with the controller used to
 * write into it.
 *
 * Use this when the stream needs to be returned before all writers are known.
 * Closing the returned controller finishes the paired stream.
 */
export function createAssistantStreamController(
  options: AssistantStreamOptions = {},
) {
  const { resolve, promise } = promiseWithResolvers<void>();
  let controller!: AssistantStreamController;
  const stream = createAssistantStream((c) => {
    controller = c;

    (controller as AssistantStreamControllerImpl).__internal_subscribeToClose(
      resolve,
    );

    return promise;
  }, options);
  return [stream, controller] as const;
}