@marimo-team/frontend 0.22.5-dev9 → 0.22.6-dev1

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

@@ -2,7 +2,13 @@
 "use no memo";
 
 import type { Column, Table } from "@tanstack/react-table";
-import { FilterIcon, MinusIcon, TextIcon, XIcon } from "lucide-react";
+import {
+  EllipsisIcon,
+  FilterIcon,
+  MinusIcon,
+  TextIcon,
+  XIcon,
+} from "lucide-react";
 import { useMemo, useRef, useState } from "react";
 import { useLocale } from "react-aria";
 import {
@@ -69,7 +75,7 @@ interface DataTableColumnHeaderProps<
 > extends React.HTMLAttributes<HTMLDivElement> {
   column: Column<TData, TValue>;
   header: React.ReactNode;
-  justify?: "left" | "center" | "right";
+  subheader?: React.ReactNode;
   calculateTopKRows?: CalculateTopKRows;
   table?: Table<TData>;
 }
@@ -77,7 +83,7 @@ interface DataTableColumnHeaderProps<
 export const DataTableColumnHeader = <TData, TValue>({
   column,
   header,
-  justify,
+  subheader,
   className,
   calculateTopKRows,
   table,
@@ -92,49 +98,51 @@ export const DataTableColumnHeader = <TData, TValue>({
 
   // No sorting or filtering
   if (!column.getCanSort() && !column.getCanFilter()) {
-    return <div className={cn(className)}>{header}</div>;
+    return (
+      <div className={cn(className)}>
+        {header}
+        {subheader}
+      </div>
+    );
   }
 
   const hasFilter = column.getFilterValue() !== undefined;
-  const hideIcon = !column.getIsSorted() && !hasFilter;
 
   return (
     <>
-      <DropdownMenu modal={false}>
-        <DropdownMenuTrigger asChild={true}>
-          <div
-            className={cn(
-              "group flex items-center my-1 space-between w-full select-none gap-2 border hover:border-border border-transparent hover:bg-(--slate-3) data-[state=open]:bg-(--slate-3) data-[state=open]:border-border rounded px-1 -mx-1",
-              justify === "right" && "flex-row-reverse",
-              className,
-            )}
-            data-testid="data-table-sort-button"
-          >
-            <span className="flex-1">{header}</span>
-            <span
-              className={cn(
-                "h-5 py-1 px-1",
-                hideIcon &&
-                  "invisible group-hover:visible data-[state=open]:visible",
-              )}
-            >
-              {renderSortFilterIcon(column)}
-            </span>
-          </div>
-        </DropdownMenuTrigger>
-        <DropdownMenuContent align="start">
-          {renderDataType(column)}
-          {renderSorts(column, table)}
-          {renderCopyColumn(column)}
-          {renderColumnPinning(column)}
-          {renderColumnWrapping(column)}
-          {renderFormatOptions(column, locale)}
-          <DropdownMenuSeparator />
-          {renderMenuItemFilter(column)}
-          {renderFilterByValues(column, setIsFilterValueOpen)}
-          {hasFilter && <ClearFilterMenuItem column={column} />}
-        </DropdownMenuContent>
-      </DropdownMenu>
+      <div
+        className={cn("group flex flex-col my-1 w-full select-none", className)}
+      >
+        <div className="flex items-center gap-1">
+          <span>{header}</span>
+          {column.getCanSort() && <SortButton column={column} />}
+          <DropdownMenu modal={false}>
+            <DropdownMenuTrigger asChild={true}>
+              <button
+                type="button"
+                className="inline-flex items-center justify-center h-5 w-5 rounded hover:bg-(--slate-4) text-muted-foreground opacity-0 group-hover:opacity-100 focus:opacity-100 group-focus-within:opacity-100 data-[state=open]:opacity-100 data-[state=open]:text-accent-foreground"
+                aria-label="Column options"
+                data-testid="data-table-column-menu-button"
+              >
+                <EllipsisIcon className="h-3.5 w-3.5" />
+              </button>
+            </DropdownMenuTrigger>
+            <DropdownMenuContent align="start">
+              {renderDataType(column)}
+              {renderSorts(column, table)}
+              {renderCopyColumn(column)}
+              {renderColumnPinning(column)}
+              {renderColumnWrapping(column)}
+              {renderFormatOptions(column, locale)}
+              <DropdownMenuSeparator />
+              {renderMenuItemFilter(column)}
+              {renderFilterByValues(column, setIsFilterValueOpen)}
+              {hasFilter && <ClearFilterMenuItem column={column} />}
+            </DropdownMenuContent>
+          </DropdownMenu>
+        </div>
+        {subheader}
+      </div>
       {isFilterValueOpen && (
         <PopoverFilterByValues
           setIsFilterValueOpen={setIsFilterValueOpen}
@@ -146,28 +154,45 @@ export const DataTableColumnHeader = <TData, TValue>({
   );
 };
 
-export const DataTableColumnHeaderWithSummary = <TData, TValue>({
+const SortButton = <TData, TValue>({
   column,
-  header,
-  summary,
-  className,
-}: DataTableColumnHeaderProps<TData, TValue> & {
-  summary: React.ReactNode;
+}: {
+  column: Column<TData, TValue>;
 }) => {
+  const sortDirection = column.getIsSorted();
+
+  const handleClick = (e: React.MouseEvent) => {
+    e.stopPropagation();
+    if (!sortDirection) {
+      column.toggleSorting(false, true); // asc
+    } else if (sortDirection === "asc") {
+      column.toggleSorting(true, true); // desc
+    } else {
+      column.clearSorting();
+    }
+  };
+
   return (
-    <div
+    <button
+      type="button"
+      onClick={handleClick}
       className={cn(
-        "flex flex-col h-full pt-0.5 pb-3 justify-between items-start",
-        className,
+        "inline-flex items-center justify-center h-5 w-5 rounded hover:bg-(--slate-4)",
+        sortDirection
+          ? "text-accent-foreground"
+          : "text-muted-foreground opacity-0 group-hover:opacity-100 focus:opacity-100 group-focus-within:opacity-100",
       )}
+      aria-label={
+        sortDirection === "asc"
+          ? "Sorted ascending, click to sort descending"
+          : sortDirection === "desc"
+            ? "Sorted descending, click to clear sort"
+            : "Sort column ascending"
+      }
+      data-testid="data-table-sort-button"
     >
-      <DataTableColumnHeader
-        column={column}
-        header={header}
-        className={className}
-      />
-      {summary}
-    </div>
+      {renderSortFilterIcon(column)}
+    </button>
   );
 };
 

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

@@ -206,41 +206,30 @@ export function generateColumns<T>({
             </div>
           ) : null;
 
-        const justify = getJustify(key);
-
-        const headerWithType = (
-          <div
+        const headerName = (
+          <span
             className={cn(
-              "flex flex-col",
-              justify === "center" && "items-center",
-              justify === "right" && "items-end",
+              "font-bold",
+              headerTitle && "underline decoration-dotted",
             )}
           >
-            <span
-              className={cn(
-                "font-bold",
-                headerTitle && "underline decoration-dotted",
-              )}
-            >
-              {key === "" ? " " : key}
-            </span>
-            {dtypeHeader}
-          </div>
+            {key === "" ? " " : key}
+          </span>
         );
 
         const headerWithTooltip = headerTitle ? (
           <Tooltip content={headerTitle} delayDuration={300}>
-            {headerWithType}
+            {headerName}
           </Tooltip>
         ) : (
-          headerWithType
+          headerName
         );
 
         const dataTableColumnHeader = (
           <DataTableColumnHeader
             header={headerWithTooltip}
+            subheader={dtypeHeader}
             column={column}
-            justify={justify}
             calculateTopKRows={calculateTopKRows}
             table={table}
           />
@@ -255,8 +244,6 @@ export function generateColumns<T>({
           <div
             className={cn(
               "flex flex-col h-full pt-0.5 pb-3 justify-between items-start",
-              justify === "center" && "items-center",
-              justify === "right" && "items-end",
             )}
           >
             {dataTableColumnHeader}
@@ -283,13 +270,13 @@ export function generateColumns<T>({
 
         const dataType = column.columnDef.meta?.dataType;
         const isNumeric = dataType === "number" || dataType === "integer";
-        const cellStyles = getCellStyleClass(
+        const cellStyles = getCellStyleClass({
           justify,
           wrapped,
           canSelectCell,
-          isCellSelected,
+          isSelected: isCellSelected,
           isNumeric,
-        );
+        });
 
         const renderedCell = renderCellValue({
           column,
@@ -448,13 +435,19 @@ function getFilterTypeForFieldType(
   }
 }
 
-function getCellStyleClass(
-  justify: "left" | "center" | "right" | undefined,
-  wrapped: boolean | undefined,
-  canSelectCell: boolean,
-  isSelected: boolean,
-  isNumeric?: boolean,
-): string {
+function getCellStyleClass({
+  justify = "left",
+  wrapped,
+  canSelectCell,
+  isSelected,
+  isNumeric = false,
+}: {
+  justify: "left" | "center" | "right" | undefined;
+  wrapped: boolean | undefined;
+  canSelectCell: boolean;
+  isSelected: boolean;
+  isNumeric?: boolean;
+}): string {
   return cn(
     canSelectCell && "cursor-pointer",
     isSelected &&

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

@@ -47,6 +47,7 @@ import { DataTableBody, renderTableHeader } from "./renderers";
 import { TableBottomBar } from "./TableBottomBar";
 import { TableTopBar } from "./TableTopBar";
 import {
+  AUTO_WIDTH_MAX_COLUMNS,
   type DataTableSelection,
   MIN_ROWS_TO_VIRTUALIZE,
   type TooManyRows,
@@ -300,7 +301,13 @@ const DataTableInternal = <TData,>({
             isAnyPanelOpen={isAnyPanelOpen}
             downloadAs={downloadAs}
           />
-          <Table className="relative" ref={tableRef}>
+          <Table
+            className={cn(
+              "relative",
+              columns.length <= AUTO_WIDTH_MAX_COLUMNS ? "w-auto" : "w-full",
+            )}
+            ref={tableRef}
+          >
             {showLoadingBar && (
               <thead className="absolute top-0 left-0 h-[3px] w-1/2 bg-primary animate-slide" />
             )}

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

@@ -27,7 +27,7 @@ import { DataTableContextMenu } from "./context-menu";
 import { CellRangeSelectionIndicator } from "./range-focus/cell-selection-indicator";
 import { useCellRangeSelection } from "./range-focus/use-cell-range-selection";
 import { useScrollIntoViewOnFocus } from "./range-focus/use-scroll-into-view";
-import { TABLE_ROW_HEIGHT_PX } from "./types";
+import { AUTO_WIDTH_MAX_COLUMNS, TABLE_ROW_HEIGHT_PX } from "./types";
 import { stringifyUnknownValue } from "./utils";
 
 export function renderTableHeader<TData>(
@@ -46,7 +46,7 @@ export function renderTableHeader<TData>(
           <TableHead
             key={header.id}
             className={cn(
-              "h-auto min-h-10 whitespace-pre align-top",
+              "h-auto min-h-10 whitespace-pre align-top border-r border-r-border/75",
               className,
             )}
             style={style}
@@ -69,6 +69,13 @@ export function renderTableHeader<TData>(
         {renderHeaderGroup(table.getLeftHeaderGroups())}
         {renderHeaderGroup(table.getCenterHeaderGroups())}
         {renderHeaderGroup(table.getRightHeaderGroups())}
+        {table.getAllColumns().length <= AUTO_WIDTH_MAX_COLUMNS && (
+          <th
+            className="w-full border-0"
+            aria-hidden="true"
+            role="presentation"
+          />
+        )}
       </TableRow>
     </TableHeader>
   );
@@ -163,7 +170,7 @@ export const DataTableBody = <TData,>({
           {...getCellDomProps(cell.id)}
           key={cell.id}
           className={cn(
-            "whitespace-pre truncate max-w-[300px] outline-hidden",
+            "whitespace-pre truncate max-w-[300px] outline-hidden border-r border-r-border/75",
             cell.column.getColumnWrapping &&
               cell.column.getColumnWrapping?.() === "wrap" &&
               COLUMN_WRAPPING_STYLES,
@@ -230,15 +237,21 @@ export const DataTableBody = <TData,>({
         {renderCells(row.getLeftVisibleCells())}
         {renderCells(row.getCenterVisibleCells())}
         {renderCells(row.getRightVisibleCells())}
+        {columns.length <= AUTO_WIDTH_MAX_COLUMNS && (
+          <td className="border-0" aria-hidden="true" role="presentation" />
+        )}
       </TableRow>
     );
   };
 
+  const hasFillerColumn = columns.length <= AUTO_WIDTH_MAX_COLUMNS;
+  const totalColSpan = columns.length + (hasFillerColumn ? 1 : 0);
+
   const renderRows = () => {
     if (rows.length === 0) {
       return (
         <TableRow>
-          <TableCell colSpan={columns.length} className="h-24 text-center">
+          <TableCell colSpan={totalColSpan} className="h-24 text-center">
             No results.
           </TableCell>
         </TableRow>
@@ -255,7 +268,7 @@ export const DataTableBody = <TData,>({
               data-virtual-spacer=""
               style={{ height: virtualItems[0].start }}
             >
-              <td colSpan={columns.length} />
+              <td colSpan={totalColSpan} />
             </tr>
           )}
           {virtualItems.map((vItem) => renderRow(rows[vItem.index]))}
@@ -266,7 +279,7 @@ export const DataTableBody = <TData,>({
                 height: totalSize - (virtualItems.at(-1)?.end ?? totalSize),
               }}
             >
-              <td colSpan={columns.length} />
+              <td colSpan={totalColSpan} />
             </tr>
           )}
         </>

package/src/components/data-table/types.ts

@@ -16,6 +16,10 @@ declare module "@tanstack/react-table" {
 export const TABLE_ROW_HEIGHT_PX = 24;
 export const TABLE_HEADER_HEIGHT_PX = 40;
 
+// Below this column count, the table uses w-auto with a filler column
+// to prevent columns from stretching unnecessarily
+export const AUTO_WIDTH_MAX_COLUMNS = 4;
+
 // Default number of visible rows when virtualizing without an explicit maxHeight.
 export const DEFAULT_VIRTUAL_ROWS = 15;
 

package/src/components/editor/Output.tsx

@@ -367,7 +367,7 @@ export const OutputArea = React.memo(
     forceExpand,
     className,
   }: OutputAreaProps) => {
-    if (output === null) {
+    if (output == null) {
       return null;
     }
     if (output.channel === "output" && output.data === "") {

package/src/components/editor/__tests__/Output.test.tsx

@@ -1,7 +1,9 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { render, screen } from "@testing-library/react";
 import { describe, expect, it } from "vitest";
-import { OutputRenderer } from "../Output";
+import { cellId } from "@/__tests__/branded";
+import { TooltipProvider } from "@/components/ui/tooltip";
+import { OutputArea, OutputRenderer } from "../Output";
 
 describe("OutputRenderer renderFallback prop", () => {
   it("should use renderFallback for unsupported mimetypes", () => {
@@ -65,6 +67,39 @@ describe("OutputRenderer renderFallback prop", () => {
   });
 });
 
+describe("OutputArea null/undefined handling", () => {
+  it("should render null when output is null", () => {
+    const { container } = render(
+      <TooltipProvider>
+        <OutputArea
+          output={null}
+          cellId={cellId("test")}
+          stale={false}
+          loading={false}
+          allowExpand={true}
+        />
+      </TooltipProvider>,
+    );
+    expect(container.innerHTML).toBe("");
+  });
+
+  it("should render null when output is undefined", () => {
+    const { container } = render(
+      <TooltipProvider>
+        <OutputArea
+          // @ts-expect-error -- testing runtime safety for undefined output
+          output={undefined}
+          cellId={cellId("test")}
+          stale={false}
+          loading={false}
+          allowExpand={true}
+        />
+      </TooltipProvider>,
+    );
+    expect(container.innerHTML).toBe("");
+  });
+});
+
 describe("OutputRenderer image and SVG rendering", () => {
   const plainSvgString =
     '<svg><rect x="0" y="0" width="10" height="10"></rect></svg>';

package/src/core/cells/__tests__/cells.test.ts

@@ -1209,6 +1209,47 @@ describe("cell reducer", () => {
     ]);
   });
 
+  it("does not crash when setStdinResponse has out-of-bounds outputIndex", () => {
+    const STDOUT: OutputMessage = {
+      channel: "stdout",
+      mimetype: "text/plain",
+      data: "hello!",
+      timestamp: 1,
+    };
+
+    // Set the cell to running with a console output
+    actions.prepareForRun({ cellId: firstCellId });
+    actions.handleCellMessage({
+      cell_id: firstCellId,
+      output: undefined,
+      console: null,
+      status: "running",
+      stale_inputs: null,
+      timestamp: new Date(20).getTime() as Seconds,
+    });
+    actions.handleCellMessage({
+      cell_id: firstCellId,
+      output: undefined,
+      console: STDOUT,
+      status: undefined,
+      stale_inputs: null,
+      timestamp: new Date(22).getTime() as Seconds,
+    });
+
+    // Try to set stdin response with an out-of-bounds index
+    // This should not crash - it should return state unchanged
+    actions.setStdinResponse({
+      response: "test",
+      cellId: firstCellId,
+      outputIndex: 999,
+    });
+
+    // Cell state should be unchanged
+    const cell = cells[0];
+    expect(cell.consoleOutputs).toHaveLength(1);
+    expect(cell.consoleOutputs[0]).toMatchObject(STDOUT);
+  });
+
   it("can receive console when the cell is idle and will clear when starts again", () => {
     const OLD_STDOUT: OutputMessage = {
       channel: "stdout",

package/src/core/cells/__tests__/collapseConsoleOutputs.test.ts

@@ -241,6 +241,44 @@ describe("collapseConsoleOutputs", () => {
     expect(result[2].data).toBe("<pre>E\nF\nG\nH\n</pre>");
   });
 
+  it("should not crash when truncating with a single output at the limit boundary", () => {
+    // Create outputs that push truncation to the exact boundary
+    const consoleOutputs: OutputMessage[] = [
+      {
+        mimetype: "text/html",
+        channel: "output",
+        data: "<div>html1</div>",
+        timestamp: 0,
+      },
+      {
+        mimetype: "text/html",
+        channel: "output",
+        data: "<div>html2</div>",
+        timestamp: 0,
+      },
+    ];
+    // With limit=1, truncation must handle edge cases gracefully
+    const result = collapseConsoleOutputs(consoleOutputs, 1);
+    expect(result[0].data).toContain("Streaming output truncated");
+  });
+
+  it("should handle truncation when cutoff indexes past the end of the array", () => {
+    // With maxLines=0, the truncation loop never runs, causing cutoff
+    // to index past the array. This exercises the `output == null`
+    // defensive branch in truncateHead().
+    const consoleOutputs: OutputMessage[] = [
+      {
+        mimetype: "text/html",
+        channel: "output",
+        data: "<div>content</div>",
+        timestamp: 0,
+      },
+    ];
+    const result = collapseConsoleOutputs(consoleOutputs, 0);
+    expect(result).toHaveLength(1);
+    expect(result[0].data).toContain("Streaming output truncated");
+  });
+
   describe("ANSI escape sequences", () => {
     it("should handle cursor movement with collapse", () => {
       const consoleOutputs: OutputMessage[] = [

package/src/core/cells/cells.ts

@@ -947,7 +947,7 @@ const {
       cellReducer: (cell) => {
         const consoleOutputs = [...cell.consoleOutputs];
         const stdinOutput = consoleOutputs[outputIndex];
-        if (stdinOutput.channel !== "stdin") {
+        if (stdinOutput == null || stdinOutput.channel !== "stdin") {
           Logger.warn("Expected stdin output");
           return cell;
         }

package/src/core/cells/collapseConsoleOutputs.tsx

@@ -108,6 +108,9 @@ function truncateHead(consoleOutputs: OutputMessage[], limit: number) {
     timestamp: -1,
   };
   const output = consoleOutputs[cutoff];
+  if (output == null) {
+    return [warningOutput, ...consoleOutputs.slice(cutoff + 1)];
+  }
   if (output.mimetype === "text/plain") {
     invariant(typeof output.data === "string", "expected string");
     const outputLines = output.data.split("\n");

package/src/core/cells/document-changes.ts

@@ -25,6 +25,7 @@ import type { NotebookDocumentTransactionRequest } from "../network/types";
 import { store } from "../state/jotai";
 import type { CellActions, NotebookState } from "./cells";
 import type { CellId } from "./ids";
+import { SCRATCH_CELL_ID } from "./ids";
 import type { CellData } from "./types";
 
 export type DocumentChange =
@@ -572,10 +573,21 @@ const flushChanges = debounce(() => {
   void getRequestClient().sendDocumentTransaction({ changes });
 }, 400);
 
+function isScratchChange(change: DocumentChange): boolean {
+  if ("cellId" in change && change.cellId === SCRATCH_CELL_ID) {
+    return true;
+  }
+  return false;
+}
+
 function enqueue(change: DocumentChange) {
   if (store.get(kioskModeAtom)) {
     return;
   }
+  // The scratchpad cell is local-only — don't sync it to the document.
+  if (isScratchChange(change)) {
+    return;
+  }
   pendingChanges.push(change);
   flushChanges();
 }