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

Streaming utilities for AI assistants

import { describe, it, expect, vi } from "vitest";
import {
  createAssistantStream,
  createAssistantStreamController,
  type AssistantStreamController,
} from "./assistant-stream";
import { createAssistantStreamResponse } from "./assistant-stream-response";
import { AssistantStream } from "../AssistantStream";
import type { AssistantStreamChunk } from "../AssistantStreamChunk";
import { DataStreamDecoder } from "../serialization/data-stream/DataStream";
import { AssistantMessageAccumulator } from "../accumulators/assistant-message-accumulator";
import type { AssistantMessage } from "../utils/types";
import { toolResultStream } from "../tool/toolResultStream";
import {
  AssistantTransportDecoder,
  AssistantTransportEncoder,
} from "../serialization/assistant-transport/AssistantTransport";

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!;
};

const collectChunks = async (
  stream: AssistantStream,
): Promise<AssistantStreamChunk[]> => {
  const chunks: AssistantStreamChunk[] = [];
  await stream.pipeTo(
    new WritableStream({
      write(chunk) {
        chunks.push(chunk);
      },
    }),
  );
  return chunks;
};

const captureUnhandledRejections = async (
  callback: () => Promise<void>,
): Promise<unknown[]> => {
  const reasons: unknown[] = [];
  const listener = (reason: unknown) => reasons.push(reason);
  process.on("unhandledRejection", listener);
  try {
    await callback();
    await new Promise((resolve) => setTimeout(resolve, 0));
    return reasons;
  } finally {
    process.off("unhandledRejection", listener);
  }
};

describe("tool-call writes after cancellation", () => {
  it("ignores setResponse after the consumer cancels", async () => {
    const unhandledRejections = await captureUnhandledRejections(async () => {
      const [stream, controller] = createAssistantStreamController();
      const toolCall = controller.addToolCallPart({
        toolCallId: "t1",
        toolName: "search",
      });

      const reader = stream.getReader();
      await reader.cancel("consumer stopped");
      await new Promise((resolve) => setTimeout(resolve, 0));

      toolCall.setResponse({ result: "ok" });
      await new Promise((resolve) => setTimeout(resolve, 0));
    });

    expect(unhandledRejections).toEqual([]);
  });

  it("ignores args appends and close after the consumer cancels", async () => {
    const unhandledRejections = await captureUnhandledRejections(async () => {
      const [stream, controller] = createAssistantStreamController();
      const toolCall = controller.addToolCallPart({
        toolCallId: "t1",
        toolName: "search",
      });

      const reader = stream.getReader();
      await reader.cancel("consumer stopped");
      await new Promise((resolve) => setTimeout(resolve, 0));

      toolCall.argsText.append('{"q":1}');
      toolCall.close();
      await new Promise((resolve) => setTimeout(resolve, 0));
    });

    expect(unhandledRejections).toEqual([]);
  });

  it("cancels while a local part waits for its start index", async () => {
    const [stream, controller] = createAssistantStreamController();
    const first = controller.addTextPart();
    first.append("before");
    first.close();
    controller.enqueue({
      type: "step-finish",
      path: [],
      finishReason: "tool-calls",
      usage: { inputTokens: 0, outputTokens: 0 },
      isContinued: false,
    });
    const second = controller.addTextPart();
    second.append("after");

    const reader = stream.getReader();
    await reader.read();
    await reader.cancel();
  });
});

describe("controller close idempotence", () => {
  it("tolerates a second close after the stream drains", async () => {
    const [stream, controller] = createAssistantStreamController();
    controller.appendText("hi");
    controller.close();

    const reader = stream.getReader();
    while (!(await reader.read()).done) {
      // drain
    }

    expect(() => controller.close()).not.toThrow();
  });
});

describe("raw chunk ordering", () => {
  const stepFinish: AssistantStreamChunk = {
    type: "step-finish",
    path: [],
    finishReason: "tool-calls",
    usage: { inputTokens: 0, outputTokens: 0 },
    isContinued: false,
  };

  it("keeps a raw step-finish before a buffered merged part", async () => {
    const merged = new ReadableStream<AssistantStreamChunk>({
      start(inner) {
        inner.enqueue({ type: "part-start", part: { type: "text" }, path: [] });
        inner.enqueue({ type: "text-delta", textDelta: "child", path: [0] });
        inner.enqueue({ type: "part-finish", path: [0] });
        inner.close();
      },
    });

    const chunks = await collectChunks(
      createAssistantStream((controller) => {
        controller.merge(merged);
        controller.enqueue(stepFinish);
      }),
    );

    expect(chunks).toEqual([
      stepFinish,
      { type: "part-start", part: { type: "text" }, path: [] },
      { type: "text-delta", textDelta: "child", path: [0] },
      { type: "part-finish", path: [0] },
    ]);
  });

  it("rejects a merged chunk whose part path has no start and cancels its source", async () => {
    const consoleError = vi
      .spyOn(console, "error")
      .mockImplementation(() => {});
    try {
      const cancelSource = vi.fn();
      const merged = new ReadableStream<AssistantStreamChunk>({
        start(inner) {
          inner.enqueue({ type: "text-delta", textDelta: "orphan", path: [0] });
        },
        cancel: cancelSource,
      });

      await expect(
        collectChunks(
          createAssistantStream((controller) => controller.merge(merged)),
        ),
      ).rejects.toThrow("Path not found");
      await vi.waitFor(() => expect(cancelSource).toHaveBeenCalledOnce());
    } finally {
      consoleError.mockRestore();
    }
  });

  it("emits synchronous raw chunks before an appended part body", async () => {
    const before: AssistantStreamChunk = {
      type: "annotations",
      path: [],
      annotations: ["before"],
    };
    const after: AssistantStreamChunk = {
      type: "annotations",
      path: [],
      annotations: ["after"],
    };

    const chunks = await collectChunks(
      createAssistantStream((controller) => {
        controller.enqueue(before);
        controller.appendText("child");
        controller.enqueue(after);
      }),
    );

    expect(chunks).toEqual([
      before,
      {
        type: "part-start",
        path: [],
        part: { type: "text" },
      },
      after,
      {
        type: "text-delta",
        path: [0],
        textDelta: "child",
      },
      {
        type: "part-finish",
        path: [0],
      },
    ]);
  });

  it("keeps tool args when step-finish follows the closed args", async () => {
    const message = await accumulate(
      createAssistantStreamResponse((controller) => {
        const tool = controller.addToolCallPart({
          toolCallId: "t1",
          toolName: "weather",
        });
        tool.argsText.append('{"city":"Paris"}');
        tool.argsText.close();
        controller.enqueue({
          type: "step-finish",
          path: [],
          finishReason: "tool-calls",
          usage: { inputTokens: 0, outputTokens: 0 },
          isContinued: false,
        });
        tool.close();
      }),
    );

    expect(message.parts).toContainEqual(
      expect.objectContaining({
        type: "tool-call",
        argsText: '{"city":"Paris"}',
      }),
    );
  });

  it("keeps tool args when message-finish follows the closed tool call", async () => {
    const message = await accumulate(
      createAssistantStreamResponse((controller) => {
        const tool = controller.addToolCallPart({
          toolCallId: "t1",
          toolName: "weather",
        });
        tool.argsText.append('{"city":"Paris"}');
        tool.argsText.close();
        tool.close();
        controller.enqueue({
          type: "message-finish",
          path: [],
          finishReason: "tool-calls",
          usage: { inputTokens: 0, outputTokens: 0 },
        });
      }),
    );

    expect(message.parts).toContainEqual(
      expect.objectContaining({
        type: "tool-call",
        argsText: '{"city":"Paris"}',
      }),
    );
  });

  it("keeps a text body ahead of a message-finish enqueued after it", async () => {
    const chunks = await collectChunks(
      createAssistantStream((controller) => {
        controller.appendText("hello");
        controller.enqueue({
          type: "message-finish",
          path: [],
          finishReason: "stop",
          usage: { inputTokens: 0, outputTokens: 0 },
        });
      }),
    );

    expect(chunks.map((c) => c.type)).toEqual([
      "part-start",
      "text-delta",
      "message-finish",
      "part-finish",
    ]);
  });

  it("emits part chunks before the error chunk when a callback throws", async () => {
    const chunks = await collectChunks(
      createAssistantStream((controller) => {
        controller.appendText("hello");
        throw new Error("failed");
      }),
    );

    const errorIndex = chunks.findIndex((chunk) => chunk.type === "error");
    const deltaIndex = chunks.findIndex((chunk) => chunk.type === "text-delta");
    expect(deltaIndex).toBeGreaterThanOrEqual(0);
    expect(errorIndex).toBeGreaterThan(deltaIndex);
  });

  it("starts a part opened after a held step-finish before its body", async () => {
    const run = (controller: AssistantStreamController) => {
      const tool = controller.addToolCallPart({
        toolCallId: "t1",
        toolName: "weather",
      });
      tool.argsText.append('{"city":"Paris"}');
      controller.enqueue(stepFinish);
      controller.appendText("next");
      tool.close();
    };

    const chunks = await collectChunks(createAssistantStream(run));
    let started = 0;
    for (const chunk of chunks) {
      if (chunk.type === "part-start" && chunk.path.length === 0) started++;
      else if (chunk.path.length > 0)
        expect(chunk.path[0]).toBeLessThan(started);
    }

    const message = await accumulate(createAssistantStreamResponse(run));
    expect(message.parts).toContainEqual(
      expect.objectContaining({ type: "text", text: "next" }),
    );
    expect(message.parts).toContainEqual(
      expect.objectContaining({
        type: "tool-call",
        argsText: '{"city":"Paris"}',
      }),
    );
  });

  it("keeps a tool call added after a held step-finish", async () => {
    const message = await accumulate(
      createAssistantStreamResponse((controller) => {
        const first = controller.addToolCallPart({
          toolCallId: "t1",
          toolName: "weather",
        });
        first.argsText.append('{"city":"Paris"}');
        controller.enqueue(stepFinish);
        const second = controller.addToolCallPart({
          toolCallId: "t2",
          toolName: "weather",
          args: { city: "Rome" },
        });
        second.setResponse({ result: "sunny" });
        first.close();
      }),
    );

    expect(message.parts).toContainEqual(
      expect.objectContaining({
        type: "tool-call",
        toolCallId: "t2",
        argsText: '{"city":"Rome"}',
        result: "sunny",
      }),
    );
  });

  it("keeps merged part paths aligned with start order across a microtask", async () => {
    let innerController!: ReadableStreamDefaultController<AssistantStreamChunk>;
    let release!: () => void;
    const ready = new Promise<void>((resolve) => (release = resolve));
    const stream = createAssistantStream(async (controller) => {
      controller.merge(
        new ReadableStream<AssistantStreamChunk>({
          start(inner) {
            innerController = inner;
          },
        }),
      );
      const text = controller.addTextPart();
      controller.enqueue({
        type: "step-finish",
        path: [],
        finishReason: "stop",
        usage: { inputTokens: 2, outputTokens: 0 },
        isContinued: false,
      });
      await ready;
      innerController.enqueue({
        type: "part-start",
        part: { type: "text" },
        path: [],
      });
      const tool = controller.addToolCallPart({
        toolCallId: "t1",
        toolName: "t",
      });
      text.close();
      innerController.enqueue({ type: "part-finish", path: [0] });
      tool.argsText.append('{"a":1}');
      void tool.close();
      innerController.close();
    });
    const reader = stream.getReader();
    const chunks: AssistantStreamChunk[] = [];
    const first = await reader.read();
    chunks.push(first.value!);
    await Promise.resolve();
    release();
    for (;;) {
      const result = await reader.read();
      if (result.done) break;
      chunks.push(result.value);
    }

    const started: string[] = [];
    for (const chunk of chunks) {
      if (chunk.type === "part-start" && chunk.path.length === 0) {
        started.push(chunk.part.type);
      } else if (chunk.path.length > 0) {
        expect(chunk.path[0]).toBeLessThan(started.length);
        if (
          chunk.type === "tool-call-args-text-finish" ||
          (chunk.type === "text-delta" && chunk.textDelta.startsWith("{"))
        ) {
          expect(started[chunk.path[0]!]).toBe("tool-call");
        }
      }
    }
    expect([...started].sort()).toEqual(["text", "text", "tool-call"]);
    expect(
      chunks
        .filter((chunk) => chunk.type === "part-finish")
        .map((chunk) => chunk.path[0])
        .sort(),
    ).toEqual([0, 1, 2]);

    let message: AssistantMessage | undefined;
    await new ReadableStream<AssistantStreamChunk>({
      start(controller) {
        for (const chunk of chunks) controller.enqueue(chunk);
        controller.close();
      },
    })
      .pipeThrough(new AssistantMessageAccumulator())
      .pipeTo(
        new WritableStream({
          write(value: AssistantMessage) {
            message = value;
          },
        }),
      );
    expect(message?.parts).toContainEqual(
      expect.objectContaining({
        type: "tool-call",
        toolCallId: "t1",
        argsText: '{"a":1}',
      }),
    );
  });

  it("emits a step-finish without waiting for an open part to close", async () => {
    const [stream, controller] = createAssistantStreamController();
    const text = controller.addTextPart();
    text.append("a");
    controller.enqueue(stepFinish);

    const reader = stream.getReader();
    const types: string[] = [];
    while (!types.includes("step-finish")) {
      const { value } = await reader.read();
      types.push(value!.type);
    }
    expect(types).toEqual(["part-start", "text-delta", "step-finish"]);

    text.append("b");
    text.close();
    controller.close();
    const rest: string[] = [];
    for (;;) {
      const { done, value } = await reader.read();
      if (done) break;
      rest.push(value.type);
    }
    expect(rest).toEqual(["text-delta", "part-finish"]);
  });

  it("cancels promptly while a step-finish is held", async () => {
    let release!: () => void;
    const released = new Promise<void>((resolve) => {
      release = resolve;
    });
    const stream = createAssistantStream(async (controller) => {
      const text = controller.addTextPart();
      const tool = controller.addToolCallPart({
        toolCallId: "t1",
        toolName: "weather",
      });
      text.append("a");
      tool.argsText.append("{");
      controller.enqueue(stepFinish);
      await released;
      text.append("late");
      tool.close();
      text.close();
    });

    const reader = stream.getReader();
    await reader.read();
    const reasons = await captureUnhandledRejections(() => reader.cancel());
    release();
    expect(reasons).toEqual([]);
  });

  it("does not hold a step-finish for a write the part dropped", async () => {
    const consoleError = vi
      .spyOn(console, "error")
      .mockImplementation(() => {});
    try {
      const [stream, controller] = createAssistantStreamController({
        strict: false,
      });
      const tool = controller.addToolCallPart({
        toolCallId: "t1",
        toolName: "weather",
      });
      tool.argsText.append("{}");
      tool.argsText.close();
      tool.argsText.append("dropped");
      controller.enqueue(stepFinish);

      const reader = stream.getReader();
      const types: string[] = [];
      while (!types.includes("step-finish")) {
        const { value } = await reader.read();
        types.push(value!.type);
      }
      expect(types).toEqual(["part-start", "text-delta", "step-finish"]);

      tool.close();
      controller.close();
      await reader.cancel();
    } finally {
      consoleError.mockRestore();
    }
  });

  it("emits the error chunk while a merged stream stays open", async () => {
    let upstream!: ReadableStreamDefaultController<AssistantStreamChunk>;
    const stream = createAssistantStream((controller) => {
      controller.merge(
        new ReadableStream({
          start(c) {
            upstream = c;
          },
        }),
      );
      controller.appendText("partial");
      throw new Error("failed");
    });

    const reader = stream.getReader();
    const types: string[] = [];
    while (!types.includes("error")) {
      const { value } = await reader.read();
      types.push(value!.type);
    }
    expect(types).toEqual(["part-start", "text-delta", "error"]);

    upstream.close();
    await reader.cancel();
  });
});

describe("createAssistantStream task settlement", () => {
  it.each(["text", "reasoning", "tool-call"] as const)(
    "ignores a background %s writer after the callback fails",
    async (type) => {
      let release!: () => void;
      const gate = new Promise<void>((resolve) => {
        release = resolve;
      });
      let lateWrite!: Promise<void>;
      const chunks = await collectChunks(
        createAssistantStream((controller) => {
          const part =
            type === "text"
              ? controller.addTextPart()
              : type === "reasoning"
                ? controller.addReasoningPart()
                : controller.addToolCallPart("search").argsText;
          lateWrite = gate.then(() => {
            part.append("late");
          });
          throw new Error("provider failed");
        }),
      );

      release();
      await expect(lateWrite).resolves.toBeUndefined();
      expect(chunks.filter((chunk) => chunk.type === "text-delta")).toEqual([]);
    },
  );

  it.each(["text", "reasoning", "tool-call"] as const)(
    "keeps strict writes after an explicit %s close when the callback fails",
    async (type) => {
      let append!: () => void;
      await collectChunks(
        createAssistantStream((controller) => {
          const part =
            type === "text"
              ? controller.addTextPart()
              : type === "reasoning"
                ? controller.addReasoningPart()
                : controller.addToolCallPart("search").argsText;
          part.close();
          append = () => part.append("late");
          throw new Error("provider failed");
        }),
      );

      expect(append).toThrow("Cannot append to a closed TextStreamController");
    },
  );

  it("finishes open text and reasoning parts and cuts off open tool calls when the callback throws", async () => {
    const chunks = await collectChunks(
      createAssistantStream(async (controller) => {
        controller.addTextPart().append("partial");
        controller.addReasoningPart().append("thinking");
        controller.addToolCallPart({ toolCallId: "t1", toolName: "search" });
        throw new Error("provider failed");
      }),
    );

    expect(chunks.filter((c) => c.type === "error")).toHaveLength(1);
    expect(chunks).toContainEqual({
      type: "text-delta",
      path: [0],
      textDelta: "partial",
    });
    expect(
      chunks.filter((c) => c.type === "part-finish").map((c) => c.path),
    ).toEqual([[0], [1]]);
    expect(chunks.map((c) => c.type)).not.toContain(
      "tool-call-args-text-finish",
    );
  });

  it.each([false, true])(
    "keeps the args finish of a tool call whose args completed when the callback throws (async: %s)",
    async (isAsync) => {
      const run: Parameters<typeof createAssistantStream>[0] = (controller) => {
        controller.addToolCallPart({
          toolCallId: "t1",
          toolName: "confirm",
          args: { id: 1 },
        });
        throw new Error("upstream failed");
      };
      const chunks = await collectChunks(
        createAssistantStream(isAsync ? async (c) => run(c) : run),
      );

      expect(chunks.filter((c) => c.type === "error")).toHaveLength(1);
      expect(chunks.map((c) => c.type)).toContain("tool-call-args-text-finish");
    },
  );

  it.each(["appendText", "appendReasoning"] as const)(
    "finishes the part %s opened when the callback throws",
    async (method) => {
      const chunks = await collectChunks(
        createAssistantStream((controller) => {
          controller[method]("partial");
          throw new Error("provider failed");
        }),
      );

      expect(chunks.map((c) => c.type)).toEqual([
        "part-start",
        "text-delta",
        "error",
        "part-finish",
      ]);
    },
  );

  it("does not run a frontend tool whose call was open when the callback threw", async () => {
    const execute = vi.fn(() => "confirmed");
    const response = new Response(
      createAssistantStream((controller) => {
        controller.appendText("partial");
        controller.addToolCallPart({ toolCallId: "t1", toolName: "confirm" });
        controller
          .addToolCallPart({ toolCallId: "t2", toolName: "confirm" })
          .argsText.append('{"orderId":');
        throw new Error("upstream failed");
      }).pipeThrough(new AssistantTransportEncoder()),
    );

    let message: AssistantMessage | undefined;
    await AssistantStream.fromResponse(
      response,
      new AssistantTransportDecoder(),
    )
      .pipeThrough(
        toolResultStream(
          { confirm: { parameters: { type: "object" }, execute } },
          new AbortController().signal,
          async () => {},
        ),
      )
      .pipeThrough(new AssistantMessageAccumulator())
      .pipeTo(
        new WritableStream({
          write(m) {
            message = m;
          },
        }),
      );

    expect(execute).not.toHaveBeenCalled();
    expect(message?.status).toMatchObject({
      type: "incomplete",
      reason: "error",
    });
    expect(message?.parts).toMatchObject([
      { type: "text", text: "partial", status: { type: "complete" } },
      { toolCallId: "t1", state: "partial-call" },
      { toolCallId: "t2", state: "partial-call", argsText: '{"orderId":' },
    ]);
  });

  it.each([true, false])(
    "preserves merged output when the outer callback throws (finished: %s)",
    async (finished) => {
      const [merged, inner] = createAssistantStreamController();
      inner.appendText("full answer");
      if (finished) inner.close();

      const pending = collectChunks(
        createAssistantStream((controller) => {
          controller.merge(merged);
          throw new Error("outer failed");
        }),
      );
      if (!finished) {
        await Promise.resolve();
        inner.appendText(" continued");
        inner.close();
      }

      const chunks = await pending;
      expect(
        chunks
          .filter((chunk) => chunk.type === "text-delta")
          .map((chunk) => chunk.textDelta)
          .join(""),
      ).toBe(finished ? "full answer" : "full answer continued");
      expect(chunks.filter((chunk) => chunk.type === "error")).toHaveLength(1);
      expect(
        chunks.filter((chunk) => chunk.type === "part-finish"),
      ).toHaveLength(1);
    },
  );

  it("stops tracking an input once it finishes", async () => {
    let tracked: Set<unknown> | undefined;
    await collectChunks(
      createAssistantStream(async (controller) => {
        tracked = (
          controller as unknown as { _state: { openInputs: Set<unknown> } }
        )._state.openInputs;
        controller.appendText("a");
        controller.addReasoningPart().close();
        controller.addToolCallPart("search").setResponse({ result: 1 });
      }),
    );

    expect(tracked?.size).toBe(0);
  });

  it("closes the outer stream when a merged stream throws with an open part", async () => {
    const chunks = await collectChunks(
      createAssistantStream((controller) => {
        controller.merge(
          createAssistantStream(async (inner) => {
            inner.addToolCallPart("search");
            throw new Error("inner failed");
          }),
        );
      }),
    );

    expect(chunks.map((c) => c.type)).toEqual(["part-start", "error"]);
  });

  it("emits callback failures without leaking an unhandled rejection", async () => {
    let chunks: AssistantStreamChunk[] = [];
    const unhandledRejections = await captureUnhandledRejections(async () => {
      chunks = await collectChunks(
        createAssistantStream(async () => {
          throw new Error("provider failed");
        }),
      );
    });

    expect(chunks).toEqual([
      {
        type: "error",
        path: [],
        error: "Error: provider failed",
      },
    ]);
    expect(unhandledRejections).toEqual([]);
  });

  it("does not settle the stream again after cancellation", async () => {
    let finishCallback!: () => void;
    const callbackPending = new Promise<void>((resolve) => {
      finishCallback = resolve;
    });

    const unhandledRejections = await captureUnhandledRejections(async () => {
      const reader = createAssistantStream(() => callbackPending).getReader();
      await reader.cancel("consumer stopped");
      finishCallback();
      await callbackPending;
    });

    expect(unhandledRejections).toEqual([]);
  });

  it("waits for merged source cleanup during cancellation", async () => {
    let finishCleanup = () => {};
    const cleanup = new Promise<void>((resolve) => {
      finishCleanup = resolve;
    });
    const cancelSource = vi.fn(() => cleanup);
    const stream = createAssistantStream((controller) => {
      controller.merge(new ReadableStream({ cancel: cancelSource }));
      return new Promise(() => {});
    });

    let cancelSettled = false;
    const cancel = stream.cancel().then(() => {
      cancelSettled = true;
    });

    await vi.waitFor(() => expect(cancelSource).toHaveBeenCalledOnce());
    expect(cancelSettled).toBe(false);

    finishCleanup();
    await expect(cancel).resolves.toBeUndefined();
    expect(cancelSettled).toBe(true);
  });

  it("reports callback failures after the controller is explicitly closed", async () => {
    const error = new Error("cleanup failed");
    const consoleError = vi
      .spyOn(console, "error")
      .mockImplementation(() => {});

    try {
      const chunks = await collectChunks(
        createAssistantStream((controller) => {
          controller.close();
          throw error;
        }),
      );

      expect(chunks).toEqual([]);
      expect(consoleError).toHaveBeenCalledOnce();
      expect(consoleError).toHaveBeenCalledWith(error);
    } finally {
      consoleError.mockRestore();
    }
  });

  it("does not report callback failures caused after cancellation", async () => {
    let failCallback!: (error: Error) => void;
    const callbackPending = new Promise<void>((_, reject) => {
      failCallback = reject;
    });
    const consoleError = vi
      .spyOn(console, "error")
      .mockImplementation(() => {});

    try {
      const unhandledRejections = await captureUnhandledRejections(async () => {
        const reader = createAssistantStream(() => callbackPending).getReader();
        await reader.cancel("consumer stopped");
        failCallback(new Error("provider stopped"));
        await callbackPending.catch(() => undefined);
      });

      expect(unhandledRejections).toEqual([]);
      expect(consoleError).not.toHaveBeenCalled();
    } finally {
      consoleError.mockRestore();
    }
  });

  it("does not settle the outer stream again after a merged stream errors", async () => {
    let finishCallback!: () => void;
    const callbackPending = new Promise<void>((resolve) => {
      finishCallback = resolve;
    });
    const streamError = new Error("merged stream failed");
    const consoleError = vi
      .spyOn(console, "error")
      .mockImplementation(() => {});

    try {
      const unhandledRejections = await captureUnhandledRejections(async () => {
        const stream = createAssistantStream(async (controller) => {
          controller.merge(
            new ReadableStream({
              start(streamController) {
                streamController.error(streamError);
              },
            }),
          );
          await callbackPending;
        });

        await expect(collectChunks(stream)).rejects.toBe(streamError);
        finishCallback();
        await callbackPending;
      });

      expect(unhandledRejections).toEqual([]);
      expect(consoleError).toHaveBeenCalledWith(streamError);
    } finally {
      consoleError.mockRestore();
    }
  });

  it("surfaces a merged stream error while sibling cleanup continues", async () => {
    let finishCleanup = () => {};
    const cleanup = new Promise<void>((resolve) => {
      finishCleanup = resolve;
    });
    const cancelSource = vi.fn(() => cleanup);
    const streamError = new Error("merged stream failed");
    const consoleError = vi
      .spyOn(console, "error")
      .mockImplementation(() => {});

    try {
      const stream = createAssistantStream((controller) => {
        controller.merge(
          new ReadableStream({
            start(streamController) {
              streamController.error(streamError);
            },
          }),
        );
        controller.merge(new ReadableStream({ cancel: cancelSource }));
        return new Promise(() => {});
      });
      const readResult = stream
        .getReader()
        .read()
        .then(
          (value) => ({ value }),
          (error: unknown) => ({ error }),
        );

      await vi.waitFor(() => expect(cancelSource).toHaveBeenCalledOnce());
      await expect(readResult).resolves.toEqual({ error: streamError });
    } finally {
      finishCleanup();
      await cleanup;
      consoleError.mockRestore();
    }
  });

  it("emits an error chunk when merging a locked stream", async () => {
    const source = new ReadableStream();
    const sourceReader = source.getReader();

    try {
      const chunks = await collectChunks(
        createAssistantStream((controller) => {
          controller.merge(source);
        }),
      );

      expect(chunks).toEqual([
        {
          type: "error",
          path: [],
          error:
            "TypeError: Cannot merge a stream that is already locked to a reader.",
        },
      ]);
    } finally {
      sourceReader.releaseLock();
    }
  });
});

describe("addToolCallPart with an immediate response", () => {
  it("completes the stream without an explicit tool close", async () => {
    const chunks = await collectChunks(
      createAssistantStream((controller) => {
        controller.addToolCallPart({
          toolName: "search",
          response: { result: "done" },
        });
      }),
    );

    expect(chunks.map((c) => c.type)).toContain("result");
    expect(chunks.at(-1)?.type).toBe("part-finish");
  });

  it("completes the stream when args accompany the response", async () => {
    const chunks = await collectChunks(
      createAssistantStream((controller) => {
        controller.addToolCallPart({
          toolName: "search",
          args: { query: "x" },
          response: { result: "done" },
        });
      }),
    );

    const deltas = chunks.filter((c) => c.type === "text-delta");
    expect(deltas.map((c) => c.textDelta).join("")).toBe('{"query":"x"}');
    expect(chunks.filter((c) => c.type === "result")).toHaveLength(1);
    expect(chunks.at(-1)?.type).toBe("part-finish");
  });

  it("emits the args ahead of the result", async () => {
    const chunks = await collectChunks(
      createAssistantStream((controller) => {
        const tool = controller.addToolCallPart({ toolName: "search" });
        tool.argsText.append('{"query":"x"}');
        tool.setResponse({ result: "done" });
      }),
    );

    expect(
      chunks.filter((c) => c.path.length === 1).map((c) => c.type),
    ).toEqual([
      "text-delta",
      "result",
      "tool-call-args-text-finish",
      "part-finish",
    ]);
  });

  it("emits the result ahead of the args finish when the args close in the same tick", async () => {
    const chunks = await collectChunks(
      createAssistantStream((controller) => {
        const tool = controller.addToolCallPart({ toolName: "search" });
        tool.argsText.append('{"query":"x"}');
        tool.argsText.close();
        tool.setResponse({ result: "done" });
      }),
    );

    expect(
      chunks.filter((c) => c.path.length === 1).map((c) => c.type),
    ).toEqual([
      "text-delta",
      "result",
      "tool-call-args-text-finish",
      "part-finish",
    ]);
  });

  it("rejects args passed alongside argsText instead of concatenating them", () => {
    const [, controller] = createAssistantStreamController();

    expect(() =>
      controller.addToolCallPart({
        toolName: "search",
        argsText: '{"query":"x"}',
        args: { query: "y" },
        response: { result: "done" },
      }),
    ).toThrow("Cannot append to a closed TextStreamController");
  });

  it("keeps only argsText when args are also passed in lenient mode", async () => {
    const consoleError = vi
      .spyOn(console, "error")
      .mockImplementation(() => {});
    try {
      const chunks = await collectChunks(
        createAssistantStream(
          (controller) => {
            controller.addToolCallPart({
              toolName: "search",
              argsText: '{"query":"x"}',
              args: { query: "y" },
              response: { result: "done" },
            });
          },
          { strict: false },
        ),
      );

      const deltas = chunks.filter((c) => c.type === "text-delta");
      expect(deltas.map((c) => c.textDelta).join("")).toBe('{"query":"x"}');
      expect(chunks.filter((c) => c.type === "result")).toHaveLength(1);
      expect(consoleError).toHaveBeenCalledTimes(1);
    } finally {
      consoleError.mockRestore();
    }
  });

  it("keeps the args across a data-stream round trip", async () => {
    const message = await accumulate(
      createAssistantStreamResponse((controller) => {
        controller.addToolCallPart({
          toolCallId: "t1",
          toolName: "search",
          args: { query: "x" },
          response: { result: "done" },
        });
      }),
    );

    expect(message.parts[0]).toMatchObject({
      type: "tool-call",
      argsText: '{"query":"x"}',
      result: "done",
    });
  });

  it("keeps working when the caller also closes explicitly", async () => {
    const chunks = await collectChunks(
      createAssistantStream((controller) => {
        const tool = controller.addToolCallPart({
          toolName: "search",
          response: { result: "done" },
        });
        tool.close();
      }),
    );

    expect(chunks.filter((c) => c.type === "result")).toHaveLength(1);
    expect(chunks.filter((c) => c.type === "part-finish")).toHaveLength(1);
    expect(chunks.at(-1)?.type).toBe("part-finish");
  });
});

describe("tool-call finish ordering", () => {
  it("keeps close order across uneven argument backlogs", async () => {
    const chunks = await collectChunks(
      createAssistantStream((controller) => {
        const first = controller.addToolCallPart("first");
        const second = controller.addToolCallPart("second");

        first.argsText.append('{"value":"');
        for (let index = 0; index < 20; index++) {
          first.argsText.append(String(index % 10));
        }
        first.argsText.append('"}');
        first.close();
        second.close();
      }),
    );

    expect(
      chunks
        .filter((chunk) => chunk.type === "part-finish")
        .map((chunk) => chunk.path),
    ).toEqual([[0], [1]]);
  });

  it("delivers finishes in tool close order", async () => {
    const chunks = await collectChunks(
      createAssistantStream((controller) => {
        const first = controller.addToolCallPart("first");
        const second = controller.addToolCallPart("second");
        second.close();
        first.close();
      }),
    );

    expect(
      chunks
        .filter((chunk) => chunk.type === "part-finish")
        .map((chunk) => chunk.path),
    ).toEqual([[1], [0]]);
  });

  it("does not block a closed tool behind an earlier open tool", async () => {
    const [stream, controller] = createAssistantStreamController();
    controller.addToolCallPart("first");
    controller.addToolCallPart("second").close();
    const reader = stream.getReader();
    let finish: AssistantStreamChunk | undefined;

    while (finish?.type !== "part-finish") {
      const next = await reader.read();
      if (next.done) break;
      finish = next.value;
    }

    expect(finish).toEqual({ type: "part-finish", path: [1] });
    await reader.cancel();
  });
});

describe("AssistantStreamController withParentId", () => {
  it("preserves a reasoning summary from addReasoningPart", async () => {
    const stream = createAssistantStream((controller) => {
      const part = controller.addReasoningPart({
        unstable_summary: "Planning",
      });
      part.append("thinking");
      part.close();
    });

    const chunks = await collectChunks(stream);

    expect(chunks).toContainEqual({
      type: "part-start",
      path: [],
      part: {
        type: "reasoning",
        unstable_summary: "Planning",
      },
    });
  });

  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");
  });
});