@marimo-team/frontend 0.16.0 → 0.16.1

package/src/components/data-table/data-table.tsx

@@ -19,11 +19,13 @@ import {
   useReactTable,
 } from "@tanstack/react-table";
 import React, { memo } from "react";
+import { useLocale } from "react-aria";
 
 import { Table } from "@/components/ui/table";
 import type { GetRowIds } from "@/plugins/impl/DataTablePlugin";
 import { cn } from "@/utils/cn";
 import type { PanelType } from "../editor/chrome/panels/context-aware-panel/context-aware-panel";
+import { CellHoverTemplateFeature } from "./cell-hover-template/feature";
 import { CellSelectionFeature } from "./cell-selection/feature";
 import type { CellSelectionState } from "./cell-selection/types";
 import { CellStylingFeature } from "./cell-styling/feature";
@@ -63,6 +65,7 @@ interface DataTableProps<TData> extends Partial<DownloadActionProps> {
   rowSelection?: RowSelectionState;
   cellSelection?: CellSelectionState;
   cellStyling?: CellStyleState | null;
+  hoverTemplate?: string | null;
   onRowSelectionChange?: OnChangeFn<RowSelectionState>;
   onCellSelectionChange?: OnChangeFn<CellSelectionState>;
   getRowIds?: GetRowIds;
@@ -103,6 +106,7 @@ const DataTableInternal = <TData,>({
   rowSelection,
   cellSelection,
   cellStyling,
+  hoverTemplate,
   paginationState,
   setPaginationState,
   downloadAs,
@@ -131,6 +135,7 @@ const DataTableInternal = <TData,>({
 }: DataTableProps<TData>) => {
   const [isSearchEnabled, setIsSearchEnabled] = React.useState<boolean>(false);
   const [showLoadingBar, setShowLoadingBar] = React.useState<boolean>(false);
+  const { locale } = useLocale();
 
   const { columnPinning, setColumnPinning } = useColumnPinning(
     freezeColumnsLeft,
@@ -176,6 +181,7 @@ const DataTableInternal = <TData,>({
       ColumnFormattingFeature,
       CellSelectionFeature,
       CellStylingFeature,
+      CellHoverTemplateFeature,
       CopyColumnFeature,
       FocusRowFeature,
     ],
@@ -199,6 +205,7 @@ const DataTableInternal = <TData,>({
           },
         }
       : {}),
+    locale: locale,
     manualPagination: manualPagination,
     getPaginationRowModel: getPaginationRowModel(),
     // sorting
@@ -233,10 +240,11 @@ const DataTableInternal = <TData,>({
           ? {}
           : // No pagination, show all rows
             { pagination: { pageIndex: 0, pageSize: data.length } }),
-      rowSelection,
-      cellSelection,
+      rowSelection: rowSelection ?? {},
+      cellSelection: cellSelection ?? [],
       cellStyling,
       columnPinning: columnPinning,
+      cellHoverTemplate: hoverTemplate,
     },
   });
 

package/src/components/data-table/date-popover.tsx

@@ -1,5 +1,6 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import React from "react";
+import { useDateFormatter, useLocale } from "react-aria";
 import { Tooltip } from "@/components/ui/tooltip";
 
 interface DatePopoverProps {
@@ -19,33 +20,17 @@ export const DatePopover: React.FC<DatePopoverProps> = ({
   }
 
   const dateObj = new Date(date);
-  const relativeTime = getRelativeTime(dateObj);
 
   const content = (
     <div className="min-w-[240px] p-1 text-sm">
-      <div className="text-muted-foreground mb-2">{relativeTime}</div>
+      <div className="text-muted-foreground mb-2">
+        <RelativeTime date={dateObj} />
+      </div>
       <div className="space-y-1">
         {type === "datetime" ? (
-          Object.entries(getTimezones(dateObj)).map(
-            ([timezone, formattedDate]) => (
-              <div
-                key={timezone}
-                className="grid grid-cols-[fit-content(40px)_1fr] gap-4 items-center justify-items-end"
-              >
-                <span className="bg-muted rounded-md py-1 px-2 w-fit ml-auto">
-                  {timezone}
-                </span>
-                <span>{formattedDate}</span>
-              </div>
-            ),
-          )
+          <TimezoneDisplay date={dateObj} />
         ) : (
-          <span>
-            {dateObj.toLocaleDateString("en-US", {
-              timeZone: "UTC",
-              dateStyle: "long",
-            })}
-          </span>
+          <DateDisplay date={dateObj} />
         )}
       </div>
     </div>
@@ -58,57 +43,79 @@ export const DatePopover: React.FC<DatePopoverProps> = ({
   );
 };
 
-function getTimezones(date: Date) {
-  const localTimezone = Intl.DateTimeFormat().resolvedOptions().timeZone;
-
+const TimezoneDisplay = ({ date }: { date: Date }) => {
+  const { locale } = useLocale();
+  const localTimezone = Intl.DateTimeFormat(locale).resolvedOptions().timeZone;
   const hasSubSeconds = date.getUTCMilliseconds() !== 0;
-  if (hasSubSeconds) {
-    return {
-      UTC: new Intl.DateTimeFormat("en-US", {
-        timeZone: "UTC",
-        year: "numeric",
-        month: "2-digit",
-        day: "2-digit",
-        hour: "2-digit",
-        minute: "2-digit",
-        second: "2-digit",
-        fractionalSecondDigits: 3,
-      }).format(date),
-      [localTimezone]: new Intl.DateTimeFormat("en-US", {
-        timeZone: localTimezone,
-        year: "numeric",
-        month: "2-digit",
-        day: "2-digit",
-        hour: "2-digit",
-        minute: "2-digit",
-        second: "2-digit",
-        fractionalSecondDigits: 3,
-      }).format(date),
-    };
-  }
 
-  return {
-    UTC: new Intl.DateTimeFormat("en-US", {
-      timeZone: "UTC",
-      dateStyle: "long",
-      timeStyle: "medium",
-    }).format(date),
-    [localTimezone]: new Intl.DateTimeFormat("en-US", {
-      timeZone: localTimezone,
-      dateStyle: "long",
-      timeStyle: "medium",
-    }).format(date),
-  };
-}
+  const utcFormatter = useDateFormatter(
+    hasSubSeconds
+      ? {
+          timeZone: "UTC",
+          year: "numeric",
+          month: "2-digit",
+          day: "2-digit",
+          hour: "2-digit",
+          minute: "2-digit",
+          second: "2-digit",
+          fractionalSecondDigits: 3,
+        }
+      : {
+          timeZone: "UTC",
+          dateStyle: "long",
+          timeStyle: "medium",
+        },
+  );
+
+  const localFormatter = useDateFormatter(
+    hasSubSeconds
+      ? {
+          timeZone: localTimezone,
+          year: "numeric",
+          month: "2-digit",
+          day: "2-digit",
+          hour: "2-digit",
+          minute: "2-digit",
+          second: "2-digit",
+          fractionalSecondDigits: 3,
+        }
+      : {
+          timeZone: localTimezone,
+          dateStyle: "long",
+          timeStyle: "medium",
+        },
+  );
+
+  return (
+    <>
+      <div className="grid grid-cols-[fit-content(40px)_1fr] gap-4 items-center justify-items-end">
+        <span className="bg-muted rounded-md py-1 px-2 w-fit ml-auto">UTC</span>
+        <span>{utcFormatter.format(date)}</span>
+      </div>
+      <div className="grid grid-cols-[fit-content(40px)_1fr] gap-4 items-center justify-items-end">
+        <span className="bg-muted rounded-md py-1 px-2 w-fit ml-auto">
+          {localTimezone}
+        </span>
+        <span>{localFormatter.format(date)}</span>
+      </div>
+    </>
+  );
+};
+
+const DateDisplay = ({ date }: { date: Date }) => {
+  const dateFormatter = useDateFormatter({
+    timeZone: "UTC",
+    dateStyle: "long",
+  });
 
-/**
- * Converts a date into a human-readable relative time string (e.g. "2 hours ago", "in 5 minutes")
- * Uses Intl.RelativeTimeFormat for localized formatting
- */
-function getRelativeTime(date: Date): string {
-  // Initialize relative time formatter with English locale and "auto" numeric style
-  // "auto" allows for strings like "yesterday" instead of "1 day ago"
-  const relativeTimeFormatter = new Intl.RelativeTimeFormat("en", {
+  return <span>{dateFormatter.format(date)}</span>;
+};
+
+const RelativeTime = ({ date }: { date: Date }) => {
+  const { locale } = useLocale();
+
+  // Initialize relative time formatter with current locale and "auto" numeric style
+  const relativeTimeFormatter = new Intl.RelativeTimeFormat(locale, {
     numeric: "auto",
   });
 
@@ -117,7 +124,6 @@ function getRelativeTime(date: Date): string {
   const differenceInSeconds = (currentTime.getTime() - date.getTime()) / 1000;
 
   // Define time units with their thresholds and conversion factors
-  // Format: [threshold before next unit, seconds in this unit, unit name]
   const timeUnits: Array<[number, number, string]> = [
     [60, 1, "second"], // Less than 60 seconds
     [60, 60, "minute"], // Less than 60 minutes
@@ -134,13 +140,17 @@ function getRelativeTime(date: Date): string {
       // Convert to fixed decimal and negate since RelativeTimeFormat expects
       // negative values for past times and positive for future times
       const roundedValue = -Number(valueInUnits.toFixed(1));
-      return relativeTimeFormatter.format(
-        roundedValue,
-        unitName as Intl.RelativeTimeFormatUnit,
+      return (
+        <span>
+          {relativeTimeFormatter.format(
+            roundedValue,
+            unitName as Intl.RelativeTimeFormatUnit,
+          )}
+        </span>
       );
     }
   }
 
   // Should never reach here due to Infinity threshold, but provide fallback
-  return relativeTimeFormatter.format(0, "second");
-}
+  return <span>{relativeTimeFormatter.format(0, "second")}</span>;
+};

package/src/components/data-table/filter-pills.tsx

@@ -7,6 +7,7 @@ import type {
   Table,
 } from "@tanstack/react-table";
 import { XIcon } from "lucide-react";
+import { type DateFormatter, useDateFormatter } from "react-aria";
 import { logNever } from "@/utils/assertNever";
 import { Badge } from "../ui/badge";
 import type { ColumnFilterValue } from "./filters";
@@ -18,12 +19,21 @@ interface Props<TData> {
 }
 
 export const FilterPills = <TData,>({ filters, table }: Props<TData>) => {
+  const timeFormatter = useDateFormatter({
+    hour: "2-digit",
+    minute: "2-digit",
+    second: "2-digit",
+  });
+
   if (!filters || filters.length === 0) {
     return null;
   }
 
   function renderFilterPill(filter: ColumnFilter) {
-    const formattedValue = formatValue(filter.value as ColumnFilterValue);
+    const formattedValue = formatValue(
+      filter.value as ColumnFilterValue,
+      timeFormatter,
+    );
     if (!formattedValue) {
       return null;
     }
@@ -52,7 +62,7 @@ export const FilterPills = <TData,>({ filters, table }: Props<TData>) => {
   );
 };
 
-function formatValue(value: ColumnFilterValue) {
+function formatValue(value: ColumnFilterValue, timeFormatter: DateFormatter) {
   if (!("type" in value)) {
     return;
   }
@@ -71,14 +81,9 @@ function formatValue(value: ColumnFilterValue) {
     return formatMinMax(value.min?.toISOString(), value.max?.toISOString());
   }
   if (value.type === "time") {
-    const formatTime = new Intl.DateTimeFormat("en-US", {
-      hour: "2-digit",
-      minute: "2-digit",
-      second: "2-digit",
-    });
     return formatMinMax(
-      value.min ? formatTime.format(value.min) : undefined,
-      value.max ? formatTime.format(value.max) : undefined,
+      value.min ? timeFormatter.format(value.min) : undefined,
+      value.max ? timeFormatter.format(value.max) : undefined,
     );
   }
   if (value.type === "datetime") {

package/src/components/data-table/header-items.tsx

@@ -34,6 +34,7 @@ import { NAMELESS_COLUMN_PREFIX } from "./columns";
 
 export function renderFormatOptions<TData, TValue>(
   column: Column<TData, TValue>,
+  locale: string,
 ) {
   const dataType: DataType | undefined = column.columnDef.meta?.dataType;
   const columnFormatOptions = dataType ? formatOptions[dataType] : [];
@@ -51,6 +52,9 @@ export function renderFormatOptions<TData, TValue>(
       </DropdownMenuSubTrigger>
       <DropdownMenuPortal>
         <DropdownMenuSubContent>
+          <div className="text-xs text-muted-foreground px-2 py-1">
+            Locale: {locale}
+          </div>
           {Boolean(currentFormat) && (
             <>
               <DropdownMenuItem
@@ -72,7 +76,7 @@ export function renderFormatOptions<TData, TValue>(
                 {option}
               </span>
               <span className="ml-auto pl-5 text-xs text-muted-foreground">
-                {formattingExample(option)}
+                {formattingExample(option, locale)}
               </span>
             </DropdownMenuItem>
           ))}

package/src/components/data-table/pagination.tsx

@@ -9,8 +9,10 @@ import {
   ChevronsLeft,
   ChevronsRight,
 } from "lucide-react";
+import { useLocale } from "react-aria";
 import { Button } from "@/components/ui/button";
 import { Events } from "@/utils/events";
+import { prettyNumber } from "@/utils/numbers";
 import { PluralWord } from "@/utils/pluralize";
 import {
   Select,
@@ -40,6 +42,8 @@ export const DataTablePagination = <TData,>({
   tableLoading,
   showPageSizeSelector,
 }: DataTablePaginationProps<TData>) => {
+  const { locale } = useLocale();
+
   const renderTotal = () => {
     const { rowSelection, cellSelection } = table.getState();
     let selected = Object.keys(rowSelection).length;
@@ -58,7 +62,7 @@ export const DataTablePagination = <TData,>({
     if (isAllPageSelected && !isAllSelected) {
       return (
         <>
-          <span>{prettyNumber(selected)} selected</span>
+          <span>{prettyNumber(selected, locale)} selected</span>
           <Button
             size="xs"
             data-testid="select-all-button"
@@ -73,7 +77,7 @@ export const DataTablePagination = <TData,>({
               }
             }}
           >
-            Select all {prettyNumber(numRows)}
+            Select all {prettyNumber(numRows, locale)}
           </Button>
         </>
       );
@@ -82,7 +86,7 @@ export const DataTablePagination = <TData,>({
     if (selected) {
       return (
         <>
-          <span>{prettyNumber(selected)} selected</span>
+          <span>{prettyNumber(selected, locale)} selected</span>
           <Button
             size="xs"
             data-testid="clear-selection-button"
@@ -107,7 +111,11 @@ export const DataTablePagination = <TData,>({
       );
     }
 
-    const rowColumnCount = prettifyRowColumnCount(numRows, totalColumns);
+    const rowColumnCount = prettifyRowColumnCount(
+      numRows,
+      totalColumns,
+      locale,
+    );
     return <span>{rowColumnCount}</span>;
   };
   const currentPage = Math.min(
@@ -199,7 +207,9 @@ export const DataTablePagination = <TData,>({
               handlePageChange(() => table.setPageIndex(page))
             }
           />
-          <span className="shrink-0">of {prettyNumber(totalPages)}</span>
+          <span className="shrink-0">
+            of {prettyNumber(totalPages, locale)}
+          </span>
         </div>
         <Button
           size="xs"
@@ -232,10 +242,6 @@ export const DataTablePagination = <TData,>({
   );
 };
 
-function prettyNumber(value: number): string {
-  return new Intl.NumberFormat().format(value);
-}
-
 export const PageSelector = ({
   currentPage,
   totalPages,
@@ -313,17 +319,18 @@ export const PageSelector = ({
   );
 };
 
-export function prettifyRowCount(rowCount: number): string {
-  return `${prettyNumber(rowCount)} ${new PluralWord("row").pluralize(rowCount)}`;
+export function prettifyRowCount(rowCount: number, locale: string): string {
+  return `${prettyNumber(rowCount, locale)} ${new PluralWord("row").pluralize(rowCount)}`;
 }
 
 export const prettifyRowColumnCount = (
   numRows: number | "too_many",
   totalColumns: number,
+  locale: string,
 ): string => {
   const rowsLabel =
-    numRows === "too_many" ? "Unknown" : prettifyRowCount(numRows);
-  const columnsLabel = `${prettyNumber(totalColumns)} ${new PluralWord("column").pluralize(totalColumns)}`;
+    numRows === "too_many" ? "Unknown" : prettifyRowCount(numRows, locale);
+  const columnsLabel = `${prettyNumber(totalColumns, locale)} ${new PluralWord("column").pluralize(totalColumns)}`;
 
   return [rowsLabel, columnsLabel].join(", ");
 };

package/src/components/data-table/renderers.tsx

@@ -93,6 +93,25 @@ export const DataTableBody = <TData,>({
     handleCellsKeyDown,
   } = useCellRangeSelection({ table });
 
+  function applyHoverTemplate(
+    template: string,
+    cells: Array<Cell<TData, unknown>>,
+  ): string {
+    const variableRegex = /{{(\w+)}}/g;
+    // Map column id -> stringified value
+    const idToValue = new Map<string, string>();
+    for (const c of cells) {
+      const v = c.getValue();
+      // Prefer empty string for nulls to keep tooltip clean
+      const s = renderUnknownValue({ value: v, nullAsEmptyString: true });
+      idToValue.set(c.column.id, s);
+    }
+    return template.replace(variableRegex, (_substr, varName: string) => {
+      const val = idToValue.get(varName);
+      return val !== undefined ? val : `{{${varName}}}`;
+    });
+  }
+
   const renderCells = (cells: Array<Cell<TData, unknown>>) => {
     return cells.map((cell) => {
       const { className, style: pinningstyle } = getPinningStyles(cell.column);
@@ -101,6 +120,7 @@ export const DataTableBody = <TData,>({
         cell.getUserStyling?.() || {},
         pinningstyle,
       );
+
       return (
         <TableCell
           tabIndex={0}
@@ -145,10 +165,18 @@ export const DataTableBody = <TData,>({
           const isRowViewedInPanel =
             rowViewerPanelOpen && viewedRowIdx === rowIndex;
 
+          // Compute hover title once per row using this row's cells (visible or hidden)
+          const hoverTemplate = table.getState().cellHoverTemplate || null;
+          const rowCells = row.getAllCells();
+          const rowTitle = hoverTemplate
+            ? applyHoverTemplate(hoverTemplate, rowCells)
+            : undefined;
+
           return (
             <TableRow
               key={row.id}
               data-state={row.getIsSelected() && "selected"}
+              title={rowTitle}
               // These classes ensure that empty rows (nulls) still render
               className={cn(
                 "border-t h-6",

package/src/components/data-table/row-viewer-panel/row-viewer.tsx

@@ -15,6 +15,7 @@ import {
   SearchIcon,
 } from "lucide-react";
 import { useRef, useState } from "react";
+import { useLocale } from "react-aria";
 import { ColumnName } from "@/components/datasources/components";
 import { CopyClipboardIcon } from "@/components/icons/copy-icon";
 import { KeyboardHotkeys } from "@/components/shortcuts/renderShortcut";
@@ -66,6 +67,7 @@ export const RowViewerPanel: React.FC<RowViewerPanelProps> = ({
   const [searchQuery, setSearchQuery] = useState("");
   const panelRef = useRef<HTMLDivElement>(null);
   const searchInputRef = useRef<HTMLInputElement>(null);
+  const { locale } = useLocale();
 
   const tooManyRows = totalRows === TOO_MANY_ROWS;
 
@@ -204,13 +206,13 @@ export const RowViewerPanel: React.FC<RowViewerPanelProps> = ({
               applyColumnFormatting: (value) => value,
             } as Column<unknown>;
 
-            const cellContent = renderCellValue(
-              mockColumn,
-              () => columnValue,
-              () => columnValue,
-              undefined,
-              "text-left break-word",
-            );
+            const cellContent = renderCellValue({
+              column: mockColumn,
+              renderValue: () => columnValue,
+              getValue: () => columnValue,
+              selectCell: undefined,
+              cellStyles: "text-left break-word",
+            });
 
             const copyValue =
               typeof columnValue === "object"
@@ -286,7 +288,7 @@ export const RowViewerPanel: React.FC<RowViewerPanelProps> = ({
         <span className="text-xs">
           {tooManyRows
             ? `Row ${rowIdx + 1}`
-            : `Row ${rowIdx + 1} of ${prettifyRowCount(totalRows)}`}
+            : `Row ${rowIdx + 1} of ${prettifyRowCount(totalRows, locale)}`}
         </span>
         <Button
           variant="outline"

package/src/components/datasources/column-preview.tsx

@@ -3,6 +3,7 @@
 import { useAtomValue } from "jotai";
 import { PlusSquareIcon } from "lucide-react";
 import React, { Suspense } from "react";
+import { useLocale } from "react-aria";
 import { maybeAddAltairImport } from "@/core/cells/add-missing-import";
 import { useCellActions } from "@/core/cells/cells";
 import { useLastFocusedCellId } from "@/core/cells/focus";
@@ -42,6 +43,7 @@ export const DatasetColumnPreview: React.FC<{
 }> = ({ table, column, preview, onAddColumnChart, sqlTableContext }) => {
   const { theme } = useTheme();
   const { previewDatasetColumn } = useRequestClient();
+  const { locale } = useLocale();
 
   const previewColumn = () => {
     previewDatasetColumn({
@@ -107,7 +109,8 @@ export const DatasetColumnPreview: React.FC<{
       refetchPreview: previewColumn,
     });
 
-  const stats = preview.stats && renderStats(preview.stats, column.type);
+  const stats =
+    preview.stats && renderStats(preview.stats, column.type, locale);
 
   const chart = preview.chart_spec && renderChart(preview.chart_spec, theme);
 
@@ -175,6 +178,7 @@ export function renderPreviewError({
 export function renderStats(
   stats: Partial<Record<ColumnHeaderStatsKey, unknown>>,
   dataType: DataType,
+  locale: string,
 ) {
   return (
     <div className="gap-x-16 gap-y-1 grid grid-cols-2-fit border rounded p-2 empty:hidden">
@@ -189,7 +193,7 @@ export function renderStats(
               {convertStatsName(key as ColumnHeaderStatsKey, dataType)}
             </span>
             <span className="text-xs font-bold text-muted-foreground tracking-wide">
-              {prettyNumber(value)}
+              {prettyNumber(value, locale)}
             </span>
             <CopyClipboardIcon
               className="h-3 w-3 invisible group-hover:visible"

package/src/components/datasources/datasources.tsx

@@ -304,10 +304,12 @@ const DatabaseItem: React.FC<{
 }> = ({ hasSearch, engineName, database, children }) => {
   const [isExpanded, setIsExpanded] = React.useState(false);
   const [isSelected, setIsSelected] = React.useState(false);
+  const [prevHasSearch, setPrevHasSearch] = React.useState(hasSearch);
 
-  React.useEffect(() => {
+  if (prevHasSearch !== hasSearch) {
+    setPrevHasSearch(hasSearch);
     setIsExpanded(hasSearch);
-  }, [hasSearch]);
+  }
 
   return (
     <>
@@ -399,14 +401,10 @@ const SchemaItem: React.FC<{
   children: React.ReactNode;
   hasSearch: boolean;
 }> = ({ databaseName, schema, children, hasSearch }) => {
-  const [isExpanded, setIsExpanded] = React.useState(false);
+  const [isExpanded, setIsExpanded] = React.useState(hasSearch);
   const [isSelected, setIsSelected] = React.useState(false);
   const uniqueValue = `${databaseName}:${schema.name}`;
 
-  React.useEffect(() => {
-    setIsExpanded(hasSearch);
-  }, [hasSearch]);
-
   if (isSchemaless(schema.name)) {
     return children;
   }
@@ -693,11 +691,9 @@ const DatasetColumnItem: React.FC<{
   const closeAllColumns = useAtomValue(closeAllColumnsAtom);
   const setExpandedColumns = useSetAtom(expandedColumnsAtom);
 
-  React.useEffect(() => {
-    if (closeAllColumns) {
-      setIsExpanded(false);
-    }
-  }, [closeAllColumns]);
+  if (closeAllColumns && isExpanded) {
+    setIsExpanded(false);
+  }
 
   if (isExpanded) {
     setExpandedColumns(

package/src/components/editor/ai/transport/chat-transport.tsx

@@ -13,11 +13,14 @@ import {
 export class StreamingChunkTransport<
   UI_MESSAGE extends UIMessage,
 > extends DefaultChatTransport<UI_MESSAGE> {
+  private onChunkReceived: (chunk: UIMessageChunk) => void;
+
   constructor(
     options: HttpChatTransportInitOptions<UI_MESSAGE> = {},
-    private onChunkReceived: (chunk: UIMessageChunk) => void,
+    onChunkReceived: (chunk: UIMessageChunk) => void,
   ) {
     super(options);
+    this.onChunkReceived = onChunkReceived;
   }
 
   protected override processResponseStream(