perf(ui): stop rescanning whole replies on every streamed chunk (#127749)

* perf(ui): scan streaming markdown incrementally * perf(ui): own streaming markdown cache invalidation Co-authored-by: vyctorbrzezowski <krzyszchweski@gmail.com> * docs: respect release-owned changelog policy Release-note context remains documented in the pull request evidence. Co-authored-by: vyctorbrzezowski <krzyszchweski@gmail.com> --------- Co-authored-by: Peter Steinberger <steipete@gmail.com>
2026-08-26 20:35:39 -06:00 · 2026-08-22 03:38:32 -03:00
parent 32cf13dc90
commit dacfd09786
7 changed files with 333 additions and 53 deletions
@@ -10,6 +10,8 @@ const FENCE_OPEN_RE = /^[ \t]{0,3}(`{3,}|~{3,})/;
 const FENCE_CONTAINER_PREFIX_RE = /^[ \t]{0,3}(?:(?:>\s?)|(?:(?:[-+*]|\d{1,9}[.)])[ \t]+))/;
 const LIST_ITEM_OPEN_RE = /^[ \t]{0,3}(?:[-+*]|\d{1,9}[.)])[ \t]+/u;
 const LINK_REFERENCE_CANDIDATE_RE = /^[ \t]*\[/u;
+const DISCLOSURE_LINE_CANDIDATE_RE = /^[ \t]*<\/?(?:details|summary)(?=[\s>])/iu;
+const STREAMING_SPLIT_CACHE_LIMIT = 8;

 type DetailsFrame = { hasSummary: boolean };
 type FenceMarker = { length: number; marker: "`" | "~" };
@@ -106,20 +108,70 @@ type StreamingMarkdownSplit = {
  tailRepairStart: number | null;
 };

-export function splitStableStreamingMarkdown(markdownLocal: string): StreamingMarkdownSplit {
-  let boundary = 0;
-  let index = 0;
-  let openFence: FenceMarker | null = null;
+type StreamingMarkdownCursor = {
+  boundary: number;
+  firstListOffset: number | null;
+  hasLinkReferenceDefinition: boolean;
+  index: number;
+  lastFenceOffset: number;
+  lineMode: "fence" | "plain" | null;
+  openFence: FenceMarker | null;
+};
+
+type StreamingMarkdownCacheEntry = {
+  cursor: StreamingMarkdownCursor;
+  markdown: string;
+};
+
+// A reused row key does not imply append-only text: rollovers, snapshots, and
+// completed citation markers can all replace the normalized Markdown prefix.
+const streamingSplitCache = new Map<string, StreamingMarkdownCacheEntry>();
+
+function findStreamingCodeSpans(markdown: string, start: number): Array<[number, number]> {
+  return findMarkdownCodeSpans(markdown.slice(start)).map(([from, to]) => [
+    from + start,
+    to + start,
+  ]);
+}
+
+function scanStableStreamingMarkdown(
+  markdownLocal: string,
+  cursor: StreamingMarkdownCursor = {
+    boundary: 0,
+    firstListOffset: null,
+    hasLinkReferenceDefinition: false,
+    index: 0,
+    lastFenceOffset: 0,
+    lineMode: null,
+    openFence: null,
+  },
+): { cursor: StreamingMarkdownCursor; result: StreamingMarkdownSplit } {
+  let { boundary, firstListOffset, hasLinkReferenceDefinition, index, lastFenceOffset } = cursor;
+  let lineMode = cursor.lineMode;
+  let openFence = cursor.openFence;
  const detailsStack: DetailsFrame[] = [];
-  const codeSpans = findMarkdownCodeSpans(markdownLocal);
-  let lastFenceOffset = 0;
-  let firstListOffset: number | null = null;
-  let hasLinkReferenceDefinition = false;
+  let codeSpans: ReturnType<typeof findMarkdownCodeSpans> | undefined;
+  let resumeCursor = cursor;

  while (index < markdownLocal.length) {
    const nextLineBreak = markdownLocal.indexOf("\n", index);
    const lineEnd = nextLineBreak === -1 ? markdownLocal.length : nextLineBreak + 1;
+    if (lineMode) {
+      index = lineEnd;
+      lineMode = nextLineBreak === -1 ? lineMode : null;
+      resumeCursor = {
+        boundary,
+        firstListOffset,
+        hasLinkReferenceDefinition,
+        index,
+        lastFenceOffset,
+        lineMode,
+        openFence,
+      };
+      continue;
+    }
    const line = markdownLocal.slice(index, nextLineBreak === -1 ? lineEnd : nextLineBreak);
+    const lineFence = openFence;

    if (openFence) {
      if (isFenceClose(line, openFence)) {
@@ -129,32 +181,52 @@ export function splitStableStreamingMarkdown(markdownLocal: string): StreamingMa
          boundary = lineEnd;
        }
      }
-      index = lineEnd;
-      continue;
-    }
-
-    if (firstListOffset === null && LIST_ITEM_OPEN_RE.test(line)) {
-      firstListOffset = index;
-    }
-
-    const openingFence = getFenceMarker(line);
-    if (openingFence) {
-      openFence = openingFence;
-      lastFenceOffset = lineEnd;
-      index = lineEnd;
-      continue;
-    }
-
-    updateDetailsStack(line, detailsStack, false, codeSpans, index);
-    if (detailsStack.length === 0) {
-      if (LINK_REFERENCE_CANDIDATE_RE.test(stripMarkdownContainerPrefixes(line).content)) {
-        hasLinkReferenceDefinition = true;
+    } else {
+      if (firstListOffset === null && LIST_ITEM_OPEN_RE.test(line)) {
+        firstListOffset = index;
      }
-      if (line.trim() === "") {
-        boundary = lineEnd;
+
+      const openingFence = getFenceMarker(line);
+      if (openingFence) {
+        openFence = openingFence;
+        lastFenceOffset = lineEnd;
+      } else {
+        const strippedLine = stripMarkdownContainerPrefixes(line).content;
+        if (DISCLOSURE_LINE_CANDIDATE_RE.test(strippedLine)) {
+          updateDetailsStack(
+            line,
+            detailsStack,
+            false,
+            (codeSpans ??= findStreamingCodeSpans(markdownLocal, firstListOffset ?? boundary)),
+            index,
+          );
+        }
+        if (detailsStack.length === 0) {
+          if (LINK_REFERENCE_CANDIDATE_RE.test(strippedLine)) {
+            hasLinkReferenceDefinition = true;
+          }
+          if (line.trim() === "") {
+            boundary = lineEnd;
+          }
+        }
      }
    }
    index = lineEnd;
+    if (
+      detailsStack.length === 0 &&
+      (nextLineBreak !== -1 || canResumeStreamingLine(line, lineFence))
+    ) {
+      lineMode = nextLineBreak === -1 ? (lineFence ? "fence" : "plain") : null;
+      resumeCursor = {
+        boundary,
+        firstListOffset,
+        hasLinkReferenceDefinition,
+        index,
+        lastFenceOffset,
+        lineMode,
+        openFence,
+      };
+    }
  }

  // A bracket-leading line can start a multiline or escaped reference label.
@@ -168,11 +240,52 @@ export function splitStableStreamingMarkdown(markdownLocal: string): StreamingMa
  }

  return {
-    boundary,
-    tailRepairStart: openFence ? null : Math.max(boundary, lastFenceOffset),
+    cursor: resumeCursor,
+    result: {
+      boundary,
+      tailRepairStart: openFence ? null : Math.max(boundary, lastFenceOffset),
+    },
  };
 }

+function canResumeStreamingLine(line: string, fence: FenceMarker | null): boolean {
+  const first = stripMarkdownContainerPrefixes(line).content.charAt(0);
+  if (!first) {
+    return false;
+  }
+  return fence ? first !== fence.marker : !/[\s`~<[\]*+\-\d>]/u.test(first);
+}
+
+export function splitStableStreamingMarkdown(
+  markdownLocal: string,
+  streamKey?: string,
+  stablePrefixLength = markdownLocal.length,
+): StreamingMarkdownSplit {
+  if (!streamKey) {
+    return scanStableStreamingMarkdown(markdownLocal).result;
+  }
+  const stableMarkdown = markdownLocal.slice(0, stablePrefixLength);
+  const cached = streamingSplitCache.get(streamKey);
+  const scanned = scanStableStreamingMarkdown(
+    stableMarkdown,
+    cached && stableMarkdown.startsWith(cached.markdown) ? cached.cursor : undefined,
+  );
+  streamingSplitCache.delete(streamKey);
+  streamingSplitCache.set(streamKey, { cursor: scanned.cursor, markdown: stableMarkdown });
+  while (streamingSplitCache.size > STREAMING_SPLIT_CACHE_LIMIT) {
+    const oldest = streamingSplitCache.keys().next().value;
+    if (oldest === undefined) {
+      break;
+    }
+    streamingSplitCache.delete(oldest);
+  }
+  // Truncation notices change on every chunk even after their capped content is
+  // fixed; retain the immutable checkpoint and rescan only that short suffix.
+  return stablePrefixLength === markdownLocal.length
+    ? scanned.result
+    : scanStableStreamingMarkdown(markdownLocal, scanned.cursor).result;
+}
+
 // Streaming-tail repair config: math is not rendered by this pipeline, so
 // completing `$$` would inject visible characters into ordinary prose.
 const streamingRemendOptions = { katex: false, linkMode: "text-only" } satisfies RemendOptions;
@@ -2,6 +2,7 @@
 import { describe, expect, it, vi } from "vitest";
 import { i18n } from "../i18n/index.ts";
 import { handleMarkdownCodeBlockClick } from "./markdown-code-blocks.ts";
+import { splitStableStreamingMarkdown } from "./markdown-streaming.ts";
 import { toSanitizedMarkdownHtml, toStreamingMarkdownHtml } from "./markdown.ts";

 function htmlFragment(html: string): HTMLElement {
@@ -856,6 +857,147 @@ PY
 });

 describe("toStreamingMarkdownHtml", () => {
+  it("keeps appended-prefix splitting below repeated full-rescan cost", () => {
+    const splitIncrementally = splitStableStreamingMarkdown as (
+      markdown: string,
+      streamKey: string,
+    ) => ReturnType<typeof splitStableStreamingMarkdown>;
+    const prefixes: string[] = [];
+    let prefix = "<details><summary>Done</summary></details>\n\n";
+    for (let index = 0; index < 96; index += 1) {
+      prefix += `${String(index).padStart(3, "0")} ${"streaming markdown ".repeat(50)}\n`;
+      prefixes.push(prefix);
+    }
+    const measure = (streamKey?: string) => {
+      const startedAt = performance.now();
+      for (const value of prefixes) {
+        if (streamKey) {
+          splitIncrementally(value, streamKey);
+        } else {
+          splitStableStreamingMarkdown(value);
+        }
+      }
+      return performance.now() - startedAt;
+    };
+    measure("line-scan-warmup");
+    const fullRescanMs = measure();
+    const incrementalMs = measure("line-scan-regression");
+
+    expect(incrementalMs).toBeLessThan(fullRescanMs / 5);
+  }, 5_000);
+
+  it("keeps chunked-prefix splits identical to full splits", () => {
+    const splitIncrementally = splitStableStreamingMarkdown as (
+      markdown: string,
+      streamKey: string,
+    ) => ReturnType<typeof splitStableStreamingMarkdown>;
+    const cases = [
+      [
+        "## Result",
+        "",
+        "A paragraph with `inline code`.",
+        "",
+        "<details>",
+        "<summary>Logs</summary>",
+        "",
+        "```ts",
+        "const value = 1;",
+        "```",
+        "",
+        "More **text**",
+        "",
+        "</details>",
+      ].join("\n"),
+      "- one\n\n  - nested\n\n[Docs][ref\\]]\n\n[ref\\]]: /docs",
+      "`` multiline\n<details> remains code\n``\n\n<details>\n<summary>Real</summary>",
+      "- item\n\n    <details>\n    <summary>Logs</summary>\n\n    still inside",
+      "1. item\n\n    <details>\n    <summary>Logs</summary>\n\n    still inside",
+    ];
+    for (const [caseIndex, markdown] of cases.entries()) {
+      for (const chunkSize of [1, 7, 64]) {
+        for (let end = chunkSize; end <= markdown.length + chunkSize; end += chunkSize) {
+          const prefix = markdown.slice(0, Math.min(end, markdown.length));
+          const key = `${caseIndex}-${chunkSize}`;
+          expect(splitIncrementally(prefix, `split-parity-${key}`)).toEqual(
+            splitStableStreamingMarkdown(prefix),
+          );
+          expect(toStreamingMarkdownHtml(prefix, {}, `html-parity-${key}`)).toBe(
+            toStreamingMarkdownHtml(prefix),
+          );
+          if (end >= markdown.length) {
+            break;
+          }
+        }
+      }
+    }
+  });
+
+  it("resets replaced streams and keeps interleaved streams independent", () => {
+    const splitIncrementally = splitStableStreamingMarkdown as (
+      markdown: string,
+      streamKey: string,
+    ) => ReturnType<typeof splitStableStreamingMarkdown>;
+    const streams = new Map([
+      ["a", "First stream\n\n```ts\nconst a = 1;"],
+      ["b", "Second stream\n\n<details>\n<summary>B</summary>"],
+    ]);
+    for (const end of [8, 16, 32, 64]) {
+      for (const [key, markdown] of streams) {
+        const prefix = markdown.slice(0, end);
+        expect(splitIncrementally(prefix, `interleaved-${key}`)).toEqual(
+          splitStableStreamingMarkdown(prefix),
+        );
+      }
+    }
+    for (const replacement of [
+      "short",
+      "Replacement\n\n- starts a different list",
+      "A much longer replacement\n\n```ts\nconst changed = true;",
+    ]) {
+      expect(splitIncrementally(replacement, "interleaved-a")).toEqual(
+        splitStableStreamingMarkdown(replacement),
+      );
+    }
+  });
+
+  it("resets an incremental cursor when a completed citation marker rewrites its prefix", () => {
+    const partial = "Intro\n\ncitevery-long-partial-citation-marker";
+    const completed = `${partial}\n\n\`\`\`ts\nconst answer = 42;`;
+
+    toStreamingMarkdownHtml(partial, {}, "citation-prefix-replacement");
+
+    expect(toStreamingMarkdownHtml(completed, {}, "citation-prefix-replacement")).toBe(
+      toStreamingMarkdownHtml(completed),
+    );
+  });
+
+  it.each(["- item", "1. item"])(
+    "keeps details inside a loose %s list continuation while streaming",
+    (item) => {
+      const markdown = `${item}\n\n    <details>\n    <summary>Logs</summary>\n\n    still inside`;
+      const fragment = htmlFragment(toStreamingMarkdownHtml(markdown, {}, `loose-list:${item}`));
+      const details = fragment.querySelector("li details");
+
+      expect(details?.querySelector("summary")?.textContent).toBe("Logs");
+      expect(details?.textContent).toContain("still inside");
+    },
+  );
+
+  it("preserves incremental parity when streamed text grows beyond the truncation cap", () => {
+    const text = Array.from(
+      { length: 210 },
+      (_, index) => `${String(index).padStart(3, "0")} ${"streamed markdown ".repeat(55)}\n`,
+    ).join("");
+
+    for (const end of [139_500, 140_050, 141_000, text.length]) {
+      const prefix = text.slice(0, end);
+
+      expect(toStreamingMarkdownHtml(prefix, {}, "truncated-stream-parity")).toBe(
+        toStreamingMarkdownHtml(prefix),
+      );
+    }
+  });
+
  it("marks a completed transcript-role header in the streaming tail", () => {
    const html = toStreamingMarkdownHtml("user[Thu 2026-07-02] question", {
      assistantTranscriptRoleHeaders: true,
@@ -483,11 +483,6 @@ function installHooks() {
  });
 }

-function formatTruncatedMarkdownInput(input: string): string {
-  const truncated = truncateText(input, MARKDOWN_CHAR_LIMIT);
-  return appendMarkdownTruncationNotice(truncated);
-}
-
 function appendMarkdownTruncationNotice(truncated: {
  text: string;
  truncated: boolean;
@@ -582,6 +577,7 @@ function toEscapedPlainTextHtml(value: string, options: MarkdownRenderEnv): stri
 export function toStreamingMarkdownHtml(
  markdownLocal: string,
  options: MarkdownRenderOptions = {},
+  streamKey?: string,
 ): string {
  const renderOptions = normalizeMarkdownRenderOptions(options);
  const rawInput = normalizeMarkdownLineBreaks(
@@ -595,9 +591,14 @@ export function toStreamingMarkdownHtml(
  if (!trimmedInput) {
    return "";
  }
-  const input = formatTruncatedMarkdownInput(trimmedInput);
+  const truncated = truncateText(trimmedInput, MARKDOWN_CHAR_LIMIT);
+  const input = appendMarkdownTruncationNotice(truncated);

-  const { boundary, tailRepairStart } = splitStableStreamingMarkdown(input);
+  const { boundary, tailRepairStart } = splitStableStreamingMarkdown(
+    input,
+    streamKey,
+    truncated.text.length,
+  );
  const stableMarkdown = input.slice(0, boundary);
  const streamingTail = input.slice(boundary);
  const stableHtml = boundary > 0 ? toSanitizedMarkdownHtml(stableMarkdown, options) : "";
@@ -935,10 +935,11 @@ suite.define(() => {
      const params = requireRecord(sendRequest.params);
      const runId = requireString(params.idempotencyKey, "chat send idempotency key");

+      const initialStream = `I will inspect the file. ${"Prior streamed output. ".repeat(20)}`;
      await gateway.emitGatewayEvent("chat", {
-        deltaText: "I will inspect the file.",
+        deltaText: initialStream,
        message: {
-          content: [{ text: "I will inspect the file.", type: "text" }],
+          content: [{ text: initialStream, type: "text" }],
          role: "assistant",
          timestamp: Date.now(),
        },
@@ -964,6 +965,22 @@ suite.define(() => {
      const toolBubble = page.locator('[data-message-id^="tool:assistant:call-read"]');
      await toolBubble.waitFor({ timeout: 10_000 });

+      const nextStream = "```ts\nconst answer = 42;";
+      await gateway.emitGatewayEvent("chat", {
+        deltaText: nextStream,
+        message: {
+          content: [{ text: nextStream, type: "text" }],
+          role: "assistant",
+          timestamp: Date.now(),
+        },
+        runId,
+        sessionKey: "main",
+        state: "delta",
+      });
+      await expect
+        .poll(() => page.locator(".chat-bubble.streaming code.language-ts").textContent())
+        .toContain("const answer = 42;");
+
      const visibleOrder = await page.locator(".chat-thread").evaluate((thread: Element) => {
        return Array.from(thread.querySelectorAll(".chat-group")).flatMap((group: Element) => {
          const text = group.textContent ?? "";
@@ -642,6 +642,7 @@ export function renderGroupedMessage(
                        opts.assistantMessageDisclosure,
                        markdownRenderOptions,
                        duplicateSuffix,
+                        opts.isStreaming ? messageKey : undefined,
                      )
                    : renderMarkdownText(
                        bodyMarkdown,
@@ -295,6 +295,7 @@ export function renderAssistantMessageMarkdown(
  disclosure: AssistantMessageDisclosure | undefined,
  markdownRenderOptions: MarkdownRenderOptions,
  duplicateSuffix?: DuplicateSuffix,
+  streamKey?: string,
 ) {
  const markdown = disclosure?.expanded
    ? (disclosure.markdown ?? previewMarkdown)
@@ -302,7 +303,7 @@ export function renderAssistantMessageMarkdown(
  const renderOptions = disclosure?.expanded
    ? { ...markdownRenderOptions, mode: "document" as const }
    : markdownRenderOptions;
-  const text = renderMarkdownText(markdown, isStreaming, renderOptions, duplicateSuffix);
+  const text = renderMarkdownText(markdown, isStreaming, renderOptions, duplicateSuffix, streamKey);
  if (!disclosure?.onRetryFullMessage) {
    return text;
  }
@@ -328,9 +329,10 @@ export function renderMarkdownText(
  isStreaming: boolean,
  markdownRenderOptions?: MarkdownRenderOptions,
  duplicateSuffix?: DuplicateSuffix,
+  streamKey?: string,
 ) {
  const rendered = isStreaming
-    ? toStreamingMarkdownHtml(markdown, markdownRenderOptions)
+    ? toStreamingMarkdownHtml(markdown, markdownRenderOptions, streamKey)
    : toSanitizedMarkdownHtml(markdown, markdownRenderOptions);
  const content = duplicateSuffix ? appendDuplicateSuffix(rendered, duplicateSuffix) : rendered;
  return html`
@@ -1663,16 +1663,20 @@ describe("grouped chat rendering", () => {
    );

    expect(markdownRenderMock).not.toHaveBeenCalled();
-    expect(streamingMarkdownRenderMock).toHaveBeenCalledWith("**live**\nreply", {
-      assistantTranscriptRoleHeaders: true,
-      codeBlockChrome: "copy",
-      codeBlockInteraction: "interactive",
-      fileLinks: true,
-      interactiveImages: false,
-      linkFavicons: false,
-      sessionLinks: true,
-      tableInteractions: "enabled",
-    });
+    expect(streamingMarkdownRenderMock).toHaveBeenCalledWith(
+      "**live**\nreply",
+      {
+        assistantTranscriptRoleHeaders: true,
+        codeBlockChrome: "copy",
+        codeBlockInteraction: "interactive",
+        fileLinks: true,
+        interactiveImages: false,
+        linkFavicons: false,
+        sessionLinks: true,
+        tableInteractions: "enabled",
+      },
+      "stream:1",
+    );
    const text = container.querySelector(".streaming-markdown");
    expect(text?.textContent).toBe("**live**\nreply");
  });