Improve running cell visibility and auto-follow streaming console output

Author: axsseldz

examples/third_party/plotly/histogram.py

@@ -8,10 +8,17 @@
 
 import marimo
 
-__generated_with = "0.20.2"
+__generated_with = "0.21.1"
 app = marimo.App(width="medium")
 
 
+@app.cell
+def _():
+    for num in range(1000000000):
+        print(f'number: {num}')
+    return
+
+
 @app.cell
 def _():
     import random
@@ -73,7 +80,6 @@ def _(numeric_hist, values):
         for row in selected_rows[:10]
         if isinstance(row.get("pointIndex"), int)
     ]
-
     return mapped_original_values, selected_rows
 
 

frontend/src/__tests__/__snapshots__/CellStatus.test.tsx.snap

@@ -343,6 +343,11 @@ exports[`CellStatusComponent > renders running state 1`] = `
     data-status="running"
     data-testid="cell-status"
   >
+    <span
+      class="running-state-label"
+    >
+      Running
+    </span>
     <span>
       0ms
     </span>

frontend/src/components/editor/cell/CellStatus.tsx

@@ -176,6 +176,7 @@ export const CellStatusComponent: React.FC<CellStatusComponentProps> = ({
           data-testid="cell-status"
           data-status="running"
         >
+          <span className="running-state-label">Running</span>
           <CellTimer
             startTime={Time.fromSeconds(runStartTimestamp) || Time.now()}
           />

frontend/src/components/editor/cell/cell-status.css

@@ -1,3 +1,14 @@
+@keyframes running-cell-status-pulse {
+  0%,
+  100% {
+    opacity: 0.65;
+  }
+
+  50% {
+    opacity: 1;
+  }
+}
+
 .cell-status-icon {
   margin-top: 4px;
   margin-left: 3px;
@@ -10,6 +21,35 @@
   color: var(--gray-11);
 }
 
+.elapsed-time.running {
+  display: inline-flex;
+  align-items: center;
+  gap: 0.35rem;
+  padding: 0.15rem 0.5rem;
+  border-radius: 999px;
+  color: var(--accent-foreground);
+  border: 1px solid
+    color-mix(in srgb, var(--accent-foreground), transparent 78%);
+  background: color-mix(in srgb, var(--accent), transparent 22%);
+}
+
+.elapsed-time.running::before {
+  content: "";
+  width: 0.45rem;
+  height: 0.45rem;
+  border-radius: 999px;
+  background: currentColor;
+  box-shadow: 0 0 0 0 color-mix(in srgb, currentColor, transparent 100%);
+  animation: running-cell-status-pulse 1600ms ease-in-out infinite;
+}
+
+.running-state-label {
+  font-size: 0.7rem;
+  font-weight: 700;
+  letter-spacing: 0.02em;
+  text-transform: uppercase;
+}
+
 #App.disconnected .elapsed-time {
   visibility: hidden;
 }

frontend/src/components/editor/notebook-cell.tsx

@@ -534,6 +534,7 @@ const EditableCellComponent = ({
     interactive: true,
     "needs-run": needsRun,
     "has-error": cellRuntime.errored,
+    running: cellRuntime.status === "running",
     stopped: cellRuntime.stopped,
     disabled: cellData.config.disabled,
     stale: cellRuntime.status === "disabled-transitively",
@@ -731,6 +732,7 @@ const EditableCellComponent = ({
             <ConsoleOutput
               consoleOutputs={cellRuntime.consoleOutputs}
               stale={consoleOutputStale}
+              running={cellRuntime.status === "running"}
               // Empty name if serialization triggered
               cellName={cellRuntime.serialization ? "_" : cellData.name}
               onRefactorWithAI={handleRefactorWithAI}
@@ -1036,6 +1038,7 @@ const SetupCellComponent = ({
     interactive: true,
     "needs-run": needsRun,
     "has-error": cellRuntime.errored,
+    running: cellRuntime.status === "running",
     stopped: cellRuntime.stopped,
   });
 
@@ -1191,6 +1194,7 @@ const SetupCellComponent = ({
             <ConsoleOutput
               consoleOutputs={cellRuntime.consoleOutputs}
               stale={consoleOutputStale}
+              running={cellRuntime.status === "running"}
               // Don't show name
               cellName={"_"}
               onRefactorWithAI={handleRefactorWithAI}

frontend/src/components/editor/output/console/ConsoleOutput.tsx

@@ -85,6 +85,7 @@ interface Props {
   className?: string;
   consoleOutputs: WithResponse<OutputMessage>[];
   stale: boolean;
+  running?: boolean;
   debuggerActive: boolean;
   onRefactorWithAI?: OnRefactorWithAI;
   onClear?: () => void;
@@ -101,6 +102,8 @@ export const ConsoleOutput = (props: Props) => {
 
 const ConsoleOutputInternal = (props: Props): React.ReactNode => {
   const ref = React.useRef<HTMLDivElement>(null);
+  const shouldFollowOutputRef = useRef(true);
+  const prevRenderedOutputCountRef = useRef(0);
   const { wrapText, setWrapText } = useWrapText();
   const [isExpanded, setIsExpanded] = useExpandedConsoleOutput(props.cellId);
   const [stdinValue, setStdinValue] = React.useState("");
@@ -111,6 +114,7 @@ const ConsoleOutputInternal = (props: Props): React.ReactNode => {
   const {
     consoleOutputs: rawConsoleOutputs,
     stale,
+    running = false,
     cellName,
     cellId,
     onSubmitDebugger,
@@ -141,27 +145,37 @@ const ConsoleOutputInternal = (props: Props): React.ReactNode => {
   // Detect overflow on resize
   const isOverflowing = useOverflowDetection(ref, hasOutputs);
 
-  // Keep scroll at the bottom if it is within 120px of the bottom,
-  // so when we add new content, it will lock to the bottom
-  //
-  // We use flex flex-col-reverse to handle this, but it doesn't
-  // always work perfectly when moved form the bottom and back.
+  const isNearBottom = (el: HTMLDivElement) => {
+    const threshold = 24;
+    const scrollOffset = el.scrollHeight - el.clientHeight;
+    const distanceFromBottom = scrollOffset - el.scrollTop;
+    return distanceFromBottom <= threshold;
+  };
+
+  // Follow newly appended console output while the cell is actively running,
+  // but stop following when the user scrolls away from the bottom.
   useLayoutEffect(() => {
     const el = ref.current;
     if (!el) {
       return;
     }
-    // N.B. This won't handle large jumps in the scroll position
-    // if there is a lot of content added at once.
-    // This is 'good enough' for now.
-    const threshold = 120;
 
-    const scrollOffset = el.scrollHeight - el.clientHeight;
-    const distanceFromBottom = scrollOffset - el.scrollTop;
-    if (distanceFromBottom < threshold) {
-      el.scrollTop = scrollOffset;
+    if (!hasOutputs) {
+      shouldFollowOutputRef.current = true;
+      prevRenderedOutputCountRef.current = 0;
+      return;
     }
-  });
+
+    const appendedOutput =
+      consoleOutputs.length > prevRenderedOutputCountRef.current;
+    prevRenderedOutputCountRef.current = consoleOutputs.length;
+
+    if (!running || !appendedOutput || !shouldFollowOutputRef.current) {
+      return;
+    }
+
+    el.scrollTop = el.scrollHeight - el.clientHeight;
+  }, [consoleOutputs, hasOutputs, running]);
 
   if (!hasOutputs && isInternalCellName(cellName)) {
     return null;
@@ -236,6 +250,9 @@ const ConsoleOutputInternal = (props: Props): React.ReactNode => {
         {...selectAllProps}
         // biome-ignore lint/a11y/noNoninteractiveTabindex: Needed to capture keypress events
         tabIndex={0}
+        onScroll={(e) => {
+          shouldFollowOutputRef.current = isNearBottom(e.currentTarget);
+        }}
         className={cn(
           "console-output-area overflow-hidden rounded-b-lg flex flex-col-reverse w-full gap-1 focus:outline-hidden",
           stale && "marimo-output-stale",

frontend/src/components/editor/output/console/__tests__/ConsoleOutput.test.tsx

@@ -28,6 +28,7 @@ describe("ConsoleOutput integration", () => {
     cellName: "test_cell",
     consoleOutputs: [] as WithResponse<OutputMessage>[],
     stale: false,
+    running: false,
     debuggerActive: false,
     onSubmitDebugger: () => {
       // noop
@@ -59,6 +60,7 @@ describe("ConsoleOutput pdb history", () => {
     cellName: "test_cell",
     consoleOutputs: [] as WithResponse<OutputMessage>[],
     stale: false,
+    running: false,
     debuggerActive: false,
     onSubmitDebugger: vi.fn(),
   };
@@ -219,6 +221,7 @@ describe("ConsoleOutput debounced clearing", () => {
     cellName: "test_cell",
     consoleOutputs: [] as WithResponse<OutputMessage>[],
     stale: false,
+    running: false,
     debuggerActive: false,
     onSubmitDebugger: vi.fn(),
   };
@@ -284,3 +287,101 @@ describe("ConsoleOutput debounced clearing", () => {
     expect(screen.queryByText("old output")).not.toBeInTheDocument();
   });
 });
+
+describe("ConsoleOutput auto-scroll", () => {
+  const createOutput = (data: string): WithResponse<OutputMessage> => ({
+    channel: "stdout",
+    mimetype: "text/plain",
+    data,
+    timestamp: 0,
+    response: undefined,
+  });
+
+  const defaultProps = {
+    cellId: cellId("cell-1"),
+    cellName: "test_cell",
+    consoleOutputs: [] as WithResponse<OutputMessage>[],
+    stale: false,
+    running: true,
+    debuggerActive: false,
+    onSubmitDebugger: vi.fn(),
+  };
+
+  const setScrollMetrics = (
+    element: HTMLElement,
+    {
+      clientHeight,
+      scrollHeight,
+      scrollTop,
+    }: { clientHeight: number; scrollHeight: number; scrollTop: number },
+  ) => {
+    Object.defineProperty(element, "clientHeight", {
+      configurable: true,
+      value: clientHeight,
+    });
+    Object.defineProperty(element, "scrollHeight", {
+      configurable: true,
+      value: scrollHeight,
+    });
+    Object.defineProperty(element, "scrollTop", {
+      configurable: true,
+      writable: true,
+      value: scrollTop,
+    });
+  };
+
+  it("follows newly appended output while running when already at the bottom", () => {
+    const { rerender } = renderWithProvider(
+      <ConsoleOutput
+        {...defaultProps}
+        consoleOutputs={[createOutput("line 1")]}
+      />,
+    );
+
+    const consoleArea = screen.getByTestId("console-output-area");
+    setScrollMetrics(consoleArea, {
+      clientHeight: 100,
+      scrollHeight: 300,
+      scrollTop: 200,
+    });
+
+    rerender(
+      <TooltipProvider>
+        <ConsoleOutput
+          {...defaultProps}
+          consoleOutputs={[createOutput("line 1"), createOutput("line 2")]}
+        />
+      </TooltipProvider>,
+    );
+
+    expect(consoleArea.scrollTop).toBe(200);
+  });
+
+  it("does not auto-scroll when the user has scrolled away from the bottom", () => {
+    const { rerender } = renderWithProvider(
+      <ConsoleOutput
+        {...defaultProps}
+        consoleOutputs={[createOutput("line 1")]}
+      />,
+    );
+
+    const consoleArea = screen.getByTestId("console-output-area");
+    setScrollMetrics(consoleArea, {
+      clientHeight: 100,
+      scrollHeight: 300,
+      scrollTop: 60,
+    });
+    fireEvent.scroll(consoleArea);
+
+    rerender(
+      <TooltipProvider>
+        <ConsoleOutput
+          {...defaultProps}
+          consoleOutputs={[createOutput("line 1"), createOutput("line 2")]}
+        />
+      </TooltipProvider>,
+    );
+
+    expect(consoleArea.scrollTop).toBe(60);
+  });
+});

frontend/src/components/editor/renderers/vertical-layout/vertical-layout.tsx

@@ -418,6 +418,7 @@ const VerticalCell = memo(
           <ConsoleOutput
             consoleOutputs={consoleOutputs}
             stale={outputStale}
+            running={false}
             cellName={name}
             onSubmitDebugger={() => null}
             cellId={cellId}

frontend/src/components/scratchpad/scratchpad.tsx

@@ -289,6 +289,7 @@ export const ScratchPad: React.FC = () => {
                 consoleOutputs={consoleOutputs}
                 className="overflow-auto"
                 stale={false}
+                running={status === "running"}
                 cellName={DEFAULT_CELL_NAME}
                 onSubmitDebugger={Functions.NOOP}
                 cellId={cellId}