fix: table re-rendering during stream

in safari this meant the animations started over on every tick

Author: ethanwinters

packages/ai-chat-components/src/components/markdown/src/markdown-renderer.ts

@@ -116,11 +116,13 @@ export interface RenderTokenTreeOptions {
   tooltipContent?: string;
   /** Function to get formatted line count text */
   getLineCountText?: ({ count }: { count: number }) => string;
-}
 
-const EMPTY_ATTRS = {};
-const EMPTY_TABLE_HEADERS: TableCellContent[] = [];
-const EMPTY_TABLE_ROWS: TableRowContent[] = [];
+  /**
+   * Force markdown tables to render in loading mode.
+   * Useful for freezing streaming table visuals until stream completion.
+   */
+  forceTableLoading?: boolean;
+}
 
 /**
  * Converts TokenTree to Lit TemplateResult.
@@ -432,6 +434,7 @@ function renderWithStaticTag(
 
       const {
         streaming,
+        forceTableLoading,
         context: parentContext,
         filterPlaceholderText,
         previousPageText,
@@ -444,8 +447,9 @@ function renderWithStaticTag(
       } = options;
 
       // Determine if we should show loading state during streaming
-      let isLoading = false;
+      let isLoading = Boolean(forceTableLoading);
       if (
+        !isLoading &&
         streaming &&
         parentContext?.parentChildren &&
         parentContext?.currentIndex !== undefined
@@ -482,33 +486,28 @@ function renderWithStaticTag(
           : null,
       });
 
-      // Extract table data or use empty placeholders for loading state
-      let headers: TableCellContent[];
-      let tableRows: TableRowContent[];
-
-      if (!isLoading) {
-        const extractedData = extractTableData(node);
+      if (isLoading) {
+        // Keep loading output stable during streaming table assembly.
+        return html`<div class="cds-aichat-table-holder">
+          <cds-aichat-table .loading=${true}></cds-aichat-table>
+        </div>`;
+      }
 
-        headers = extractedData.headers.map((cell) =>
-          createCellContent(cell, { isInThead: true }),
-        );
+      const extractedData = extractTableData(node);
 
-        tableRows = extractedData.rows.map((row) => ({
-          cells: row.map((cell) => createCellContent(cell)),
-        }));
-      } else {
-        // Use static empty arrays to prevent re-renders during streaming
-        headers = EMPTY_TABLE_HEADERS;
-        tableRows = EMPTY_TABLE_ROWS;
-      }
+      const headers: TableCellContent[] = extractedData.headers.map((cell) =>
+        createCellContent(cell, { isInThead: true }),
+      );
 
-      const tableAttrs = isLoading ? EMPTY_ATTRS : attrs;
+      const tableRows: TableRowContent[] = extractedData.rows.map((row) => ({
+        cells: row.map((cell) => createCellContent(cell)),
+      }));
 
       return html`<div class="cds-aichat-table-holder">
         <cds-aichat-table
           .headers=${headers}
           .rows=${tableRows}
-          .loading=${isLoading}
+          .loading=${false}
           .filterPlaceholderText=${filterPlaceholderText || "Filter table..."}
           .previousPageText=${previousPageText || "Previous page"}
           .nextPageText=${nextPageText || "Next page"}
@@ -519,7 +518,7 @@ function renderWithStaticTag(
           DEFAULT_PAGINATION_SUPPLEMENTAL_TEXT}
           .getPaginationStatusText=${getPaginationStatusText ||
           DEFAULT_PAGINATION_STATUS_TEXT}
-          ...=${tableAttrs}
+          ...=${attrs}
         ></cds-aichat-table>
       </div>`;
     }

packages/ai-chat-components/src/components/markdown/src/markdown.ts

@@ -20,6 +20,20 @@ import { renderTokenTree } from "./markdown-renderer.js";
 import { consoleError } from "./utils.js";
 import { markdownTemplate } from "./markdown.template.js";
 
+function hasTrailingTableToken(node: TokenTree): boolean {
+  const children = node.children || [];
+  for (let index = 0; index < children.length; index++) {
+    const child = children[index];
+    if (child.token.tag === "table" && index === children.length - 1) {
+      return true;
+    }
+    if (hasTrailingTableToken(child)) {
+      return true;
+    }
+  }
+  return false;
+}
+
 /**
  * Markdown component
  * @element cds-aichat-markdown
@@ -144,6 +158,8 @@ class CDSAIChatMarkdown extends LitElement {
    */
   private renderTask: Promise<void> | null = null;
 
+  private hasRenderedStreamingTableLoadingFrame = false;
+
   connectedCallback() {
     super.connectedCallback();
     this.needsReparse = true;
@@ -235,23 +251,61 @@ class CDSAIChatMarkdown extends LitElement {
    */
   private renderMarkdown = async () => {
     try {
+      let nextTokenTree = this.tokenTree;
+
       if (this.needsReparse) {
         // First, we take the markdown we were given and use the markdown-it parser to turn is into a tree we can
         // transform into Lit components and compare smartly to avoid re-renders of components that were already
         // rendered when the markdown is updated (likely by streaming, but possibly by an edit somewhere in the
         // middle). It takes the current tokenTree as an argument for quick diffing to avoid re-creating parts
         // of the tree.
-        this.tokenTree = markdownToTokenTree(
+        nextTokenTree = markdownToTokenTree(
           this._slottedMarkdown,
           this.tokenTree,
           !this.removeHTML,
         );
         this.needsReparse = false;
       }
 
+      const hasStreamingTailTable =
+        Boolean(this.streaming) && hasTrailingTableToken(nextTokenTree);
+
+      if (nextTokenTree !== this.tokenTree) {
+        this.tokenTree = nextTokenTree;
+      }
+
+      if (hasStreamingTailTable) {
+        if (!this.hasRenderedStreamingTableLoadingFrame) {
+          this.renderedContent = renderTokenTree(nextTokenTree, {
+            sanitize: this.sanitizeHTML,
+            streaming: this.streaming,
+            highlight: this.highlight,
+            // Table strings
+            filterPlaceholderText: this.filterPlaceholderText,
+            previousPageText: this.previousPageText,
+            nextPageText: this.nextPageText,
+            itemsPerPageText: this.itemsPerPageText,
+            downloadLabelText: this.downloadLabelText,
+            locale: this.locale,
+            getPaginationSupplementalText: this.getPaginationSupplementalText,
+            getPaginationStatusText: this.getPaginationStatusText,
+            // Code snippet strings
+            feedback: this.feedback,
+            showLessText: this.showLessText,
+            showMoreText: this.showMoreText,
+            tooltipContent: this.tooltipContent,
+            getLineCountText: this.getLineCountText,
+          });
+          this.hasRenderedStreamingTableLoadingFrame = true;
+        }
+        return;
+      }
+
+      this.hasRenderedStreamingTableLoadingFrame = false;
+
       // Next we take that tree and transform it into Lit content to be rendered into the template.
       // this.renderedContent is what is rendered in the template directly.
-      this.renderedContent = renderTokenTree(this.tokenTree, {
+      this.renderedContent = renderTokenTree(nextTokenTree, {
         sanitize: this.sanitizeHTML,
         streaming: this.streaming,
         highlight: this.highlight,

packages/ai-chat-components/src/components/table/src/table.ts

@@ -228,6 +228,7 @@ class CDSAIChatTable extends LitElement {
    * @internal
    */
   private tableRuntimePromise: Promise<TableRuntimeModule> | null = null;
+  private hasRenderedLoadingFrame = false;
 
   connectedCallback() {
     super.connectedCallback();
@@ -248,6 +249,25 @@ class CDSAIChatTable extends LitElement {
     this._setPageSize();
   }
 
+  /**
+   * While loading we only need to paint the skeleton once. Streaming updates can
+   * mutate rows/headers rapidly, so skip subsequent updates until loading ends.
+   */
+  protected shouldUpdate(changedProperties: PropertyValues<this>) {
+    if (changedProperties.has("loading")) {
+      if (this.loading) {
+        this.hasRenderedLoadingFrame = false;
+      }
+      return true;
+    }
+
+    if (this.loading && this.hasRenderedLoadingFrame) {
+      return false;
+    }
+
+    return true;
+  }
+
   /**
    * Updates the CSS custom property `--cds-chat-table-width` based on the parent element's width.
    * Also calculates the default page size on first run based on the measured width.
@@ -286,19 +306,31 @@ class CDSAIChatTable extends LitElement {
    * @protected
    */
   protected willUpdate(changedProperties: PropertyValues<this>) {
+    const loadingJustFinished =
+      changedProperties.has("loading") &&
+      changedProperties.get("loading") === true &&
+      this.loading === false;
+
     // If the headers or rows has recently updated and both are defined than we should validate the table
     // data. This will likely only happen on the web components first render cycle when the props go from undefined to
     // defined.
     if (
-      (changedProperties.has("headers") || changedProperties.has("rows")) &&
+      (changedProperties.has("headers") ||
+        changedProperties.has("rows") ||
+        loadingJustFinished) &&
+      !this.loading &&
       this.headers !== undefined &&
       this.rows !== undefined
     ) {
       this._calcIsTableValid();
     }
 
     // If the value of tableRows updated then initialize the internal rows arrays.
-    if (changedProperties.has("rows") && this.rows !== undefined) {
+    if (
+      (changedProperties.has("rows") || loadingJustFinished) &&
+      !this.loading &&
+      this.rows !== undefined
+    ) {
       this._initializeRowsArrays();
       this._setPageSize();
     }
@@ -579,11 +611,13 @@ class CDSAIChatTable extends LitElement {
     // This could be used while we wait for a md stream containing a table to complete.
     const runtime = this.tableRuntime;
     if (this.loading || !runtime) {
+      this.hasRenderedLoadingFrame = true;
       if (!runtime) {
         void this.ensureTableRuntime();
       }
       return tableSkeletonTemplate(this._currentPageSize);
     }
+    this.hasRenderedLoadingFrame = false;
 
     const { tableTemplate, tablePaginationTemplate } = runtime;