Jyotisha/frontend/tests/home-streaming-render-split.test.ts

import assert from "node:assert/strict";
import { createElement, createRef } from "react";
import { renderToString } from "react-dom/server";
import test from "node:test";

import {
  LatestAssistantEntry,
  SettledMessageList,
  UnsplitChatTranscript,
  type ChatTranscriptActions,
  type ChatTranscriptProps,
} from "../src/components/chat-transcript.tsx";
import type { ChatMessage } from "../src/lib/chat-message-view.ts";
import { createStreamFrameBuffer, type StreamFrameScheduler } from "../src/lib/stream-frame-buffer.ts";
import { settledChatMessageViews, streamingChatMessageView } from "../src/lib/chat-message-view.ts";
import {
  disableHomeStreamingRenderProbe,
  enableHomeStreamingRenderProbe,
  homeStreamingRenderProbeSnapshot,
  resetHomeStreamingRenderProbe,
} from "../src/lib/home-streaming-render-probe.ts";

function propsFor(messages: readonly ChatMessage[]): Omit<ChatTranscriptProps, "loading" | "streamingText"> {
  const actionsRef = createRef<ChatTranscriptActions>() as ChatTranscriptProps["actionsRef"];
  actionsRef.current = {
    onFeedback() {},
    onCopy() {},
    onRegenerate() {},
    onFollowUp() {},
    onStartRectification() {},
  };
  return {
    messages,
    sessionId: "session-test",
    sessionType: "consultation",
    theme: "career",
    messageFeedback: {},
    copiedMessageKey: null,
    cancellationPending: false,
    productEntrypointsDisabled: true,
    actionsRef,
  };
}

test("the split architecture renders settled history once while streaming tokens", () => {
  const messages: ChatMessage[] = [
    { role: "user", text: "问题一" },
    { role: "assistant", text: "回答一，需要足够长才能生成追问卡片。" },
    { role: "user", text: "问题二" },
    { role: "assistant", text: "回答二，需要足够长才能生成追问卡片。" },
    { role: "user", text: "请继续说明这个月的安排。" },
  ];
  const tokens = ["甲", "甲乙", "甲乙丙", "甲乙丙丁", "甲乙丙丁戊"];
  const base = propsFor(messages);

  resetHomeStreamingRenderProbe();
  enableHomeStreamingRenderProbe();
  renderToString(createElement(SettledMessageList, { ...base, loading: true }));
  for (const streamingText of tokens) {
    const streamingMessage = streamingChatMessageView(messages, true, streamingText);
    assert.ok(streamingMessage);
    renderToString(createElement(LatestAssistantEntry, {
      ...base,
      loading: true,
      message: streamingMessage,
      views: [...settledChatMessageViews(messages), streamingMessage],
      index: messages.length,
      cancellationPending: false,
    }));
  }
  const split = homeStreamingRenderProbeSnapshot();
  disableHomeStreamingRenderProbe();

  resetHomeStreamingRenderProbe();
  enableHomeStreamingRenderProbe();
  for (const streamingText of tokens) {
    renderToString(createElement(UnsplitChatTranscript, { ...base, loading: true, streamingText }));
  }
  const unsplit = homeStreamingRenderProbeSnapshot();
  disableHomeStreamingRenderProbe();

  assert.equal(split.settledListRenders, 1);
  // Former assertion: `streamingRowRenders === tokens.length`. That was a snapshot of the
  // status quo (one commit per network token), not the goal; this test drives renders by
  // hand, so the count equals the number of hand-driven renders and must never exceed it.
  assert.ok(split.streamingRowRenders <= tokens.length);
  assert.ok(split.streamingRowRenders >= 1);
  assert.equal(unsplit.unsplitListRenders, tokens.length);
  assert.ok(unsplit.settledRowRenders > split.settledRowRenders);
});

test("frame coalescing renders the streaming row once per frame, not once per token", () => {
  const messages: ChatMessage[] = [{ role: "user", text: "请继续说明这个月的安排。" }];
  const frames: Array<() => void> = [];
  const scheduler: StreamFrameScheduler = {
    requestFrame(callback) {
      frames.push(callback);
      return frames.length;
    },
    cancelFrame() { frames.length = 0; },
    requestTimeout() { return 0; },
    cancelTimeout() {},
    hidden: () => false,
  };

  resetHomeStreamingRenderProbe();
  enableHomeStreamingRenderProbe();
  const buffer = createStreamFrameBuffer<null>({
    initialMeta: null,
    scheduler,
    flush: (frame) => {
      const streamingMessage = streamingChatMessageView(messages, true, frame.answer);
      assert.ok(streamingMessage);
      renderToString(createElement(LatestAssistantEntry, {
        ...propsFor(messages),
        loading: true,
        message: streamingMessage,
        views: [...settledChatMessageViews(messages), streamingMessage],
        index: messages.length,
      }));
    },
  });
  // 200 one-character tokens arrive four per frame across fifty frames.
  let answer = "";
  for (let index = 0; index < 200; index += 1) {
    answer += "字";
    buffer.setAnswer(answer);
    if (index % 4 === 3) for (const callback of frames.splice(0)) callback();
  }
  buffer.settle();
  const probe = homeStreamingRenderProbeSnapshot();
  disableHomeStreamingRenderProbe();

  assert.ok(probe.streamingRowRenders <= 51, `rendered ${probe.streamingRowRenders} times for 200 tokens`);
  assert.ok(probe.streamingRowRenders * 3 <= 200);
});