assistant-stream 0.3.16 → 0.3.17

package/dist/core/modules/assistant-stream.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"assistant-stream.d.ts","names":[],"sources":["../../../src/core/modules/assistant-stream.ts"],"mappings":";;;;;;;;;KAuBK,gBAAA;EACH,UAAA;EACA,QAAA;EACA,QAAA;EACA,IAAA,GAAO,kBAAA;EACP,QAAA,GAAW,gBAAA,CAAiB,iBAAA;AAAA;;;;;;;;KAUlB,yBAAA;EAVV,oEAYA,UAAA,CAAW,SAAA,iBAZiB;EAc5B,eAAA,CAAgB,cAAA,iBAd6B;EAgB7C,YAAA,CAAa,OAAA,EAAS,UAAA,SANa;EAQnC,UAAA,CAAW,OAAA,EAAS,QAAA,SAFE;EAItB,UAAA,CAAW,OAAA,EAAS,QAAA;EAAA;;;;;;;EAQpB,WAAA,IAAe,oBAAA;EAgC0C;;;;;;EAzBzD,eAAA,CAAgB,OAAA,WAAkB,wBAAA;EAnBZ;;;;;;EA0BtB,eAAA,CAAgB,OAAA,EAAS,gBAAA,GAAmB,wBAAA,EAtBjC;EAwBX,OAAA,CAAQ,KAAA,EAAO,oBAAA;EAhBA;;;;;;EAuBf,KAAA,CAAM,MAAA,EAAQ,eAAA,SAT8B;EAW5C,KAAA;EATe;;;;;;EAgBf,YAAA,CAAa,QAAA,WAAmB,yBAAA;AAAA;;;AAAyB;AAkM3D;;;;;;iBAAgB,qBAAA,CACd,QAAA,GAAW,UAAA,EAAY,yBAAA,KAA8B,WAAA,gBACpD,eAAA;;;;;;;;iBAiCa,+BAAA,CAAA,aAA+B,eAAA,EAAA,yBAAA;AAA/C;;;;AAA+C;AAsB/C;;AAtBA,iBAsBgB,6BAAA,CACd,QAAA,GAAW,UAAA,EAAY,yBAAA,KAA8B,WAAA,gBAAwB,QAAA"}
+{"version":3,"file":"assistant-stream.d.ts","names":[],"sources":["../../../src/core/modules/assistant-stream.ts"],"mappings":";;;;;;;;;KAuBK,gBAAA;EACH,UAAA;EACA,QAAA;EACA,QAAA;EACA,IAAA,GAAO,kBAAA;EACP,QAAA,GAAW,gBAAA,CAAiB,iBAAA;AAAA;;;;;;;;KAUlB,yBAAA;EAVV,oEAYA,UAAA,CAAW,SAAA,iBAZiB;EAc5B,eAAA,CAAgB,cAAA,iBAd6B;EAgB7C,YAAA,CAAa,OAAA,EAAS,UAAA,SANa;EAQnC,UAAA,CAAW,OAAA,EAAS,QAAA,SAFE;EAItB,UAAA,CAAW,OAAA,EAAS,QAAA;EAAA;;;;;;;EAQpB,WAAA,IAAe,oBAAA;EAgC0C;;;;;;EAzBzD,eAAA,CAAgB,OAAA,WAAkB,wBAAA;EAnBZ;;;;;;EA0BtB,eAAA,CAAgB,OAAA,EAAS,gBAAA,GAAmB,wBAAA,EAtBjC;EAwBX,OAAA,CAAQ,KAAA,EAAO,oBAAA;EAhBA;;;;;;EAuBf,KAAA,CAAM,MAAA,EAAQ,eAAA,SAT8B;EAW5C,KAAA;EATe;;;;;;EAgBf,YAAA,CAAa,QAAA,WAAmB,yBAAA;AAAA;;;AAAyB;AA2M3D;;;;;;iBAAgB,qBAAA,CACd,QAAA,GAAW,UAAA,EAAY,yBAAA,KAA8B,WAAA,gBACpD,eAAA;;;;;;;;iBAiCa,+BAAA,CAAA,aAA+B,eAAA,EAAA,yBAAA;AAA/C;;;;AAA+C;AAsB/C;;AAtBA,iBAsBgB,6BAAA,CACd,QAAA,GAAW,UAAA,EAAY,yBAAA,KAA8B,WAAA,gBAAwB,QAAA"}

package/dist/core/modules/assistant-stream.js

@@ -42,27 +42,29 @@ var AssistantStreamControllerImpl = class AssistantStreamControllerImpl {
 		this._state.merger.addStream(stream.pipeThrough(new PathMergeEncoder(this._state.contentCounter)));
 	}
 	appendText(textDelta) {
-		if (this._state.append?.kind !== "text") this._state.append = {
+		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) {
-		if (this._state.append?.kind !== "reasoning") this._state.append = {
+		if (this._state.append?.kind !== "reasoning" || this._state.append.parentId !== this._parentId) this._state.append = {
 			kind: "reasoning",
+			parentId: this._parentId,
 			controller: this.addReasoningPart()
 		};
 		this._state.append.controller.append(textDelta);
 	}
 	addTextPart() {
 		const [stream, controller] = createTextStreamController();
-		this._addPart({ type: "text" }, stream);
+		this._addPart(this._withParentIdOption({ type: "text" }), stream);
 		return controller;
 	}
 	addReasoningPart() {
 		const [stream, controller] = createTextStreamController();
-		this._addPart({ type: "reasoning" }, stream);
+		this._addPart(this._withParentIdOption({ type: "reasoning" }), stream);
 		return controller;
 	}
 	addToolCallPart(options) {

package/dist/core/modules/assistant-stream.js.map

@@ -1 +1 @@
-{"version":3,"file":"assistant-stream.js","names":[],"sources":["../../../src/core/modules/assistant-stream.ts"],"sourcesContent":["import { AssistantStream } from \"../AssistantStream\";\nimport type { AssistantStreamChunk, PartInit } from \"../AssistantStreamChunk\";\nimport { createMergeStream } from \"../utils/stream/merge\";\nimport { createTextStreamController, type TextStreamController } from \"./text\";\nimport {\n  createToolCallStreamController,\n  type ToolCallStreamController,\n} from \"./tool-call\";\nimport { Counter } from \"../utils/Counter\";\nimport {\n  PathAppendEncoder,\n  PathMergeEncoder,\n} from \"../utils/stream/path-utils\";\nimport { DataStreamEncoder } from \"../serialization/data-stream/DataStream\";\nimport type { DataPart, FilePart, SourcePart } from \"../utils/types\";\nimport { generateId } from \"../utils/generateId\";\nimport type {\n  ReadonlyJSONObject,\n  ReadonlyJSONValue,\n} from \"../../utils/json/json-value\";\nimport type { ToolResponseLike } from \"../tool/ToolResponse\";\nimport { promiseWithResolvers } from \"../../utils/promiseWithResolvers\";\n\ntype ToolCallPartInit = {\n  toolCallId?: string;\n  toolName: string;\n  argsText?: string;\n  args?: ReadonlyJSONObject;\n  response?: ToolResponseLike<ReadonlyJSONValue>;\n};\n\n/**\n * Imperative writer for constructing an {@link AssistantStream}.\n *\n * The controller handles part boundaries for common streaming operations. Use\n * `appendText` and `appendReasoning` for simple token streams, or open explicit\n * parts with `addTextPart` and `addToolCallPart` when you need direct control.\n */\nexport type AssistantStreamController = {\n  /** Appends text to the current text part, opening one if needed. */\n  appendText(textDelta: string): void;\n  /** Appends reasoning text to the current reasoning part, opening one if needed. */\n  appendReasoning(reasoningDelta: string): void;\n  /** Appends a source citation part to the stream. */\n  appendSource(options: SourcePart): void;\n  /** Appends a file part to the stream. */\n  appendFile(options: FilePart): void;\n  /** Appends a named data part to the stream. */\n  appendData(options: DataPart): void;\n  /**\n   * Opens a text part and returns its writer.\n   *\n   * Close the returned controller when the text part is complete. Opening a new\n   * part through this controller closes any implicit text or reasoning append\n   * part first.\n   */\n  addTextPart(): TextStreamController;\n  /**\n   * Opens a tool-call part by tool name and returns its writer.\n   *\n   * A tool call id is generated automatically. Use the object overload when the\n   * caller already has an id, initial args, args text, or response.\n   */\n  addToolCallPart(options: string): ToolCallStreamController;\n  /**\n   * Opens a tool-call part and returns its writer.\n   *\n   * Use this overload to provide a stable `toolCallId`, initial arguments,\n   * streamed argument text, or an immediate {@link ToolResponseLike}.\n   */\n  addToolCallPart(options: ToolCallPartInit): ToolCallStreamController;\n  /** Enqueues a raw protocol chunk. Prefer higher-level helpers when possible. */\n  enqueue(chunk: AssistantStreamChunk): void;\n  /**\n   * Merges another assistant stream into this stream.\n   *\n   * Paths from the merged stream are remapped so its parts appear at the next\n   * available position in this controller's output.\n   */\n  merge(stream: AssistantStream): void;\n  /** Closes any active part and finishes the stream. */\n  close(): void;\n  /**\n   * Returns a controller that writes child parts with `parentId` attached.\n   *\n   * Use this for nested or related parts that should be associated with an\n   * existing message or part in downstream renderers.\n   */\n  withParentId(parentId: string): AssistantStreamController;\n};\n\n// Shared state between controller instances\ntype AssistantStreamControllerState = {\n  merger: ReturnType<typeof createMergeStream>;\n  append?:\n    | {\n        controller: TextStreamController;\n        kind: \"text\" | \"reasoning\";\n      }\n    | undefined;\n  contentCounter: Counter;\n  closeSubscriber?: () => void;\n};\n\nclass AssistantStreamControllerImpl implements AssistantStreamController {\n  private readonly _state: AssistantStreamControllerState;\n  private _parentId?: string;\n\n  constructor(state?: AssistantStreamControllerState) {\n    this._state = state || {\n      merger: createMergeStream(),\n      contentCounter: new Counter(),\n    };\n  }\n\n  get __internal_isClosed() {\n    return this._state.merger.isSealed();\n  }\n\n  __internal_getReadable() {\n    return this._state.merger.readable;\n  }\n\n  __internal_subscribeToClose(callback: () => void) {\n    this._state.closeSubscriber = callback;\n  }\n\n  private _addPart(part: PartInit, stream: AssistantStream) {\n    if (this._state.append) {\n      this._state.append.controller.close();\n      this._state.append = undefined;\n    }\n\n    this.enqueue({\n      type: \"part-start\",\n      part,\n      path: [],\n    });\n    this._state.merger.addStream(\n      stream.pipeThrough(\n        new PathAppendEncoder(this._state.contentCounter.value),\n      ),\n    );\n  }\n\n  merge(stream: AssistantStream) {\n    this._state.merger.addStream(\n      stream.pipeThrough(new PathMergeEncoder(this._state.contentCounter)),\n    );\n  }\n\n  appendText(textDelta: string) {\n    if (this._state.append?.kind !== \"text\") {\n      this._state.append = {\n        kind: \"text\",\n        controller: this.addTextPart(),\n      };\n    }\n    this._state.append.controller.append(textDelta);\n  }\n\n  appendReasoning(textDelta: string) {\n    if (this._state.append?.kind !== \"reasoning\") {\n      this._state.append = {\n        kind: \"reasoning\",\n        controller: this.addReasoningPart(),\n      };\n    }\n    this._state.append.controller.append(textDelta);\n  }\n\n  addTextPart() {\n    const [stream, controller] = createTextStreamController();\n    this._addPart({ type: \"text\" }, stream);\n    return controller;\n  }\n\n  addReasoningPart() {\n    const [stream, controller] = createTextStreamController();\n    this._addPart({ type: \"reasoning\" }, stream);\n    return controller;\n  }\n\n  addToolCallPart(\n    options: string | ToolCallPartInit,\n  ): ToolCallStreamController {\n    const opt = typeof options === \"string\" ? { toolName: options } : options;\n    const toolName = opt.toolName;\n    const toolCallId = opt.toolCallId ?? generateId();\n\n    const [stream, controller] = createToolCallStreamController();\n    this._addPart(\n      {\n        type: \"tool-call\",\n        toolName,\n        toolCallId,\n        ...(this._parentId && { parentId: this._parentId }),\n      },\n      stream,\n    );\n\n    if (opt.argsText !== undefined) {\n      controller.argsText.append(opt.argsText);\n      controller.argsText.close();\n    }\n    if (opt.args !== undefined) {\n      controller.argsText.append(JSON.stringify(opt.args));\n      controller.argsText.close();\n    }\n    if (opt.response !== undefined) {\n      controller.setResponse(opt.response);\n    }\n\n    return controller;\n  }\n\n  private _finishedPartStream(): AssistantStream {\n    return new ReadableStream({\n      start(controller) {\n        controller.enqueue({ type: \"part-finish\", path: [] });\n        controller.close();\n      },\n    });\n  }\n\n  private _withParentIdOption<T>(options: T): T {\n    if (!this._parentId) return options;\n    return { ...options, parentId: this._parentId };\n  }\n\n  appendSource(options: SourcePart) {\n    this._addPart(\n      this._withParentIdOption(options),\n      this._finishedPartStream(),\n    );\n  }\n\n  appendFile(options: FilePart) {\n    this._addPart(\n      this._withParentIdOption(options),\n      this._finishedPartStream(),\n    );\n  }\n\n  appendData(options: DataPart) {\n    this._addPart(\n      this._withParentIdOption(options),\n      this._finishedPartStream(),\n    );\n  }\n\n  enqueue(chunk: AssistantStreamChunk) {\n    this._state.merger.enqueue(chunk);\n\n    if (chunk.type === \"part-start\" && chunk.path.length === 0) {\n      this._state.contentCounter.up();\n    }\n  }\n\n  withParentId(parentId: string): AssistantStreamController {\n    const controller = new AssistantStreamControllerImpl(this._state);\n    controller._parentId = parentId;\n    return controller;\n  }\n\n  close() {\n    this._state.append?.controller?.close();\n    this._state.merger.seal();\n\n    this._state.closeSubscriber?.();\n  }\n}\n\n/**\n * Creates an {@link AssistantStream} and writes to it with an\n * {@link AssistantStreamController}.\n *\n * The callback may write synchronously or asynchronously. If it throws, an\n * `error` chunk is emitted before the error is rethrown; when the callback\n * settles, the stream is closed automatically unless the controller was\n * already closed.\n */\nexport function createAssistantStream(\n  callback: (controller: AssistantStreamController) => PromiseLike<void> | void,\n): AssistantStream {\n  const controller = new AssistantStreamControllerImpl();\n\n  const runTask = async () => {\n    try {\n      await callback(controller);\n    } catch (e) {\n      if (!controller.__internal_isClosed) {\n        controller.enqueue({\n          type: \"error\",\n          path: [],\n          error: String(e),\n        });\n      }\n      throw e;\n    } finally {\n      if (!controller.__internal_isClosed) {\n        controller.close();\n      }\n    }\n  };\n  runTask();\n\n  return controller.__internal_getReadable();\n}\n\n/**\n * Creates an {@link AssistantStream} together with the controller used to\n * write into it.\n *\n * Use this when the stream needs to be returned before all writers are known.\n * Closing the returned controller finishes the paired stream.\n */\nexport function createAssistantStreamController() {\n  const { resolve, promise } = promiseWithResolvers<void>();\n  let controller!: AssistantStreamController;\n  const stream = createAssistantStream((c) => {\n    controller = c;\n\n    (controller as AssistantStreamControllerImpl).__internal_subscribeToClose(\n      resolve,\n    );\n\n    return promise;\n  });\n  return [stream, controller] as const;\n}\n\n/**\n * Creates a `Response` whose body is an encoded {@link AssistantStream}.\n *\n * This is the HTTP-route convenience form of {@link createAssistantStream}; it\n * uses {@link DataStreamEncoder} so the response can be consumed by matching\n * assistant-ui data stream decoders.\n */\nexport function createAssistantStreamResponse(\n  callback: (controller: AssistantStreamController) => PromiseLike<void> | void,\n) {\n  return AssistantStream.toResponse(\n    createAssistantStream(callback),\n    new DataStreamEncoder(),\n  );\n}\n"],"mappings":";;;;;;;;;;AAwGA,IAAM,gCAAN,MAAM,8BAAmE;CACvE;CACA;CAEA,YAAY,OAAwC;EAClD,KAAK,SAAS,SAAS;GACrB,QAAQ,kBAAkB;GAC1B,gBAAgB,IAAI,QAAQ;EAC9B;CACF;CAEA,IAAI,sBAAsB;EACxB,OAAO,KAAK,OAAO,OAAO,SAAS;CACrC;CAEA,yBAAyB;EACvB,OAAO,KAAK,OAAO,OAAO;CAC5B;CAEA,4BAA4B,UAAsB;EAChD,KAAK,OAAO,kBAAkB;CAChC;CAEA,SAAiB,MAAgB,QAAyB;EACxD,IAAI,KAAK,OAAO,QAAQ;GACtB,KAAK,OAAO,OAAO,WAAW,MAAM;GACpC,KAAK,OAAO,SAAS,KAAA;EACvB;EAEA,KAAK,QAAQ;GACX,MAAM;GACN;GACA,MAAM,CAAC;EACT,CAAC;EACD,KAAK,OAAO,OAAO,UACjB,OAAO,YACL,IAAI,kBAAkB,KAAK,OAAO,eAAe,KAAK,CACxD,CACF;CACF;CAEA,MAAM,QAAyB;EAC7B,KAAK,OAAO,OAAO,UACjB,OAAO,YAAY,IAAI,iBAAiB,KAAK,OAAO,cAAc,CAAC,CACrE;CACF;CAEA,WAAW,WAAmB;EAC5B,IAAI,KAAK,OAAO,QAAQ,SAAS,QAC/B,KAAK,OAAO,SAAS;GACnB,MAAM;GACN,YAAY,KAAK,YAAY;EAC/B;EAEF,KAAK,OAAO,OAAO,WAAW,OAAO,SAAS;CAChD;CAEA,gBAAgB,WAAmB;EACjC,IAAI,KAAK,OAAO,QAAQ,SAAS,aAC/B,KAAK,OAAO,SAAS;GACnB,MAAM;GACN,YAAY,KAAK,iBAAiB;EACpC;EAEF,KAAK,OAAO,OAAO,WAAW,OAAO,SAAS;CAChD;CAEA,cAAc;EACZ,MAAM,CAAC,QAAQ,cAAc,2BAA2B;EACxD,KAAK,SAAS,EAAE,MAAM,OAAO,GAAG,MAAM;EACtC,OAAO;CACT;CAEA,mBAAmB;EACjB,MAAM,CAAC,QAAQ,cAAc,2BAA2B;EACxD,KAAK,SAAS,EAAE,MAAM,YAAY,GAAG,MAAM;EAC3C,OAAO;CACT;CAEA,gBACE,SAC0B;EAC1B,MAAM,MAAM,OAAO,YAAY,WAAW,EAAE,UAAU,QAAQ,IAAI;EAClE,MAAM,WAAW,IAAI;EACrB,MAAM,aAAa,IAAI,cAAc,WAAW;EAEhD,MAAM,CAAC,QAAQ,cAAc,+BAA+B;EAC5D,KAAK,SACH;GACE,MAAM;GACN;GACA;GACA,GAAI,KAAK,aAAa,EAAE,UAAU,KAAK,UAAU;EACnD,GACA,MACF;EAEA,IAAI,IAAI,aAAa,KAAA,GAAW;GAC9B,WAAW,SAAS,OAAO,IAAI,QAAQ;GACvC,WAAW,SAAS,MAAM;EAC5B;EACA,IAAI,IAAI,SAAS,KAAA,GAAW;GAC1B,WAAW,SAAS,OAAO,KAAK,UAAU,IAAI,IAAI,CAAC;GACnD,WAAW,SAAS,MAAM;EAC5B;EACA,IAAI,IAAI,aAAa,KAAA,GACnB,WAAW,YAAY,IAAI,QAAQ;EAGrC,OAAO;CACT;CAEA,sBAA+C;EAC7C,OAAO,IAAI,eAAe,EACxB,MAAM,YAAY;GAChB,WAAW,QAAQ;IAAE,MAAM;IAAe,MAAM,CAAC;GAAE,CAAC;GACpD,WAAW,MAAM;EACnB,EACF,CAAC;CACH;CAEA,oBAA+B,SAAe;EAC5C,IAAI,CAAC,KAAK,WAAW,OAAO;EAC5B,OAAO;GAAE,GAAG;GAAS,UAAU,KAAK;EAAU;CAChD;CAEA,aAAa,SAAqB;EAChC,KAAK,SACH,KAAK,oBAAoB,OAAO,GAChC,KAAK,oBAAoB,CAC3B;CACF;CAEA,WAAW,SAAmB;EAC5B,KAAK,SACH,KAAK,oBAAoB,OAAO,GAChC,KAAK,oBAAoB,CAC3B;CACF;CAEA,WAAW,SAAmB;EAC5B,KAAK,SACH,KAAK,oBAAoB,OAAO,GAChC,KAAK,oBAAoB,CAC3B;CACF;CAEA,QAAQ,OAA6B;EACnC,KAAK,OAAO,OAAO,QAAQ,KAAK;EAEhC,IAAI,MAAM,SAAS,gBAAgB,MAAM,KAAK,WAAW,GACvD,KAAK,OAAO,eAAe,GAAG;CAElC;CAEA,aAAa,UAA6C;EACxD,MAAM,aAAa,IAAI,8BAA8B,KAAK,MAAM;EAChE,WAAW,YAAY;EACvB,OAAO;CACT;CAEA,QAAQ;EACN,KAAK,OAAO,QAAQ,YAAY,MAAM;EACtC,KAAK,OAAO,OAAO,KAAK;EAExB,KAAK,OAAO,kBAAkB;CAChC;AACF;;;;;;;;;;AAWA,SAAgB,sBACd,UACiB;CACjB,MAAM,aAAa,IAAI,8BAA8B;CAErD,MAAM,UAAU,YAAY;EAC1B,IAAI;GACF,MAAM,SAAS,UAAU;EAC3B,SAAS,GAAG;GACV,IAAI,CAAC,WAAW,qBACd,WAAW,QAAQ;IACjB,MAAM;IACN,MAAM,CAAC;IACP,OAAO,OAAO,CAAC;GACjB,CAAC;GAEH,MAAM;EACR,UAAU;GACR,IAAI,CAAC,WAAW,qBACd,WAAW,MAAM;EAErB;CACF;CACA,QAAQ;CAER,OAAO,WAAW,uBAAuB;AAC3C;;;;;;;;AASA,SAAgB,kCAAkC;CAChD,MAAM,EAAE,SAAS,YAAY,qBAA2B;CACxD,IAAI;CAUJ,OAAO,CATQ,uBAAuB,MAAM;EAC1C,aAAa;EAEb,WAA8C,4BAC5C,OACF;EAEA,OAAO;CACT,CACa,GAAG,UAAU;AAC5B;;;;;;;;AASA,SAAgB,8BACd,UACA;CACA,OAAO,gBAAgB,WACrB,sBAAsB,QAAQ,GAC9B,IAAI,kBAAkB,CACxB;AACF"}
+{"version":3,"file":"assistant-stream.js","names":[],"sources":["../../../src/core/modules/assistant-stream.ts"],"sourcesContent":["import { AssistantStream } from \"../AssistantStream\";\nimport type { AssistantStreamChunk, PartInit } from \"../AssistantStreamChunk\";\nimport { createMergeStream } from \"../utils/stream/merge\";\nimport { createTextStreamController, type TextStreamController } from \"./text\";\nimport {\n  createToolCallStreamController,\n  type ToolCallStreamController,\n} from \"./tool-call\";\nimport { Counter } from \"../utils/Counter\";\nimport {\n  PathAppendEncoder,\n  PathMergeEncoder,\n} from \"../utils/stream/path-utils\";\nimport { DataStreamEncoder } from \"../serialization/data-stream/DataStream\";\nimport type { DataPart, FilePart, SourcePart } from \"../utils/types\";\nimport { generateId } from \"../utils/generateId\";\nimport type {\n  ReadonlyJSONObject,\n  ReadonlyJSONValue,\n} from \"../../utils/json/json-value\";\nimport type { ToolResponseLike } from \"../tool/ToolResponse\";\nimport { promiseWithResolvers } from \"../../utils/promiseWithResolvers\";\n\ntype ToolCallPartInit = {\n  toolCallId?: string;\n  toolName: string;\n  argsText?: string;\n  args?: ReadonlyJSONObject;\n  response?: ToolResponseLike<ReadonlyJSONValue>;\n};\n\n/**\n * Imperative writer for constructing an {@link AssistantStream}.\n *\n * The controller handles part boundaries for common streaming operations. Use\n * `appendText` and `appendReasoning` for simple token streams, or open explicit\n * parts with `addTextPart` and `addToolCallPart` when you need direct control.\n */\nexport type AssistantStreamController = {\n  /** Appends text to the current text part, opening one if needed. */\n  appendText(textDelta: string): void;\n  /** Appends reasoning text to the current reasoning part, opening one if needed. */\n  appendReasoning(reasoningDelta: string): void;\n  /** Appends a source citation part to the stream. */\n  appendSource(options: SourcePart): void;\n  /** Appends a file part to the stream. */\n  appendFile(options: FilePart): void;\n  /** Appends a named data part to the stream. */\n  appendData(options: DataPart): void;\n  /**\n   * Opens a text part and returns its writer.\n   *\n   * Close the returned controller when the text part is complete. Opening a new\n   * part through this controller closes any implicit text or reasoning append\n   * part first.\n   */\n  addTextPart(): TextStreamController;\n  /**\n   * Opens a tool-call part by tool name and returns its writer.\n   *\n   * A tool call id is generated automatically. Use the object overload when the\n   * caller already has an id, initial args, args text, or response.\n   */\n  addToolCallPart(options: string): ToolCallStreamController;\n  /**\n   * Opens a tool-call part and returns its writer.\n   *\n   * Use this overload to provide a stable `toolCallId`, initial arguments,\n   * streamed argument text, or an immediate {@link ToolResponseLike}.\n   */\n  addToolCallPart(options: ToolCallPartInit): ToolCallStreamController;\n  /** Enqueues a raw protocol chunk. Prefer higher-level helpers when possible. */\n  enqueue(chunk: AssistantStreamChunk): void;\n  /**\n   * Merges another assistant stream into this stream.\n   *\n   * Paths from the merged stream are remapped so its parts appear at the next\n   * available position in this controller's output.\n   */\n  merge(stream: AssistantStream): void;\n  /** Closes any active part and finishes the stream. */\n  close(): void;\n  /**\n   * Returns a controller that writes child parts with `parentId` attached.\n   *\n   * Use this for nested or related parts that should be associated with an\n   * existing message or part in downstream renderers.\n   */\n  withParentId(parentId: string): AssistantStreamController;\n};\n\n// Shared state between controller instances\ntype AssistantStreamControllerState = {\n  merger: ReturnType<typeof createMergeStream>;\n  append?:\n    | {\n        controller: TextStreamController;\n        kind: \"text\" | \"reasoning\";\n        parentId: string | undefined;\n      }\n    | undefined;\n  contentCounter: Counter;\n  closeSubscriber?: () => void;\n};\n\nclass AssistantStreamControllerImpl implements AssistantStreamController {\n  private readonly _state: AssistantStreamControllerState;\n  private _parentId?: string;\n\n  constructor(state?: AssistantStreamControllerState) {\n    this._state = state || {\n      merger: createMergeStream(),\n      contentCounter: new Counter(),\n    };\n  }\n\n  get __internal_isClosed() {\n    return this._state.merger.isSealed();\n  }\n\n  __internal_getReadable() {\n    return this._state.merger.readable;\n  }\n\n  __internal_subscribeToClose(callback: () => void) {\n    this._state.closeSubscriber = callback;\n  }\n\n  private _addPart(part: PartInit, stream: AssistantStream) {\n    if (this._state.append) {\n      this._state.append.controller.close();\n      this._state.append = undefined;\n    }\n\n    this.enqueue({\n      type: \"part-start\",\n      part,\n      path: [],\n    });\n    this._state.merger.addStream(\n      stream.pipeThrough(\n        new PathAppendEncoder(this._state.contentCounter.value),\n      ),\n    );\n  }\n\n  merge(stream: AssistantStream) {\n    this._state.merger.addStream(\n      stream.pipeThrough(new PathMergeEncoder(this._state.contentCounter)),\n    );\n  }\n\n  appendText(textDelta: string) {\n    if (\n      this._state.append?.kind !== \"text\" ||\n      this._state.append.parentId !== this._parentId\n    ) {\n      this._state.append = {\n        kind: \"text\",\n        parentId: this._parentId,\n        controller: this.addTextPart(),\n      };\n    }\n    this._state.append.controller.append(textDelta);\n  }\n\n  appendReasoning(textDelta: string) {\n    if (\n      this._state.append?.kind !== \"reasoning\" ||\n      this._state.append.parentId !== this._parentId\n    ) {\n      this._state.append = {\n        kind: \"reasoning\",\n        parentId: this._parentId,\n        controller: this.addReasoningPart(),\n      };\n    }\n    this._state.append.controller.append(textDelta);\n  }\n\n  addTextPart() {\n    const [stream, controller] = createTextStreamController();\n    this._addPart(this._withParentIdOption({ type: \"text\" }), stream);\n    return controller;\n  }\n\n  addReasoningPart() {\n    const [stream, controller] = createTextStreamController();\n    this._addPart(this._withParentIdOption({ type: \"reasoning\" }), stream);\n    return controller;\n  }\n\n  addToolCallPart(\n    options: string | ToolCallPartInit,\n  ): ToolCallStreamController {\n    const opt = typeof options === \"string\" ? { toolName: options } : options;\n    const toolName = opt.toolName;\n    const toolCallId = opt.toolCallId ?? generateId();\n\n    const [stream, controller] = createToolCallStreamController();\n    this._addPart(\n      {\n        type: \"tool-call\",\n        toolName,\n        toolCallId,\n        ...(this._parentId && { parentId: this._parentId }),\n      },\n      stream,\n    );\n\n    if (opt.argsText !== undefined) {\n      controller.argsText.append(opt.argsText);\n      controller.argsText.close();\n    }\n    if (opt.args !== undefined) {\n      controller.argsText.append(JSON.stringify(opt.args));\n      controller.argsText.close();\n    }\n    if (opt.response !== undefined) {\n      controller.setResponse(opt.response);\n    }\n\n    return controller;\n  }\n\n  private _finishedPartStream(): AssistantStream {\n    return new ReadableStream({\n      start(controller) {\n        controller.enqueue({ type: \"part-finish\", path: [] });\n        controller.close();\n      },\n    });\n  }\n\n  private _withParentIdOption<T>(options: T): T {\n    if (!this._parentId) return options;\n    return { ...options, parentId: this._parentId };\n  }\n\n  appendSource(options: SourcePart) {\n    this._addPart(\n      this._withParentIdOption(options),\n      this._finishedPartStream(),\n    );\n  }\n\n  appendFile(options: FilePart) {\n    this._addPart(\n      this._withParentIdOption(options),\n      this._finishedPartStream(),\n    );\n  }\n\n  appendData(options: DataPart) {\n    this._addPart(\n      this._withParentIdOption(options),\n      this._finishedPartStream(),\n    );\n  }\n\n  enqueue(chunk: AssistantStreamChunk) {\n    this._state.merger.enqueue(chunk);\n\n    if (chunk.type === \"part-start\" && chunk.path.length === 0) {\n      this._state.contentCounter.up();\n    }\n  }\n\n  withParentId(parentId: string): AssistantStreamController {\n    const controller = new AssistantStreamControllerImpl(this._state);\n    controller._parentId = parentId;\n    return controller;\n  }\n\n  close() {\n    this._state.append?.controller?.close();\n    this._state.merger.seal();\n\n    this._state.closeSubscriber?.();\n  }\n}\n\n/**\n * Creates an {@link AssistantStream} and writes to it with an\n * {@link AssistantStreamController}.\n *\n * The callback may write synchronously or asynchronously. If it throws, an\n * `error` chunk is emitted before the error is rethrown; when the callback\n * settles, the stream is closed automatically unless the controller was\n * already closed.\n */\nexport function createAssistantStream(\n  callback: (controller: AssistantStreamController) => PromiseLike<void> | void,\n): AssistantStream {\n  const controller = new AssistantStreamControllerImpl();\n\n  const runTask = async () => {\n    try {\n      await callback(controller);\n    } catch (e) {\n      if (!controller.__internal_isClosed) {\n        controller.enqueue({\n          type: \"error\",\n          path: [],\n          error: String(e),\n        });\n      }\n      throw e;\n    } finally {\n      if (!controller.__internal_isClosed) {\n        controller.close();\n      }\n    }\n  };\n  runTask();\n\n  return controller.__internal_getReadable();\n}\n\n/**\n * Creates an {@link AssistantStream} together with the controller used to\n * write into it.\n *\n * Use this when the stream needs to be returned before all writers are known.\n * Closing the returned controller finishes the paired stream.\n */\nexport function createAssistantStreamController() {\n  const { resolve, promise } = promiseWithResolvers<void>();\n  let controller!: AssistantStreamController;\n  const stream = createAssistantStream((c) => {\n    controller = c;\n\n    (controller as AssistantStreamControllerImpl).__internal_subscribeToClose(\n      resolve,\n    );\n\n    return promise;\n  });\n  return [stream, controller] as const;\n}\n\n/**\n * Creates a `Response` whose body is an encoded {@link AssistantStream}.\n *\n * This is the HTTP-route convenience form of {@link createAssistantStream}; it\n * uses {@link DataStreamEncoder} so the response can be consumed by matching\n * assistant-ui data stream decoders.\n */\nexport function createAssistantStreamResponse(\n  callback: (controller: AssistantStreamController) => PromiseLike<void> | void,\n) {\n  return AssistantStream.toResponse(\n    createAssistantStream(callback),\n    new DataStreamEncoder(),\n  );\n}\n"],"mappings":";;;;;;;;;;AAyGA,IAAM,gCAAN,MAAM,8BAAmE;CACvE;CACA;CAEA,YAAY,OAAwC;EAClD,KAAK,SAAS,SAAS;GACrB,QAAQ,kBAAkB;GAC1B,gBAAgB,IAAI,QAAQ;EAC9B;CACF;CAEA,IAAI,sBAAsB;EACxB,OAAO,KAAK,OAAO,OAAO,SAAS;CACrC;CAEA,yBAAyB;EACvB,OAAO,KAAK,OAAO,OAAO;CAC5B;CAEA,4BAA4B,UAAsB;EAChD,KAAK,OAAO,kBAAkB;CAChC;CAEA,SAAiB,MAAgB,QAAyB;EACxD,IAAI,KAAK,OAAO,QAAQ;GACtB,KAAK,OAAO,OAAO,WAAW,MAAM;GACpC,KAAK,OAAO,SAAS,KAAA;EACvB;EAEA,KAAK,QAAQ;GACX,MAAM;GACN;GACA,MAAM,CAAC;EACT,CAAC;EACD,KAAK,OAAO,OAAO,UACjB,OAAO,YACL,IAAI,kBAAkB,KAAK,OAAO,eAAe,KAAK,CACxD,CACF;CACF;CAEA,MAAM,QAAyB;EAC7B,KAAK,OAAO,OAAO,UACjB,OAAO,YAAY,IAAI,iBAAiB,KAAK,OAAO,cAAc,CAAC,CACrE;CACF;CAEA,WAAW,WAAmB;EAC5B,IACE,KAAK,OAAO,QAAQ,SAAS,UAC7B,KAAK,OAAO,OAAO,aAAa,KAAK,WAErC,KAAK,OAAO,SAAS;GACnB,MAAM;GACN,UAAU,KAAK;GACf,YAAY,KAAK,YAAY;EAC/B;EAEF,KAAK,OAAO,OAAO,WAAW,OAAO,SAAS;CAChD;CAEA,gBAAgB,WAAmB;EACjC,IACE,KAAK,OAAO,QAAQ,SAAS,eAC7B,KAAK,OAAO,OAAO,aAAa,KAAK,WAErC,KAAK,OAAO,SAAS;GACnB,MAAM;GACN,UAAU,KAAK;GACf,YAAY,KAAK,iBAAiB;EACpC;EAEF,KAAK,OAAO,OAAO,WAAW,OAAO,SAAS;CAChD;CAEA,cAAc;EACZ,MAAM,CAAC,QAAQ,cAAc,2BAA2B;EACxD,KAAK,SAAS,KAAK,oBAAoB,EAAE,MAAM,OAAO,CAAC,GAAG,MAAM;EAChE,OAAO;CACT;CAEA,mBAAmB;EACjB,MAAM,CAAC,QAAQ,cAAc,2BAA2B;EACxD,KAAK,SAAS,KAAK,oBAAoB,EAAE,MAAM,YAAY,CAAC,GAAG,MAAM;EACrE,OAAO;CACT;CAEA,gBACE,SAC0B;EAC1B,MAAM,MAAM,OAAO,YAAY,WAAW,EAAE,UAAU,QAAQ,IAAI;EAClE,MAAM,WAAW,IAAI;EACrB,MAAM,aAAa,IAAI,cAAc,WAAW;EAEhD,MAAM,CAAC,QAAQ,cAAc,+BAA+B;EAC5D,KAAK,SACH;GACE,MAAM;GACN;GACA;GACA,GAAI,KAAK,aAAa,EAAE,UAAU,KAAK,UAAU;EACnD,GACA,MACF;EAEA,IAAI,IAAI,aAAa,KAAA,GAAW;GAC9B,WAAW,SAAS,OAAO,IAAI,QAAQ;GACvC,WAAW,SAAS,MAAM;EAC5B;EACA,IAAI,IAAI,SAAS,KAAA,GAAW;GAC1B,WAAW,SAAS,OAAO,KAAK,UAAU,IAAI,IAAI,CAAC;GACnD,WAAW,SAAS,MAAM;EAC5B;EACA,IAAI,IAAI,aAAa,KAAA,GACnB,WAAW,YAAY,IAAI,QAAQ;EAGrC,OAAO;CACT;CAEA,sBAA+C;EAC7C,OAAO,IAAI,eAAe,EACxB,MAAM,YAAY;GAChB,WAAW,QAAQ;IAAE,MAAM;IAAe,MAAM,CAAC;GAAE,CAAC;GACpD,WAAW,MAAM;EACnB,EACF,CAAC;CACH;CAEA,oBAA+B,SAAe;EAC5C,IAAI,CAAC,KAAK,WAAW,OAAO;EAC5B,OAAO;GAAE,GAAG;GAAS,UAAU,KAAK;EAAU;CAChD;CAEA,aAAa,SAAqB;EAChC,KAAK,SACH,KAAK,oBAAoB,OAAO,GAChC,KAAK,oBAAoB,CAC3B;CACF;CAEA,WAAW,SAAmB;EAC5B,KAAK,SACH,KAAK,oBAAoB,OAAO,GAChC,KAAK,oBAAoB,CAC3B;CACF;CAEA,WAAW,SAAmB;EAC5B,KAAK,SACH,KAAK,oBAAoB,OAAO,GAChC,KAAK,oBAAoB,CAC3B;CACF;CAEA,QAAQ,OAA6B;EACnC,KAAK,OAAO,OAAO,QAAQ,KAAK;EAEhC,IAAI,MAAM,SAAS,gBAAgB,MAAM,KAAK,WAAW,GACvD,KAAK,OAAO,eAAe,GAAG;CAElC;CAEA,aAAa,UAA6C;EACxD,MAAM,aAAa,IAAI,8BAA8B,KAAK,MAAM;EAChE,WAAW,YAAY;EACvB,OAAO;CACT;CAEA,QAAQ;EACN,KAAK,OAAO,QAAQ,YAAY,MAAM;EACtC,KAAK,OAAO,OAAO,KAAK;EAExB,KAAK,OAAO,kBAAkB;CAChC;AACF;;;;;;;;;;AAWA,SAAgB,sBACd,UACiB;CACjB,MAAM,aAAa,IAAI,8BAA8B;CAErD,MAAM,UAAU,YAAY;EAC1B,IAAI;GACF,MAAM,SAAS,UAAU;EAC3B,SAAS,GAAG;GACV,IAAI,CAAC,WAAW,qBACd,WAAW,QAAQ;IACjB,MAAM;IACN,MAAM,CAAC;IACP,OAAO,OAAO,CAAC;GACjB,CAAC;GAEH,MAAM;EACR,UAAU;GACR,IAAI,CAAC,WAAW,qBACd,WAAW,MAAM;EAErB;CACF;CACA,QAAQ;CAER,OAAO,WAAW,uBAAuB;AAC3C;;;;;;;;AASA,SAAgB,kCAAkC;CAChD,MAAM,EAAE,SAAS,YAAY,qBAA2B;CACxD,IAAI;CAUJ,OAAO,CATQ,uBAAuB,MAAM;EAC1C,aAAa;EAEb,WAA8C,4BAC5C,OACF;EAEA,OAAO;CACT,CACa,GAAG,UAAU;AAC5B;;;;;;;;AASA,SAAgB,8BACd,UACA;CACA,OAAO,gBAAgB,WACrB,sBAAsB,QAAQ,GAC9B,IAAI,kBAAkB,CACxB;AACF"}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "assistant-stream",
-  "version": "0.3.16",
+  "version": "0.3.17",
   "description": "Streaming utilities for AI assistants",
   "keywords": [
     "ai",

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

@@ -0,0 +1,104 @@
+import { describe, it, expect } from "vitest";
+import { createAssistantStreamResponse } from "./assistant-stream";
+import { AssistantStream } from "../AssistantStream";
+import { DataStreamDecoder } from "../serialization/data-stream/DataStream";
+import { AssistantMessageAccumulator } from "../accumulators/assistant-message-accumulator";
+import type { AssistantMessage } from "../utils/types";
+
+const accumulate = async (response: Response): Promise<AssistantMessage> => {
+  const stream = AssistantStream.fromResponse(
+    response,
+    new DataStreamDecoder(),
+  );
+  let last: AssistantMessage | undefined;
+  await stream.pipeThrough(new AssistantMessageAccumulator()).pipeTo(
+    new WritableStream({
+      write(message) {
+        last = message;
+      },
+    }),
+  );
+  return last!;
+};
+
+describe("AssistantStreamController withParentId", () => {
+  it("attaches parentId to text parts across a data-stream round trip", async () => {
+    const response = createAssistantStreamResponse((controller) => {
+      controller.appendText("intro");
+      const group = controller.withParentId("group-1");
+      group.appendSource({
+        type: "source",
+        sourceType: "url",
+        id: "s1",
+        url: "https://example.com",
+        title: "Example",
+      });
+      group.appendText("grouped text");
+    });
+
+    const message = await accumulate(response);
+    const intro = message.parts.find(
+      (p) => p.type === "text" && p.text === "intro",
+    );
+    const grouped = message.parts.find(
+      (p) => p.type === "text" && p.text === "grouped text",
+    );
+    const source = message.parts.find((p) => p.type === "source");
+
+    expect(intro?.parentId).toBeUndefined();
+    expect(grouped?.parentId).toBe("group-1");
+    expect(source?.parentId).toBe("group-1");
+  });
+
+  it("attaches parentId to reasoning parts across a data-stream round trip", async () => {
+    const response = createAssistantStreamResponse((controller) => {
+      controller.appendReasoning("thinking out loud");
+      const group = controller.withParentId("group-1");
+      group.appendReasoning("grouped reasoning");
+    });
+
+    const message = await accumulate(response);
+    const ungrouped = message.parts.find(
+      (p) => p.type === "reasoning" && p.text === "thinking out loud",
+    );
+    const grouped = message.parts.find(
+      (p) => p.type === "reasoning" && p.text === "grouped reasoning",
+    );
+
+    expect(ungrouped?.parentId).toBeUndefined();
+    expect(grouped?.parentId).toBe("group-1");
+  });
+
+  it("opens a new text part when withParentId switches between ids", async () => {
+    const response = createAssistantStreamResponse((controller) => {
+      controller.withParentId("group-1").appendText("first");
+      controller.withParentId("group-2").appendText("second");
+    });
+
+    const message = await accumulate(response);
+    const first = message.parts.find(
+      (p) => p.type === "text" && p.text === "first",
+    );
+    const second = message.parts.find(
+      (p) => p.type === "text" && p.text === "second",
+    );
+
+    expect(first?.parentId).toBe("group-1");
+    expect(second?.parentId).toBe("group-2");
+  });
+
+  it("attaches parentId on addTextPart called directly inside a withParentId scope", async () => {
+    const response = createAssistantStreamResponse((controller) => {
+      const part = controller.withParentId("group-1").addTextPart();
+      part.append("explicit");
+      part.close();
+    });
+
+    const message = await accumulate(response);
+    const text = message.parts.find(
+      (p) => p.type === "text" && p.text === "explicit",
+    );
+
+    expect(text?.parentId).toBe("group-1");
+  });
+});

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

@@ -96,6 +96,7 @@ type AssistantStreamControllerState = {
     | {
         controller: TextStreamController;
         kind: "text" | "reasoning";
+        parentId: string | undefined;
       }
     | undefined;
   contentCounter: Counter;
@@ -150,9 +151,13 @@ class AssistantStreamControllerImpl implements AssistantStreamController {
   }
 
   appendText(textDelta: string) {
-    if (this._state.append?.kind !== "text") {
+    if (
+      this._state.append?.kind !== "text" ||
+      this._state.append.parentId !== this._parentId
+    ) {
       this._state.append = {
         kind: "text",
+        parentId: this._parentId,
         controller: this.addTextPart(),
       };
     }
@@ -160,9 +165,13 @@ class AssistantStreamControllerImpl implements AssistantStreamController {
   }
 
   appendReasoning(textDelta: string) {
-    if (this._state.append?.kind !== "reasoning") {
+    if (
+      this._state.append?.kind !== "reasoning" ||
+      this._state.append.parentId !== this._parentId
+    ) {
       this._state.append = {
         kind: "reasoning",
+        parentId: this._parentId,
         controller: this.addReasoningPart(),
       };
     }
@@ -171,13 +180,13 @@ class AssistantStreamControllerImpl implements AssistantStreamController {
 
   addTextPart() {
     const [stream, controller] = createTextStreamController();
-    this._addPart({ type: "text" }, stream);
+    this._addPart(this._withParentIdOption({ type: "text" }), stream);
     return controller;
   }
 
   addReasoningPart() {
     const [stream, controller] = createTextStreamController();
-    this._addPart({ type: "reasoning" }, stream);
+    this._addPart(this._withParentIdOption({ type: "reasoning" }), stream);
     return controller;
   }