assistant-stream/src/core/accumulators/assistant-message-accumulator.test.ts

Streaming utilities for AI assistants

import { describe, it, expect, vi, onTestFinished } from "vitest";
import {
  AssistantMessageAccumulator,
  createInitialMessage,
} from "./assistant-message-accumulator";
import type { AssistantStreamChunk } from "../AssistantStreamChunk";
import type { AssistantMessage } from "../utils/types";

const collectStream = async (
  chunks: AssistantStreamChunk[],
): Promise<AssistantMessage[]> => {
  const accumulator = new AssistantMessageAccumulator();
  const messages: AssistantMessage[] = [];

  const source = new ReadableStream<AssistantStreamChunk>({
    start(controller) {
      for (const chunk of chunks) {
        controller.enqueue(chunk);
      }
      controller.close();
    },
  });

  await source.pipeThrough(accumulator).pipeTo(
    new WritableStream({
      write(message) {
        messages.push(message);
      },
    }),
  );

  return messages;
};

describe("AssistantMessageAccumulator reasoning summaries", () => {
  it("copies reasoning summaries from part-start chunks", async () => {
    const messages = await collectStream([
      {
        type: "part-start",
        path: [0],
        part: { type: "reasoning", unstable_summary: "Planning" },
      },
      { type: "text-delta", path: [0], textDelta: "thinking" },
      { type: "part-finish", path: [0] },
    ]);

    expect(messages.at(-1)?.parts).toEqual([
      expect.objectContaining({
        type: "reasoning",
        text: "thinking",
        unstable_summary: "Planning",
      }),
    ]);
  });

  it("preserves an empty reasoning summary", async () => {
    const messages = await collectStream([
      {
        type: "part-start",
        path: [0],
        part: { type: "reasoning", unstable_summary: "" },
      },
    ]);

    expect(messages.at(-1)?.parts[0]).toEqual(
      expect.objectContaining({ type: "reasoning", unstable_summary: "" }),
    );
  });
});

describe("AssistantMessageAccumulator document sources", () => {
  it("defaults a missing document source media type", async () => {
    const messages = await collectStream([
      {
        type: "part-start",
        path: [0],
        part: {
          type: "source",
          sourceType: "document",
          id: "d1",
          title: "Q3 report",
        },
      } as unknown as AssistantStreamChunk,
    ]);

    expect(messages.at(-1)?.parts).toEqual([
      expect.objectContaining({
        type: "source",
        sourceType: "document",
        id: "d1",
        title: "Q3 report",
        mediaType: "application/octet-stream",
      }),
    ]);
  });
});

describe("AssistantMessageAccumulator tool argument status", () => {
  it("marks arguments complete before the result and preserves a settled result", async () => {
    const messages = await collectStream([
      {
        type: "part-start",
        path: [],
        part: { type: "tool-call", toolCallId: "tc-1", toolName: "search" },
      },
      { type: "text-delta", path: [0], textDelta: '{"q":"test"}' },
      { type: "tool-call-args-text-finish", path: [0] },
      { type: "result", path: [0], result: "found", isError: false },
      { type: "tool-call-args-text-finish", path: [0] },
    ]);

    expect(messages[1]?.parts[0]).toMatchObject({
      state: "partial-call",
      status: { type: "running", isArgsComplete: false },
    });
    expect(messages[2]?.parts[0]).toMatchObject({
      state: "call",
      args: { q: "test" },
      status: { type: "running", isArgsComplete: true },
    });
    expect(messages.at(-1)?.parts[0]).toMatchObject({
      state: "result",
      result: "found",
      status: { type: "complete", reason: "stop" },
    });
  });

  it("preserves completed status when argument text finishes after the part", async () => {
    const messages = await collectStream([
      {
        type: "part-start",
        path: [],
        part: { type: "tool-call", toolCallId: "tc-1", toolName: "search" },
      },
      { type: "text-delta", path: [0], textDelta: '{"q":"test"}' },
      { type: "part-finish", path: [0] },
      { type: "tool-call-args-text-finish", path: [0] },
    ]);

    expect(messages[3]?.parts[0]).toMatchObject({
      state: "call",
      status: { type: "complete", reason: "unknown" },
    });
  });

  it("keeps a preliminary result running until a final result arrives", async () => {
    const messages = await collectStream([
      {
        type: "part-start",
        path: [],
        part: { type: "tool-call", toolCallId: "tc-1", toolName: "search" },
      },
      { type: "text-delta", path: [0], textDelta: "{}" },
      { type: "tool-call-args-text-finish", path: [0] },
      {
        type: "result",
        path: [0],
        result: "interim",
        isError: false,
        isPreliminary: true,
        modelContent: [{ type: "text", text: "Search in progress" }],
      },
      { type: "part-finish", path: [0] },
      { type: "result", path: [0], result: "done", isError: false },
      { type: "part-finish", path: [0] },
    ]);

    const preliminaryMessage = messages.find(
      (message) =>
        message.parts[0]?.type === "tool-call" &&
        message.parts[0].isPreliminary,
    );
    expect(preliminaryMessage?.parts[0]).toMatchObject({
      state: "call",
      result: "interim",
      isPreliminary: true,
      modelContent: [{ type: "text", text: "Search in progress" }],
      status: { type: "running", isArgsComplete: true },
    });
    expect(messages.at(-1)?.parts[0]).toMatchObject({
      state: "result",
      result: "done",
      status: { type: "complete" },
    });
    expect(messages.at(-1)?.parts[0]).not.toHaveProperty("isPreliminary");
    expect(messages.at(-1)?.parts[0]).not.toHaveProperty("modelContent");
  });

  it("clears model content when a later preliminary result omits it", async () => {
    const messages = await collectStream([
      {
        type: "part-start",
        path: [],
        part: { type: "tool-call", toolCallId: "tc-1", toolName: "search" },
      },
      { type: "tool-call-args-text-finish", path: [0] },
      {
        type: "result",
        path: [0],
        result: "searching",
        isError: false,
        isPreliminary: true,
        modelContent: [{ type: "text", text: "Search in progress" }],
      },
      {
        type: "result",
        path: [0],
        result: "one match so far",
        isError: false,
        isPreliminary: true,
      },
    ]);

    expect(messages.at(-1)?.parts[0]).toMatchObject({
      state: "call",
      result: "one match so far",
      isPreliminary: true,
      status: { type: "running", isArgsComplete: true },
    });
    expect(messages.at(-1)?.parts[0]).not.toHaveProperty("modelContent");
  });

  it("clears stale artifact and messages when a later result omits them", async () => {
    const messages = await collectStream([
      {
        type: "part-start",
        path: [],
        part: { type: "tool-call", toolCallId: "tc-1", toolName: "search" },
      },
      { type: "tool-call-args-text-finish", path: [0] },
      {
        type: "result",
        path: [0],
        result: "searching",
        isError: false,
        isPreliminary: true,
        artifact: { reportId: 7 },
        messages: [
          {
            role: "assistant",
            content: [{ type: "text", text: "in progress" }],
          },
        ],
        modelContent: [{ type: "text", text: "Search in progress" }],
      },
      {
        type: "result",
        path: [0],
        result: "done",
        isError: false,
      },
    ]);

    expect(messages.at(-1)?.parts[0]).toMatchObject({
      state: "result",
      result: "done",
      status: { type: "complete", reason: "stop" },
    });
    expect(messages.at(-1)?.parts[0]).not.toHaveProperty("modelContent");
    expect(messages.at(-1)?.parts[0]).not.toHaveProperty("artifact");
    expect(messages.at(-1)?.parts[0]).not.toHaveProperty("messages");
  });

  it("ignores a preliminary result after the final one", async () => {
    const messages = await collectStream([
      {
        type: "part-start",
        path: [],
        part: { type: "tool-call", toolCallId: "tc-1", toolName: "search" },
      },
      { type: "tool-call-args-text-finish", path: [0] },
      { type: "result", path: [0], result: "done", isError: false },
      {
        type: "result",
        path: [0],
        result: "late",
        isError: false,
        isPreliminary: true,
      },
    ]);

    expect(messages.at(-1)?.parts[0]).toMatchObject({
      state: "result",
      result: "done",
      status: { type: "complete", reason: "stop" },
    });
    expect(messages.at(-1)?.parts[0]).not.toHaveProperty("isPreliminary");
  });

  it("keeps a preliminary result pending when the stream ends", async () => {
    const messages = await collectStream([
      {
        type: "part-start",
        path: [],
        part: { type: "tool-call", toolCallId: "tc-1", toolName: "search" },
      },
      { type: "tool-call-args-text-finish", path: [0] },
      {
        type: "result",
        path: [0],
        result: "interim",
        isError: false,
        isPreliminary: true,
      },
      { type: "part-finish", path: [0] },
    ]);

    expect(messages.at(-1)).toMatchObject({
      status: { type: "requires-action", reason: "tool-calls" },
      parts: [
        {
          state: "call",
          result: "interim",
          isPreliminary: true,
          status: { type: "running", isArgsComplete: true },
        },
      ],
    });
  });
});

describe("AssistantMessageAccumulator timing", () => {
  it("should include timing on message-finish", async () => {
    const chunks: AssistantStreamChunk[] = [
      { type: "step-start", path: [], messageId: "msg-1" },
      { type: "part-start", path: [0], part: { type: "text" } },
      { type: "text-delta", path: [0], textDelta: "Hello " },
      { type: "text-delta", path: [0], textDelta: "world" },
      { type: "part-finish", path: [0] },
      {
        type: "step-finish",
        path: [],
        finishReason: "stop",
        usage: { inputTokens: 10, outputTokens: 5 },
        isContinued: false,
      },
      {
        type: "message-finish",
        path: [],
        finishReason: "stop",
        usage: { inputTokens: 10, outputTokens: 5 },
      },
    ];

    const messages = await collectStream(chunks);
    const last = messages.at(-1)!;

    expect(last.metadata.timing).toBeDefined();
    const timing = last.metadata.timing!;
    expect(timing.streamStartTime).toBeTypeOf("number");
    expect(timing.totalChunks).toBe(7);
    expect(timing.toolCallCount).toBe(0);
    expect(timing.tokenCount).toBe(5);
    expect(timing.firstTokenTime).toBeTypeOf("number");
    expect(timing.totalStreamTime).toBeTypeOf("number");
    expect(timing.totalStreamTime).toBeGreaterThanOrEqual(0);
  });

  it.each([
    { stepTokens: [], finalTokens: 20, expected: 20 },
    { stepTokens: [7], finalTokens: 20, expected: 20 },
    { stepTokens: [3, 4], finalTokens: 0, expected: 7 },
    { stepTokens: [3, 4], finalTokens: undefined, expected: 7 },
  ])(
    "uses $expected tokens for steps $stepTokens and final usage $finalTokens",
    async ({ stepTokens, finalTokens, expected }) => {
      const chunks: AssistantStreamChunk[] = [
        { type: "part-start", path: [], part: { type: "text" } },
        { type: "text-delta", path: [0], textDelta: "test" },
        { type: "part-finish", path: [0] },
      ];
      for (const outputTokens of stepTokens) {
        chunks.push({
          type: "step-finish",
          path: [],
          finishReason: "stop",
          usage: { inputTokens: 5, outputTokens },
          isContinued: false,
        });
      }
      if (finalTokens !== undefined) {
        chunks.push({
          type: "message-finish",
          path: [],
          finishReason: "stop",
          usage: { inputTokens: 5, outputTokens: finalTokens },
        });
      }
      chunks.push({ type: "annotations", path: [], annotations: ["done"] });

      const messages = await collectStream(chunks);
      const timed = messages.filter((m) => m.metadata.timing !== undefined);

      expect(timed.length).toBeGreaterThan(0);
      for (const message of timed) {
        expect(message.metadata.timing?.tokenCount).toBe(expected);
      }
    },
  );

  it("falls back to the step total when message-finish omits usage", async () => {
    const messages = await collectStream([
      { type: "part-start", path: [], part: { type: "text" } },
      { type: "text-delta", path: [0], textDelta: "test" },
      { type: "part-finish", path: [0] },
      {
        type: "step-finish",
        path: [],
        finishReason: "stop",
        usage: { inputTokens: 5, outputTokens: 9 },
        isContinued: false,
      },
      {
        type: "message-finish",
        path: [],
        finishReason: "stop",
      } as unknown as AssistantStreamChunk,
    ]);

    expect(messages.at(-1)?.metadata.timing?.tokenCount).toBe(9);
  });

  it("should track tool calls in timing", async () => {
    const chunks: AssistantStreamChunk[] = [
      {
        type: "part-start",
        path: [0],
        part: {
          type: "tool-call",
          toolCallId: "tc-1",
          toolName: "search",
        },
      },
      { type: "text-delta", path: [0], textDelta: '{"q":"test"}' },
      { type: "tool-call-args-text-finish", path: [0] },
      {
        type: "result",
        path: [0],
        result: "found",
        isError: false,
      },
      { type: "part-finish", path: [0] },
      {
        type: "message-finish",
        path: [],
        finishReason: "stop",
        usage: { inputTokens: 0, outputTokens: 0 },
      },
    ];

    const messages = await collectStream(chunks);
    const last = messages.at(-1)!;

    expect(last.metadata.timing).toBeDefined();
    expect(last.metadata.timing!.toolCallCount).toBe(1);
  });

  it("should record per-tool-call timing on the part", async () => {
    const before = Date.now();
    const chunks: AssistantStreamChunk[] = [
      {
        type: "part-start",
        path: [0],
        part: {
          type: "tool-call",
          toolCallId: "tc-1",
          toolName: "search",
        },
      },
      { type: "text-delta", path: [0], textDelta: '{"q":"test"}' },
      { type: "tool-call-args-text-finish", path: [0] },
      {
        type: "result",
        path: [0],
        result: "found",
        isError: false,
      },
      { type: "part-finish", path: [0] },
      {
        type: "message-finish",
        path: [],
        finishReason: "stop",
        usage: { inputTokens: 0, outputTokens: 0 },
      },
    ];

    const messages = await collectStream(chunks);
    const after = Date.now();

    const runningPart = messages
      .find((m) =>
        m.parts.some((p) => p.type === "tool-call" && p.state !== "result"),
      )!
      .parts.find((p) => p.type === "tool-call")!;
    expect(runningPart.timing).toBeDefined();
    expect(runningPart.timing!.startedAt).toBeGreaterThanOrEqual(before);
    expect(runningPart.timing!.completedAt).toBeUndefined();

    const settledPart = messages
      .at(-1)!
      .parts.find((p) => p.type === "tool-call")!;
    expect(settledPart.timing).toBeDefined();
    expect(settledPart.timing!.completedAt).toBeGreaterThanOrEqual(
      settledPart.timing!.startedAt,
    );
    expect(settledPart.timing!.completedAt).toBeLessThanOrEqual(after);
  });

  it("should record reasoning timing on the part", async () => {
    const before = Date.now();
    const chunks: AssistantStreamChunk[] = [
      { type: "part-start", path: [0], part: { type: "reasoning" } },
      { type: "text-delta", path: [0], textDelta: "thinking" },
      { type: "part-finish", path: [0] },
      { type: "part-start", path: [1], part: { type: "text" } },
      { type: "text-delta", path: [1], textDelta: "answer" },
      { type: "part-finish", path: [1] },
      {
        type: "message-finish",
        path: [],
        finishReason: "stop",
        usage: { inputTokens: 0, outputTokens: 0 },
      },
    ];

    const messages = await collectStream(chunks);
    const after = Date.now();

    const streamingPart = messages
      .find((m) =>
        m.parts.some(
          (p) => p.type === "reasoning" && p.status.type === "running",
        ),
      )!
      .parts.find((p) => p.type === "reasoning")!;
    expect(streamingPart.timing!.startedAt).toBeGreaterThanOrEqual(before);
    expect(streamingPart.timing!.completedAt).toBeUndefined();

    const last = messages.at(-1)!;
    const settledPart = last.parts.find((p) => p.type === "reasoning")!;
    expect(settledPart.timing!.completedAt).toBeGreaterThanOrEqual(
      settledPart.timing!.startedAt,
    );
    expect(settledPart.timing!.completedAt).toBeLessThanOrEqual(after);
    expect(last.parts.find((p) => p.type === "text")).not.toHaveProperty(
      "timing",
    );
  });

  it("should include timing on flush when stream closes without message-finish", async () => {
    const chunks: AssistantStreamChunk[] = [
      { type: "part-start", path: [0], part: { type: "text" } },
      { type: "text-delta", path: [0], textDelta: "partial" },
    ];

    const messages = await collectStream(chunks);
    const last = messages.at(-1)!;

    // flush should have produced a message with timing
    expect(last.metadata.timing).toBeDefined();
    expect(last.metadata.timing!.totalChunks).toBe(2);
  });

  it("should estimate tokens from text when no usage provided", async () => {
    const chunks: AssistantStreamChunk[] = [
      { type: "part-start", path: [0], part: { type: "text" } },
      {
        type: "text-delta",
        path: [0],
        textDelta: "1234567890123456789012345678901234567890",
      }, // 40 chars
      { type: "part-finish", path: [0] },
      {
        type: "message-finish",
        path: [],
        finishReason: "stop",
        usage: { inputTokens: 0, outputTokens: 0 },
      },
    ];

    const messages = await collectStream(chunks);
    const last = messages.at(-1)!;

    expect(last.metadata.timing!.tokenCount).toBe(10); // 40/4
  });

  it("should record firstTokenTime on text-delta", async () => {
    const chunks: AssistantStreamChunk[] = [
      {
        type: "part-start",
        path: [0],
        part: {
          type: "tool-call",
          toolCallId: "tc-1",
          toolName: "search",
        },
      },
      { type: "text-delta", path: [0], textDelta: "{}" },
      { type: "tool-call-args-text-finish", path: [0] },
      {
        type: "result",
        path: [0],
        result: "ok",
        isError: false,
      },
      { type: "part-finish", path: [0] },
      {
        type: "message-finish",
        path: [],
        finishReason: "stop",
        usage: { inputTokens: 0, outputTokens: 0 },
      },
    ];

    const messages = await collectStream(chunks);
    const last = messages.at(-1)!;

    // text-delta on tool-call args still records first token
    expect(last.metadata.timing).toBeDefined();
    expect(last.metadata.timing!.firstTokenTime).toBeTypeOf("number");
  });

  it("does not record firstTokenTime when a text-delta is dropped for an out-of-range path", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    const messages = await collectStream([
      { type: "part-start", path: [0], part: { type: "text" } },
      { type: "text-delta", path: [5], textDelta: "x" },
    ]);
    const last = messages.at(-1)!;

    expect(last.metadata.timing).toBeDefined();
    expect(last.metadata.timing!.firstTokenTime).toBeUndefined();
    warn.mockRestore();
  });

  it("does not record firstTokenTime when a text-delta is dropped for a nested path", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    const messages = await collectStream([
      { type: "part-start", path: [0], part: { type: "text" } },
      { type: "text-delta", path: [0, 1], textDelta: "x" },
    ]);
    const last = messages.at(-1)!;

    expect(last.metadata.timing).toBeDefined();
    expect(last.metadata.timing!.firstTokenTime).toBeUndefined();
    warn.mockRestore();
  });

  it("does not record firstTokenTime when a text-delta is dropped for a wrong part type", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    const messages = await collectStream([
      {
        type: "part-start",
        path: [0],
        part: {
          type: "source",
          sourceType: "url",
          id: "s1",
          url: "https://example.com",
        },
      },
      { type: "text-delta", path: [0], textDelta: "x" },
    ]);
    const last = messages.at(-1)!;

    expect(last.metadata.timing).toBeDefined();
    expect(last.metadata.timing!.firstTokenTime).toBeUndefined();
    warn.mockRestore();
  });

  it("records firstTokenTime on the first applied text-delta, not a preceding dropped one", async () => {
    vi.useFakeTimers();
    vi.setSystemTime(1000);

    const accumulator = new AssistantMessageAccumulator();
    const writer = accumulator.writable.getWriter();
    const messages: AssistantMessage[] = [];
    const reading = (async () => {
      const reader = accumulator.readable.getReader();
      for (;;) {
        const { done, value } = await reader.read();
        if (done) break;
        messages.push(value);
      }
    })();

    await writer.write({
      type: "part-start",
      path: [0],
      part: { type: "text" },
    });
    vi.setSystemTime(2000);
    await writer.write({ type: "text-delta", path: [5], textDelta: "x" });
    vi.setSystemTime(3000);
    await writer.write({ type: "text-delta", path: [0], textDelta: "hi" });
    await writer.close();
    await reading;
    vi.useRealTimers();

    const last = messages.at(-1)!;
    expect(last.metadata.timing).toBeDefined();
    expect(last.metadata.timing!.firstTokenTime).toBe(2000);
  });
});

describe("AssistantMessageAccumulator error chunks", () => {
  it("defaults missing code to unknown", async () => {
    const messages = await collectStream([
      { type: "error", path: [], error: "stream failed" },
    ]);
    const last = messages.at(-1)!;

    expect(last.status).toEqual({
      type: "incomplete",
      reason: "error",
      error: { code: "unknown", message: "stream failed" },
    });
  });

  it("defaults a missing error message", async () => {
    const messages = await collectStream([
      { type: "error", path: [] } as unknown as AssistantStreamChunk,
    ]);
    const last = messages.at(-1)!;

    expect(last.status).toEqual({
      type: "incomplete",
      reason: "error",
      error: { code: "unknown", message: "unknown error" },
    });
  });

  it("carries code and severity through unchanged", async () => {
    const messages = await collectStream([
      {
        type: "error",
        path: [],
        error: "rate limited",
        code: "provider",
        severity: "warning",
      },
    ]);
    const last = messages.at(-1)!;

    expect(last.status).toEqual({
      type: "incomplete",
      reason: "error",
      error: {
        code: "provider",
        message: "rate limited",
        severity: "warning",
      },
    });
  });

  it("drops out-of-taxonomy severity at ingestion", async () => {
    const messages = await collectStream([
      {
        type: "error",
        path: [],
        error: "boom",
        code: "unknown",
        severity: "fatal",
      } as unknown as AssistantStreamChunk,
    ]);
    const last = messages.at(-1)!;

    expect(last.status).toEqual({
      type: "incomplete",
      reason: "error",
      error: {
        code: "unknown",
        message: "boom",
      },
    });
    expect(last.status).toMatchObject({
      error: expect.not.objectContaining({ severity: expect.anything() }),
    });
  });
});

describe("AssistantMessageAccumulator part path bounds", () => {
  it("drops a text-delta whose path is out of range with a warning", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    const messages = await collectStream([
      { type: "part-start", path: [0], part: { type: "text" } },
      { type: "text-delta", path: [5], textDelta: "x" },
      { type: "text-delta", path: [0], textDelta: "hi" },
    ]);
    const last = messages.at(-1)!;

    expect(last.parts).toHaveLength(1);
    expect(last.parts[0]).toMatchObject({ type: "text", text: "hi" });
    expect(warn).toHaveBeenCalledWith(
      "Dropped text-delta chunk: no part at path [5]",
    );
    warn.mockRestore();
  });

  it("drops an out-of-range part-finish instead of fabricating a part", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    const messages = await collectStream([
      { type: "part-start", path: [0], part: { type: "text" } },
      { type: "part-finish", path: [3] },
    ]);
    const last = messages.at(-1)!;

    expect(last.parts).toHaveLength(1);
    expect(last.parts[0]!.type).toBe("text");
    warn.mockRestore();
  });

  it("drops part chunks when no parts exist", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    const messages = await collectStream([
      { type: "text-delta", path: [0], textDelta: "x" },
      { type: "part-start", path: [0], part: { type: "text" } },
      { type: "text-delta", path: [0], textDelta: "ok" },
    ]);
    const last = messages.at(-1)!;

    expect(last.parts).toHaveLength(1);
    expect(last.parts[0]).toMatchObject({ type: "text", text: "ok" });
    warn.mockRestore();
  });

  it("drops chunks with nested paths", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    const messages = await collectStream([
      { type: "part-start", path: [0], part: { type: "text" } },
      { type: "text-delta", path: [0, 1], textDelta: "x" },
      { type: "text-delta", path: [0], textDelta: "ok" },
    ]);
    const last = messages.at(-1)!;

    expect(last.parts).toHaveLength(1);
    expect(last.parts[0]).toMatchObject({ type: "text", text: "ok" });
    warn.mockRestore();
  });

  it("drops an out-of-range result without touching existing tool calls", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    const messages = await collectStream([
      {
        type: "part-start",
        path: [0],
        part: { type: "tool-call", toolCallId: "t1", toolName: "f" },
      },
      { type: "text-delta", path: [0], textDelta: "{}" },
      { type: "result", path: [7], result: { ok: true }, isError: false },
      { type: "result", path: [0], result: { ok: true }, isError: false },
    ]);
    const last = messages.at(-1)!;

    expect(last.parts).toHaveLength(1);
    expect(last.parts[0]).toMatchObject({
      type: "tool-call",
      state: "result",
      result: { ok: true },
    });
    warn.mockRestore();
  });
});

describe("AssistantMessageAccumulator wrong part type chunks", () => {
  const sourcePart = {
    type: "source",
    sourceType: "url",
    id: "s1",
    url: "https://example.com",
  } as const;

  it("drops a text-delta addressed to a source part", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    const messages = await collectStream([
      { type: "part-start", path: [0], part: sourcePart },
      { type: "text-delta", path: [0], textDelta: "x" },
    ]);
    const last = messages.at(-1)!;

    expect(last.parts[0]).toMatchObject({ type: "source", id: "s1" });
    expect(warn).toHaveBeenCalledWith(
      "Dropped text-delta chunk: part is neither text nor tool-call",
    );
    warn.mockRestore();
  });

  it("drops a result addressed to a text part", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    const messages = await collectStream([
      { type: "part-start", path: [0], part: { type: "text" } },
      { type: "text-delta", path: [0], textDelta: "hi" },
      { type: "result", path: [0], result: { ok: true }, isError: false },
    ]);
    const last = messages.at(-1)!;

    expect(last.parts[0]).toMatchObject({ type: "text", text: "hi" });
    expect(warn).toHaveBeenCalledWith(
      "Dropped result chunk: part is not a tool-call",
    );
    warn.mockRestore();
  });

  it("drops an args-text-finish addressed to a text part", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    const messages = await collectStream([
      { type: "part-start", path: [0], part: { type: "text" } },
      { type: "tool-call-args-text-finish", path: [0] },
    ]);
    const last = messages.at(-1)!;

    expect(last.parts[0]).toMatchObject({ type: "text" });
    expect(warn).toHaveBeenCalledWith(
      "Dropped tool-call-args-text-finish chunk: part is not a tool-call",
    );
    warn.mockRestore();
  });

  it("inserts a placeholder for an unsupported part-start so later part indices stay aligned", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    const messages = await collectStream([
      {
        type: "part-start",
        path: [0],
        part: { type: "bogus" },
      } as unknown as AssistantStreamChunk,
      { type: "part-start", path: [1], part: { type: "text" } },
      { type: "text-delta", path: [1], textDelta: "ok" },
      { type: "part-finish", path: [0] },
    ]);
    const last = messages.at(-1)!;

    expect(last.parts).toHaveLength(2);
    expect(last.parts[0]).toMatchObject({
      type: "reasoning",
      text: "",
      status: { type: "complete", reason: "unknown" },
    });
    expect(last.parts[1]).toMatchObject({ type: "text", text: "ok" });
    expect(warn).toHaveBeenCalledWith(
      "Unsupported part-start type bogus: inserting an empty reasoning part to preserve part indices",
    );
    warn.mockRestore();
  });
});

describe("AssistantMessageAccumulator warn dedup", () => {
  it("warns once per drop class per instance", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    await collectStream([
      { type: "part-start", path: [0], part: { type: "text" } },
      { type: "text-delta", path: [5], textDelta: "a" },
      { type: "text-delta", path: [5], textDelta: "b" },
      { type: "text-delta", path: [6], textDelta: "c" },
      { type: "part-finish", path: [9] },
    ]);

    expect(warn).toHaveBeenCalledTimes(2);
    expect(warn).toHaveBeenCalledWith(
      "Dropped text-delta chunk: no part at path [5]",
    );
    expect(warn).toHaveBeenCalledWith(
      "Dropped part-finish chunk: no part at path [9]",
    );
    warn.mockRestore();
  });

  it("resets dedup per accumulator instance", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    const stream = [
      { type: "part-start", path: [0], part: { type: "text" } },
      { type: "text-delta", path: [5], textDelta: "x" },
    ] satisfies AssistantStreamChunk[];
    await collectStream(stream);
    await collectStream(stream);

    expect(warn).toHaveBeenCalledTimes(2);
    warn.mockRestore();
  });

  it("dedupes repeated wrong-part-type drops", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    await collectStream([
      {
        type: "part-start",
        path: [0],
        part: { type: "source", sourceType: "url", id: "s", url: "https://x" },
      },
      { type: "text-delta", path: [0], textDelta: "a" },
      { type: "text-delta", path: [0], textDelta: "b" },
    ]);

    expect(warn).toHaveBeenCalledTimes(1);
    warn.mockRestore();
  });
});

describe("AssistantMessageAccumulator warn dedup key independence", () => {
  it("keeps drop classes as independent keys within one stream", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    await collectStream([
      {
        type: "part-start",
        path: [0],
        part: { type: "source", sourceType: "url", id: "s", url: "https://x" },
      },
      { type: "text-delta", path: [0], textDelta: "wrong-part" },
      { type: "text-delta", path: [7], textDelta: "no-part" },
    ]);

    expect(warn).toHaveBeenCalledTimes(2);
    warn.mockRestore();
  });

  it("keys unsupported part-starts per type", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    const messages = await collectStream([
      {
        type: "part-start",
        path: [0],
        part: { type: "bogus" },
      } as unknown as AssistantStreamChunk,
      {
        type: "part-start",
        path: [1],
        part: { type: "bogus" },
      } as unknown as AssistantStreamChunk,
      {
        type: "part-start",
        path: [2],
        part: { type: "video" },
      } as unknown as AssistantStreamChunk,
    ]);
    const last = messages.at(-1)!;

    expect(last.parts).toHaveLength(3);
    expect(warn).toHaveBeenCalledTimes(2);
    warn.mockRestore();
  });

  it("caps the number of warned keys per instance", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    await collectStream(
      Array.from(
        { length: 20 },
        (_, i) =>
          ({
            type: "part-start",
            path: [i],
            part: { type: `bogus-${i}` },
          }) as unknown as AssistantStreamChunk,
      ),
    );

    expect(warn).toHaveBeenCalledTimes(16);
    warn.mockRestore();
  });
});

describe("AssistantMessageAccumulator content alias", () => {
  const contentIsAccessor = (message: AssistantMessage) =>
    Object.getOwnPropertyDescriptor(message, "content")?.get !== undefined;

  it("aliases content to parts on every message it emits", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    try {
      const messages = await collectStream([
        { type: "part-start", path: [0], part: { type: "text" } },
        { type: "text-delta", path: [0], textDelta: "hi" },
        {
          type: "part-start",
          path: [1],
          part: { type: "tool-call", toolCallId: "tc-1", toolName: "search" },
        },
        {
          type: "part-start",
          path: [2],
          part: {
            type: "source",
            sourceType: "url",
            id: "s-1",
            url: "https://example.com",
          },
        },
        {
          type: "part-start",
          path: [3],
          part: { type: "file", mimeType: "image/png", data: "AAAA" },
        },
        {
          type: "part-start",
          path: [4],
          part: { type: "data", name: "chart", data: { a: 1 } },
        },
        {
          type: "part-start",
          path: [5],
          part: { type: "totally-unknown" },
        } as unknown as AssistantStreamChunk,
      ]);

      expect(messages.at(-1)?.parts).toHaveLength(6);
      for (const message of messages) {
        expect(message.content).toBe(message.parts);
      }
    } finally {
      warn.mockRestore();
    }
  });

  it("aliases content to parts on the initial message", () => {
    const message = createInitialMessage();
    expect(contentIsAccessor(message)).toBe(true);
    expect(message.content).toBe(message.parts);
    // Published as unstable_createInitialMessage, so key order is observable
    // through Object.keys and JSON serialization.
    expect(Object.keys(message)).toEqual([
      "role",
      "status",
      "parts",
      "content",
      "metadata",
    ]);
  });

  it("keeps content pointing at parts across status and metadata updates", async () => {
    const messages = await collectStream([
      { type: "part-start", path: [0], part: { type: "text" } },
      { type: "text-delta", path: [0], textDelta: "hi" },
      { type: "annotations", path: [], annotations: ["a"] },
      { type: "step-start", path: [], messageId: "m-1" },
      {
        type: "message-finish",
        path: [],
        finishReason: "stop",
        usage: { inputTokens: 1, outputTokens: 1 },
      },
    ]);

    for (const message of messages) {
      expect(message.content).toBe(message.parts);
    }
  });
});

describe("AssistantMessageAccumulator tool arguments", () => {
  it("keeps nested values and partial arguments across small deltas", async () => {
    const input = '{"query":"pizza","filters":{"limit":2}}';
    const messages = await collectStream([
      {
        type: "part-start",
        path: [0],
        part: { type: "tool-call", toolCallId: "call-1", toolName: "search" },
      },
      ...[...input].map((text) => ({
        type: "text-delta" as const,
        path: [0],
        textDelta: text,
      })),
      { type: "tool-call-args-text-finish", path: [0] },
    ]);

    const partial = messages.slice(0, -1).find((message) => {
      const part = message.parts[0];
      return (
        part?.type === "tool-call" &&
        part.args.query === "pizza" &&
        part.argsText !== input
      );
    });

    expect(partial?.parts[0]).toMatchObject({
      type: "tool-call",
      args: { query: "pizza" },
    });

    expect(messages.at(-1)?.parts[0]).toMatchObject({
      type: "tool-call",
      argsText: input,
      args: { query: "pizza", filters: { limit: 2 } },
    });
  });

  it("exposes tool arguments as plain cloneable objects", async () => {
    const messages = await collectStream([
      {
        type: "part-start",
        path: [0],
        part: { type: "tool-call", toolCallId: "call-1", toolName: "search" },
      },
      { type: "text-delta", path: [0], textDelta: '{"query":"pizza",' },
      { type: "text-delta", path: [0], textDelta: '"limit":2}' },
      { type: "tool-call-args-text-finish", path: [0] },
    ]);

    const part = messages.at(-1)?.parts[0];
    expect(part?.type).toBe("tool-call");
    if (part?.type !== "tool-call") throw new Error("Expected a tool call");

    expect(Object.getPrototypeOf(part.args)).toBe(Object.prototype);
    expect(structuredClone(part.args)).toEqual({ query: "pizza", limit: 2 });

    const mutableArgs = part.args as Record<string, unknown>;
    mutableArgs.query = "updated";
    expect(mutableArgs.query).toBe("updated");
  });
});

describe("AssistantMessageAccumulator tool-call approvals", () => {
  it("sets an approval on a tool call and drops one aimed at another part", async () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    onTestFinished(() => warn.mockRestore());
    const messages = await collectStream([
      {
        type: "part-start",
        path: [],
        part: { type: "tool-call", toolCallId: "call-1", toolName: "deploy" },
      },
      { type: "part-start", path: [], part: { type: "text" } },
      { type: "tool-call-approval", path: [0], approval: { id: "a1" } },
      { type: "tool-call-approval", path: [1], approval: { id: "a2" } },
    ]);

    expect(messages.at(-1)!.parts).toEqual([
      expect.objectContaining({ type: "tool-call", approval: { id: "a1" } }),
      expect.not.objectContaining({ approval: expect.anything() }),
    ]);
    expect(warn).toHaveBeenCalledWith(
      "Dropped tool-call-approval chunk: part is not a tool-call",
    );
  });
});