Building a Streaming AI Chat UI That Stays Smooth: SSE, ReadableStream, Incremental Markdown and a Stop Button

export type ChatEvent =
  | { type: "status"; value: "thinking" | "writing" }
  | { type: "text-delta"; text: string }
  | { type: "tool-call"; id: string; name: string; input: unknown }
  | { type: "tool-result"; id: string; output: unknown; isError: boolean }
  | { type: "error"; message: string; retryable: boolean }
  | { type: "done"; reason: "end_turn" | "max_tokens" | "aborted" | "other" };

const encoder = new TextEncoder();

/** Encodes one event in the text/event-stream format. */
export function encodeEvent(event: ChatEvent): Uint8Array {
  // JSON.stringify never emits raw newlines, so a single data: line is enough
  return encoder.encode(`event: ${event.type}\ndata: ${JSON.stringify(event)}\n\n`);
}

/** A comment line — the client ignores it, but it keeps proxies from cutting the connection. */
export const HEARTBEAT = encoder.encode(": ping\n\n");

import Anthropic from "@anthropic-ai/sdk";
import { encodeEvent, HEARTBEAT, type ChatEvent } from "@/lib/chat-protocol";
import { SUPPORT_TOOLS, runSupportTool } from "@/lib/support-tools";
import { requireSession } from "@/lib/auth";

export const maxDuration = 60;

const client = new Anthropic(); // reads ANTHROPIC_API_KEY — only exists on the server
const MODEL = process.env.CHAT_MODEL;
const MAX_TOOL_ROUNDS = 3;

type ClientMessage = { role: "user" | "assistant"; text: string };

function isRetryable(err: unknown): boolean {
  if (!(err instanceof Anthropic.APIError)) return false;
  // status undefined: connection error or an error event mid-stream (e.g. overloaded)
  return err.status === undefined || err.status === 429 || err.status >= 500;
}

export async function POST(req: Request) {
  const session = await requireSession(req); // throws 401 if not signed in
  if (!MODEL) return Response.json({ error: "CHAT_MODEL is not set" }, { status: 500 });

  const { messages } = (await req.json()) as { messages: ClientMessage[] };
  // Only accept user/assistant roles and text — never let the client send a "system" role
  const history: Anthropic.MessageParam[] = messages
    .filter((m) => m.role === "user" || m.role === "assistant")
    .map((m) => ({ role: m.role, content: String(m.text).slice(0, 8_000) }));

  // Two ways to know the client has left: req.signal, or the response stream being cancelled
  const upstream = new AbortController();
  req.signal.addEventListener("abort", () => upstream.abort(), { once: true });
  let heartbeat: ReturnType<typeof setInterval> | undefined;

  const body = new ReadableStream<Uint8Array>({
    async start(controller) {
      const send = (event: ChatEvent) => {
        if (!upstream.signal.aborted) controller.enqueue(encodeEvent(event));
      };
      heartbeat = setInterval(() => {
        if (!upstream.signal.aborted) controller.enqueue(HEARTBEAT);
      }, 15_000);

      try {
        for (let round = 0; round < MAX_TOOL_ROUNDS; round++) {
          const stream = client.messages.stream(
            { model: MODEL, max_tokens: 4096, system: SYSTEM_PROMPT, tools: SUPPORT_TOOLS, messages: history },
            { signal: upstream.signal },
          );

          for await (const event of stream) {
            if (event.type === "content_block_start") {
              if (event.content_block.type === "thinking") send({ type: "status", value: "thinking" });
              if (event.content_block.type === "text") send({ type: "status", value: "writing" });
            } else if (event.type === "content_block_delta" && event.delta.type === "text_delta") {
              send({ type: "text-delta", text: event.delta.text });
            }
          }

          const message = await stream.finalMessage();
          const calls = message.content.filter((b): b is Anthropic.ToolUseBlock => b.type === "tool_use");
          if (message.stop_reason !== "tool_use" || calls.length === 0) {
            const r = message.stop_reason;
            send({ type: "done", reason: r === "end_turn" || r === "max_tokens" ? r : "other" });
            return;
          }

          history.push({ role: "assistant", content: message.content });
          const results: Anthropic.ToolResultBlockParam[] = [];
          for (const call of calls) {
            send({ type: "tool-call", id: call.id, name: call.name, input: call.input });
            // The tool runs with the session's permissions; never trust the model's input about "who owns the order"
            const { output, isError } = await runSupportTool(call, session, upstream.signal);
            send({ type: "tool-result", id: call.id, output, isError });
            results.push({ type: "tool_result", tool_use_id: call.id, content: JSON.stringify(output), is_error: isError });
          }
          history.push({ role: "user", content: results });
        }
        send({ type: "error", message: "Too many tool steps", retryable: false });
      } catch (err) {
        if (upstream.signal.aborted) return; // the user hit Stop — not an error
        send({ type: "error", message: "The assistant is unavailable", retryable: isRetryable(err) });
      } finally {
        clearInterval(heartbeat);
        try { controller.close(); } catch { /* the stream was already cancelled */ }
      }
    },
    cancel() {
      clearInterval(heartbeat);
      upstream.abort();
    },
  });

  return new Response(body, {
    headers: {
      "Content-Type": "text/event-stream; charset=utf-8",
      "Cache-Control": "no-cache, no-transform",
      "X-Accel-Buffering": "no",
    },
  });
}

const SYSTEM_PROMPT =
  "You are the support assistant for our store. Use lookup_order for order questions. Answer in the user's language.";

{
  "functions": {
    "app/api/**/*": { "supportsCancellation": true }
  }
}

import type { ChatEvent } from "./chat-protocol";

export type SSEMessage = { event: string; data: string };

export function createSSEParser(): TransformStream<string, SSEMessage> {
  let buffer = "";
  let pendingCR = false; // the previous chunk ended with \r — the \n may be in the next chunk
  let eventName = "";
  let data: string[] = [];

  const processLine = (line: string, out: TransformStreamDefaultController<SSEMessage>) => {
    if (line === "") {
      if (data.length > 0) out.enqueue({ event: eventName || "message", data: data.join("\n") });
      eventName = "";
      data = [];
      return;
    }
    if (line.startsWith(":")) return; // comment, e.g. a heartbeat
    const colon = line.indexOf(":");
    const field = colon === -1 ? line : line.slice(0, colon);
    let value = colon === -1 ? "" : line.slice(colon + 1);
    if (value.startsWith(" ")) value = value.slice(1);
    if (field === "event") eventName = value;
    else if (field === "data") data.push(value);
  };

  return new TransformStream({
    transform(chunk, out) {
      if (pendingCR && chunk.startsWith("\n")) chunk = chunk.slice(1); // CRLF split in half
      pendingCR = chunk.endsWith("\r");
      buffer += chunk;
      const lines = buffer.split(/\r\n|\r|\n/);
      buffer = lines.pop() ?? ""; // an incomplete line, wait for the next chunk
      for (const line of lines) processLine(line, out);
    },
    // flush: an unfinished event when the stream closes is discarded, per the spec
  });
}

export async function* readChatEvents(response: Response): AsyncGenerator<ChatEvent> {
  if (!response.body) throw new Error("Response has no body");
  const reader = response.body
    .pipeThrough(new TextDecoderStream())
    .pipeThrough(createSSEParser())
    .getReader();
  try {
    while (true) {
      const { value, done } = await reader.read();
      if (done) return;
      yield JSON.parse(value.data) as ChatEvent;
    }
  } finally {
    reader.cancel().catch(() => {}); // the user broke out early → cancel the whole pipeline
  }
}

import type { ChatEvent } from "./chat-protocol";

export type Part =
  | { type: "text"; text: string }
  | { type: "tool-call"; id: string; name: string; input: unknown; state: "running" | "done" | "error"; output?: unknown }
  | { type: "error"; message: string; retryable: boolean };

export type Message = {
  id: string;
  role: "user" | "assistant";
  status: "submitted" | "thinking" | "streaming" | "done" | "stopped" | "error";
  parts: Part[];
};

export type Action =
  | { type: "submit"; user: Message; assistantId: string }
  | { type: "events"; id: string; events: ChatEvent[] } // a whole frame of events at once
  | { type: "stopped"; id: string }
  | { type: "failed"; id: string; message: string; retryable: boolean }
  | { type: "retry"; id: string };

function applyEvent(msg: Message, e: ChatEvent): Message {
  switch (e.type) {
    case "status":
      return { ...msg, status: e.value === "thinking" ? "thinking" : "streaming" };
    case "text-delta": {
      const parts = msg.parts.slice();
      const last = parts.at(-1);
      // Append to the last text part; after a tool call, open a new text part
      if (last?.type === "text") parts[parts.length - 1] = { ...last, text: last.text + e.text };
      else parts.push({ type: "text", text: e.text });
      return { ...msg, status: "streaming", parts };
    }
    case "tool-call":
      return { ...msg, parts: [...msg.parts, { type: "tool-call", id: e.id, name: e.name, input: e.input, state: "running" }] };
    case "tool-result":
      return {
        ...msg,
        parts: msg.parts.map((p) =>
          p.type === "tool-call" && p.id === e.id ? { ...p, state: e.isError ? "error" : "done", output: e.output } : p,
        ),
      };
    case "error":
      return { ...msg, status: "error", parts: [...msg.parts, { type: "error", message: e.message, retryable: e.retryable }] };
    case "done":
      return { ...msg, status: "done" };
  }
}

export function chatReducer(state: Message[], action: Action): Message[] {
  switch (action.type) {
    case "submit":
      return [...state, action.user, { id: action.assistantId, role: "assistant", status: "submitted", parts: [] }];
    case "events":
      return state.map((m) => (m.id === action.id ? action.events.reduce(applyEvent, m) : m));
    case "stopped":
      return state.map((m) => (m.id === action.id && m.status !== "done" ? { ...m, status: "stopped" } : m));
    case "failed":
      return state.map((m) =>
        m.id === action.id
          ? { ...m, status: "error", parts: [...m.parts, { type: "error", message: action.message, retryable: action.retryable }] }
          : m,
      );
    case "retry":
      return state.map((m) => (m.id === action.id ? { ...m, status: "submitted", parts: [] } : m));
  }
}

/** The only thing sent to the server: role + text. */
export function toWire(messages: Message[]) {
  return messages
    .map((m) => ({ role: m.role, text: m.parts.flatMap((p) => (p.type === "text" ? [p.text] : [])).join("") }))
    .filter((m) => m.text.length > 0);
}

"use client";
import { useCallback, useEffect, useReducer, useRef } from "react";
import type { ChatEvent } from "@/lib/chat-protocol";
import { chatReducer, toWire, type Message } from "@/lib/chat-state";
import { readChatEvents } from "@/lib/sse";

class HttpError extends Error {
  constructor(readonly status: number) {
    super(`HTTP ${status}`);
  }
}
const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));
const ACTIVE = new Set<Message["status"]>(["submitted", "thinking", "streaming"]);

export function useChatStream(endpoint = "/api/chat") {
  const [messages, dispatch] = useReducer(chatReducer, []);
  const messagesRef = useRef(messages);
  useEffect(() => {
    messagesRef.current = messages; // send/retry read through the ref → stable identity for memo
  });
  const abortRef = useRef<AbortController | null>(null);
  const pending = useRef<{ id: string; events: ChatEvent[] } | null>(null);
  const frame = useRef<number | null>(null);

  const flush = useCallback(() => {
    if (frame.current !== null) cancelAnimationFrame(frame.current);
    frame.current = null;
    const batch = pending.current;
    pending.current = null;
    if (batch) dispatch({ type: "events", id: batch.id, events: batch.events });
  }, []);

  const enqueue = useCallback(
    (id: string, event: ChatEvent) => {
      (pending.current ??= { id, events: [] }).events.push(event);
      frame.current ??= requestAnimationFrame(flush); // at most one dispatch per frame
    },
    [flush],
  );

  const run = useCallback(
    async (assistantId: string, history: Message[]) => {
      const controller = new AbortController();
      abortRef.current = controller;
      let received = 0;
      let finished = false;

      for (let attempt = 0; ; attempt++) {
        try {
          const res = await fetch(endpoint, {
            method: "POST",
            headers: { "Content-Type": "application/json" },
            body: JSON.stringify({ messages: toWire(history) }),
            signal: controller.signal,
          });
          if (!res.ok) throw new HttpError(res.status);
          for await (const event of readChatEvents(res)) {
            received++;
            if (event.type === "done" || event.type === "error") finished = true;
            enqueue(assistantId, event);
          }
          // A stream that closes quietly without "done" = cut off midway (proxy timeout, deploy...)
          if (!finished) throw new Error("Stream ended unexpectedly");
          break;
        } catch (err) {
          if (controller.signal.aborted) break;
          const retryable = !(err instanceof HttpError) || err.status === 429 || err.status >= 500;
          // Only auto-retry when the user hasn't seen any text yet — avoids repeated text
          if (received === 0 && retryable && attempt < 2) {
            await sleep(500 * 2 ** attempt);
            continue;
          }
          flush();
          dispatch({ type: "failed", id: assistantId, message: "Connection lost", retryable });
          break;
        }
      }
      flush(); // when the tab is hidden rAF doesn't run — flush synchronously at the end
      if (controller.signal.aborted) dispatch({ type: "stopped", id: assistantId });
      if (abortRef.current === controller) abortRef.current = null;
    },
    [endpoint, enqueue, flush],
  );

  const send = useCallback(
    (text: string) => {
      abortRef.current?.abort(); // only one answer at a time
      const user: Message = { id: crypto.randomUUID(), role: "user", status: "done", parts: [{ type: "text", text }] };
      const assistantId = crypto.randomUUID();
      dispatch({ type: "submit", user, assistantId });
      void run(assistantId, [...messagesRef.current, user]);
    },
    [run],
  );

  const retry = useCallback(
    (assistantId: string) => {
      const all = messagesRef.current;
      const index = all.findIndex((m) => m.id === assistantId);
      if (index === -1) return;
      dispatch({ type: "retry", id: assistantId });
      void run(assistantId, all.slice(0, index));
    },
    [run],
  );

  const stop = useCallback(() => abortRef.current?.abort(), []);
  const last = messages.at(-1);
  const isStreaming = last?.role === "assistant" && ACTIVE.has(last.status);

  return { messages, send, stop, retry, isStreaming };
}

import { marked } from "marked";

export type BlockState = { source: string; stable: string[]; stableLength: number; tail: string };
export const EMPTY_BLOCKS: BlockState = { source: "", stable: [], stableLength: 0, tail: "" };

export function splitBlocks(prev: BlockState, text: string): BlockState {
  // The text no longer continues the already-stable part (retry, edited message) → start over
  const base = text.startsWith(prev.source.slice(0, prev.stableLength)) ? prev : EMPTY_BLOCKS;
  const tokens = marked.lexer(text.slice(base.stableLength));

  let last = tokens.length - 1; // the last content token = the running tail
  while (last >= 0 && tokens[last].type === "space") last--;

  const stable = base.stable.slice();
  let consumed = 0;
  for (let i = 0; i < last; i++) {
    if (tokens[i].type !== "space") stable.push(tokens[i].raw);
    consumed += tokens[i].raw.length;
  }
  const stableLength = base.stableLength + consumed;
  return { source: text, stable, stableLength, tail: text.slice(stableLength) };
}

const FENCE = /^ {0,3}(`{3,}|~{3,})/gm;

/** Make the tail "look finished" — for rendering only, the data is not modified. */
export function prepareTail(tail: string): { markdown: string; inCode: boolean } {
  if ((tail.match(FENCE)?.length ?? 0) % 2 === 1) {
    return { markdown: tail + "\n" + "`".repeat(3), inCode: true }; // unclosed fence
  }
  let md = tail.replace(/!?\[([^\]]*)\]\([^)]*$/, "$1"); // half-written link/image → keep just the text
  if ((md.match(/\*\*/g)?.length ?? 0) % 2 === 1) md += "**"; // unclosed bold
  return { markdown: md, inCode: false };
}

"use client";
import { memo, useState } from "react";
import Markdown, { type Components } from "react-markdown";
import remarkGfm from "remark-gfm";
import rehypeHighlight from "rehype-highlight";
import { EMPTY_BLOCKS, prepareTail, splitBlocks } from "@/lib/markdown-blocks";

const REMARK = [remarkGfm];
const REHYPE = [rehypeHighlight];

function externalHost(href: string): string | null {
  try {
    return new URL(href).host; // a relative URL throws → treat as internal
  } catch {
    return null;
  }
}

const components: Components = {
  a: ({ href, children }) => (
    <a href={href} target="_blank" rel="noopener noreferrer nofollow">
      {children}
      {href && externalHost(href) && <span className="link-host"> ({externalHost(href)})</span>}
    </a>
  ),
  // Don't auto-load model-generated images: an image URL can carry conversation data out
  img: ({ alt }) => <span className="img-placeholder">[{alt || "image"}]</span>,
};

const Block = memo(function Block({ source, live }: { source: string; live: boolean }) {
  return (
    <Markdown remarkPlugins={REMARK} rehypePlugins={live ? [] : REHYPE} components={components} skipHtml>
      {source}
    </Markdown>
  );
});

export function StreamingMarkdown({ text, streaming }: { text: string; streaming: boolean }) {
  const [blocks, setBlocks] = useState(() => splitBlocks(EMPTY_BLOCKS, text));
  let current = blocks;
  if (blocks.source !== text) {
    current = splitBlocks(blocks, text);
    setBlocks(current); // React's "store information from previous renders" pattern
  }
  const tail = streaming ? prepareTail(current.tail).markdown : current.tail;

  return (
    <div className="prose">
      {current.stable.map((source, i) => (
        <Block key={i} source={source} live={false} />
      ))}
      {tail && <Block source={tail} live={streaming} />}
    </div>
  );
}

"use client";
import { useCallback, useEffect, useRef, useState } from "react";

const THRESHOLD = 48; // px

export function useStickToBottom<T extends HTMLElement>() {
  const scrollRef = useRef<T>(null);
  const contentRef = useRef<HTMLDivElement>(null);
  const pinned = useRef(true);
  const [showJump, setShowJump] = useState(false);

  useEffect(() => {
    const scroller = scrollRef.current;
    const content = contentRef.current;
    if (!scroller || !content) return;
    let lastTop = scroller.scrollTop;

    const onScroll = () => {
      const distance = scroller.scrollHeight - scroller.scrollTop - scroller.clientHeight;
      if (scroller.scrollTop < lastTop && distance > THRESHOLD) pinned.current = false; // the user scrolled up
      if (distance <= THRESHOLD) pinned.current = true;
      lastTop = scroller.scrollTop;
      setShowJump(!pinned.current);
    };
    // Catch every height change: new text, code blocks, tool cards, image placeholders
    const observer = new ResizeObserver(() => {
      if (pinned.current) scroller.scrollTop = scroller.scrollHeight;
    });

    observer.observe(content);
    scroller.addEventListener("scroll", onScroll, { passive: true });
    return () => {
      observer.disconnect();
      scroller.removeEventListener("scroll", onScroll);
    };
  }, []);

  const jumpToBottom = useCallback(() => {
    pinned.current = true;
    scrollRef.current?.scrollTo({ top: scrollRef.current.scrollHeight, behavior: "smooth" });
  }, []);

  return { scrollRef, contentRef, showJump, jumpToBottom };
}

import type { Message } from "@/lib/chat-state";

const STATUS_TEXT: Partial<Record<Message["status"], string>> = {
  thinking: "The assistant is thinking",
  streaming: "The assistant is answering",
  stopped: "The answer was stopped",
  error: "The answer ran into an error",
};

export function ChatAnnouncer({ message }: { message?: Message }) {
  let text = "";
  if (message?.role === "assistant") {
    text =
      message.status === "done"
        ? message.parts
            .flatMap((p) => (p.type === "text" ? [p.text] : []))
            .join("")
            .replace(/[*_`#>|]/g, "") // strip markdown symbols so they aren't read aloud
        : (STATUS_TEXT[message.status] ?? "");
  }
  return (
    <div aria-live="polite" aria-atomic="true" className="sr-only">
      {text}
    </div>
  );
}

import type { Part } from "@/lib/chat-state";

type ToolCall = Extract<Part, { type: "tool-call" }>;
type Order = { id: string; status: "processing" | "shipping" | "delivered"; eta?: string };

const ORDER_STATUS: Record<Order["status"], string> = {
  processing: "Processing",
  shipping: "Shipping",
  delivered: "Delivered",
};

export function ToolPart({ part }: { part: ToolCall }) {
  if (part.state === "running") {
    const orderId = (part.input as { orderId?: string }).orderId;
    return (
      <div className="tool-card" aria-busy="true">
        {part.name === "lookup_order" ? `Looking up order ${orderId ?? ""}…` : "Working…"}
      </div>
    );
  }
  if (part.state === "error") {
    return <div className="tool-card tool-card--error">The lookup failed — the assistant will explain below.</div>;
  }
  if (part.name !== "lookup_order") return null; // a tool with no dedicated UI: let the model explain it in text

  const order = part.output as Order; // in real code: validate with zod before trusting it
  return (
    <dl className="order-card">
      <dt>Order</dt>
      <dd>{order.id}</dd>
      <dt>Status</dt>
      <dd>{ORDER_STATUS[order.status]}</dd>
      {order.eta && (
        <>
          <dt>Estimated delivery</dt>
          <dd>{order.eta}</dd>
        </>
      )}
    </dl>
  );
}

// app/api/chat/route.ts — the AI SDK 7 version
import { anthropic } from "@ai-sdk/anthropic";
import { convertToModelMessages, createUIMessageStreamResponse, streamText, toUIMessageStream, type UIMessage } from "ai";

export async function POST(req: Request) {
  const { messages }: { messages: UIMessage[] } = await req.json();
  const result = streamText({
    model: anthropic(process.env.CHAT_MODEL!),
    instructions: "You are the support assistant for our store.",
    messages: await convertToModelMessages(messages),
    abortSignal: req.signal, // you still have to wire abort up to the upstream yourself
  });
  return createUIMessageStreamResponse({ stream: toUIMessageStream({ stream: result.stream }) });
}

// components/SupportChat.tsx (client)
// const { messages, sendMessage, status, stop, regenerate } = useChat({
//   transport: new DefaultChatTransport({ api: "/api/chat" }),
//   throttle: 50,
// });

"use client";
import { memo, useState } from "react";
import { useChatStream } from "@/hooks/useChatStream";
import { useStickToBottom } from "@/hooks/useStickToBottom";
import type { Message } from "@/lib/chat-state";
import { ChatAnnouncer } from "./ChatAnnouncer";
import { StreamingMarkdown } from "./StreamingMarkdown";
import { ToolPart } from "./ToolPart";

const MessageItem = memo(function MessageItem({ message, onRetry }: { message: Message; onRetry: (id: string) => void }) {
  const active = message.status === "submitted" || message.status === "thinking" || message.status === "streaming";
  return (
    <article aria-label={message.role === "user" ? "You" : "Assistant"} aria-busy={active} className={`msg msg--${message.role}`}>
      {message.status === "thinking" && <p className="msg-status">Thinking…</p>}
      {message.parts.map((part, i) => {
        if (part.type === "tool-call") return <ToolPart key={part.id} part={part} />;
        if (part.type === "error")
          return (
            <p key={i} className="msg-error">
              The assistant ran into a problem.{" "}
              {part.retryable && <button onClick={() => onRetry(message.id)}>Retry</button>}
            </p>
          );
        return message.role === "user" ? (
          <p key={i}>{part.text}</p>
        ) : (
          <StreamingMarkdown key={i} text={part.text} streaming={active && i === message.parts.length - 1} />
        );
      })}
      {message.status === "stopped" && <p className="msg-status">Stopped.</p>}
    </article>
  );
});

export function SupportChat() {
  const { messages, send, stop, retry, isStreaming } = useChatStream();
  const { scrollRef, contentRef, showJump, jumpToBottom } = useStickToBottom<HTMLDivElement>();
  const [input, setInput] = useState("");

  return (
    <section className="chat" aria-label="Customer support assistant">
      <div ref={scrollRef} className="chat-scroll">
        <div ref={contentRef}>
          {messages.map((m) => (
            <MessageItem key={m.id} message={m} onRetry={retry} />
          ))}
        </div>
      </div>
      {showJump && <button className="chat-jump" onClick={jumpToBottom}>Jump to latest</button>}
      <ChatAnnouncer message={messages.at(-1)} />
      <form
        onSubmit={(e) => {
          e.preventDefault();
          const text = input.trim();
          if (!text || isStreaming) return;
          send(text);
          setInput("");
        }}
      >
        <label htmlFor="chat-input" className="sr-only">Your question</label>
        <textarea id="chat-input" value={input} onChange={(e) => setInput(e.target.value)} />
        <button type="submit" disabled={isStreaming}>Send</button>
        <button type="button" onClick={stop} disabled={!isStreaming}>Stop</button>
      </form>
    </section>
  );
}

import { describe, expect, it } from "vitest";
import { encodeEvent, HEARTBEAT, type ChatEvent } from "./chat-protocol";
import { readChatEvents } from "./sse";

const EVENTS: ChatEvent[] = [
  { type: "status", value: "writing" },
  { type: "text-delta", text: "Your order is on its way 🚚" },
  { type: "done", reason: "end_turn" },
];

/** Fake body: cuts the data at the given byte positions. */
function fakeResponse(bytes: Uint8Array, cuts: number[]): Response {
  const body = new ReadableStream<Uint8Array>({
    start(controller) {
      let prev = 0;
      for (const cut of [...cuts, bytes.length]) {
        controller.enqueue(bytes.slice(prev, cut));
        prev = cut;
      }
      controller.close();
    },
  });
  return new Response(body, { headers: { "Content-Type": "text/event-stream" } });
}

async function collect(res: Response): Promise<ChatEvent[]> {
  const out: ChatEvent[] = [];
  for await (const event of readChatEvents(res)) out.push(event);
  return out;
}

describe("readChatEvents", async () => {
  const parts = [encodeEvent(EVENTS[0]), HEARTBEAT, encodeEvent(EVENTS[1]), encodeEvent(EVENTS[2])];
  const bytes = new Uint8Array(await new Blob(parts).arrayBuffer());

  it("reassembles events correctly no matter which byte the chunk is cut at", async () => {
    for (let cut = 1; cut < bytes.length; cut++) {
      expect(await collect(fakeResponse(bytes, [cut]))).toEqual(EVENTS);
    }
  });

  it("survives chunks that are a single byte each", async () => {
    const cuts = Array.from({ length: bytes.length - 1 }, (_, i) => i + 1);
    expect(await collect(fakeResponse(bytes, cuts))).toEqual(EVENTS);
  });
});