fix: flush console output buffer before marking cell idle (#9164)

Console messages (stdout/stderr) are batched by a background thread with
a 10ms buffer for performance. Without an explicit flush, messages
arrive
at the frontend after the cell is marked idle and completed-run is sent.
A subsequent run clears the console (console=[]) before the user sees
the previous output.

Add a FlushMarker sentinel to the console message queue. When the
buffered writer encounters it, it flushes immediately and signals the
caller. Implement flush() on ThreadSafeStdout/ThreadSafeStderr
(previously
no-ops with a TODO) and add a _flush_console post-execution hook that
runs just before _set_status_idle.

Closes #4821

Author: mscolnick

marimo/_messaging/console_output_worker.py

@@ -1,8 +1,9 @@
 # Copyright 2026 Marimo. All rights reserved.
 from __future__ import annotations
 
+import threading
 import time
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Literal
 
 from marimo._messaging.cell_output import CellChannel, CellOutput
@@ -29,6 +30,13 @@ class ConsoleMsg:
     mimetype: ConsoleMimeType
 
 
+@dataclass
+class FlushMarker:
+    """Sentinel that tells the buffered writer to flush immediately."""
+
+    done: threading.Event = field(default_factory=threading.Event)
+
+
 def _write_console_output(
     stream: Stream,
     stream_type: StreamT,
@@ -70,8 +78,24 @@ def _add_output_to_buffer(
         outputs_buffered_per_cell[cell_id] = [console_output]
 
 
+def _flush_outputs(
+    outputs_buffered_per_cell: dict[CellId_t, list[ConsoleMsg]],
+    stream: Stream,
+) -> None:
+    for cell_id, buffer in outputs_buffered_per_cell.items():
+        for output in buffer:
+            _write_console_output(
+                stream,
+                output.stream,
+                cell_id,
+                output.data,
+                output.mimetype,
+            )
+    outputs_buffered_per_cell.clear()
+
+
 def buffered_writer(
-    msg_queue: deque[ConsoleMsg | None],
+    msg_queue: deque[ConsoleMsg | FlushMarker | None],
     stream: Stream,
     cv: Condition,
 ) -> None:
@@ -84,6 +108,7 @@ def buffered_writer(
     was noticeably faster than the builtin queue.Queue in testing.)
 
     A `None` passed to `msg_queue` signals the writer should terminate.
+    A `FlushMarker` forces an immediate flush and signals the caller.
     """
 
     # only have a non-None timer when there's at least one output buffered
@@ -93,6 +118,7 @@ def buffered_writer(
 
     outputs_buffered_per_cell: dict[CellId_t, list[ConsoleMsg]] = {}
     while True:
+        flush_marker: FlushMarker | None = None
         with cv:
             # We wait for messages until the timer (if any) expires
             while timer is None or timer > 0:
@@ -106,23 +132,21 @@ def buffered_writer(
                     msg = msg_queue.popleft()
                     if msg is None:
                         return
+                    if isinstance(msg, FlushMarker):
+                        flush_marker = msg
+                        break
                     _add_output_to_buffer(msg, outputs_buffered_per_cell)
+                if flush_marker is not None:
+                    break
                 if outputs_buffered_per_cell and timer is None:
                     # start the timeout timer
                     timer = TIMEOUT_S
                 elif timer is not None:
                     time_waited = time.time() - time_started_waiting
                     timer -= time_waited
 
-        # the timer has expired: flush the outputs
-        for cell_id, buffer in outputs_buffered_per_cell.items():
-            for output in buffer:
-                _write_console_output(
-                    stream,
-                    output.stream,
-                    cell_id,
-                    output.data,
-                    output.mimetype,
-                )
-        outputs_buffered_per_cell = {}
+        # the timer has expired (or a flush was requested): flush the outputs
+        _flush_outputs(outputs_buffered_per_cell, stream)
+        if flush_marker is not None:
+            flush_marker.done.set()
         timer = None

marimo/_messaging/streams.py

@@ -14,7 +14,11 @@
 
 from marimo import _loggers
 from marimo._messaging.cell_output import CellChannel
-from marimo._messaging.console_output_worker import ConsoleMsg, buffered_writer
+from marimo._messaging.console_output_worker import (
+    ConsoleMsg,
+    FlushMarker,
+    buffered_writer,
+)
 from marimo._messaging.mimetypes import ConsoleMimeType
 from marimo._messaging.types import (
     KernelMessage,
@@ -32,6 +36,12 @@
 LOGGER = _loggers.marimo_logger()
 
 
+# Maximum time to block waiting for the buffered console writer to flush.
+# The flush hook runs on the hot path between cell execution and idle, so
+# we bound the wait even though flushes normally complete in <10ms.
+_FLUSH_CONSOLE_TIMEOUT_S = 5.0
+
+
 # Byte limits on outputs exist for two reasons
 #
 # 1. We use a multiprocessing.Connection object to send outputs from
@@ -106,7 +116,9 @@ def __init__(
         if self.redirect_console:
             # Console outputs are buffered
             self.console_msg_cv = threading.Condition(threading.Lock())
-            self.console_msg_queue: deque[ConsoleMsg | None] = deque()
+            self.console_msg_queue: deque[ConsoleMsg | FlushMarker | None] = (
+                deque()
+            )
             self.buffered_console_thread = threading.Thread(
                 target=buffered_writer,
                 args=(self.console_msg_queue, self, self.console_msg_cv),
@@ -133,14 +145,41 @@ def write(self, data: KernelMessage) -> None:
                     e,
                 )
 
+    def flush_console(self) -> None:
+        """Force the buffered console writer to flush immediately.
+
+        Blocks until all pending console messages have been sent to the
+        frontend, or until a short timeout elapses if the writer thread
+        is no longer alive.  This ensures that stderr/stdout output
+        produced during cell execution is delivered before the cell is
+        marked idle.
+        """
+        if not self.redirect_console:
+            return
+        # If the buffered writer isn't alive (e.g., shutdown in progress),
+        # enqueuing a marker would never be drained. Bail out early.
+        if not self.buffered_console_thread.is_alive():
+            return
+        marker = FlushMarker()
+        with self.console_msg_cv:
+            self.console_msg_queue.append(marker)
+            self.console_msg_cv.notify()
+        # Bounded wait: if the writer dies or stalls, don't block the
+        # caller indefinitely.
+        if not marker.done.wait(timeout=_FLUSH_CONSOLE_TIMEOUT_S):
+            LOGGER.warning(
+                "Timed out waiting for console flush after %ss",
+                _FLUSH_CONSOLE_TIMEOUT_S,
+            )
+
     def stop(self) -> None:
         """Teardown resources created by the stream."""
         # Sending `None` through the queue signals the console thread to exit.
         # We don't join the thread in case its processing outputs still; don't
         # want to block the entire program.
         if self.redirect_console:
-            self.console_msg_queue.append(None)
             with self.console_msg_cv:
+                self.console_msg_queue.append(None)
                 self.console_msg_cv.notify()
 
 
@@ -263,8 +302,7 @@ def seekable(self) -> bool:
         return False
 
     def flush(self) -> None:
-        # TODO(akshayka): maybe force the buffered writer to write
-        return
+        self._stream.flush_console()
 
     def _write_with_mimetype(
         self, data: str, mimetype: ConsoleMimeType
@@ -280,15 +318,15 @@ def _write_with_mimetype(
                 "Warning: marimo truncated a very large console output.\n"
             )
             data = data[: int(max_bytes)] + " ... "
-        self._stream.console_msg_queue.append(
-            ConsoleMsg(
-                stream=CellChannel.STDOUT,
-                cell_id=self._stream.cell_id,
-                data=data,
-                mimetype=mimetype,
-            )
-        )
         with self._stream.console_msg_cv:
+            self._stream.console_msg_queue.append(
+                ConsoleMsg(
+                    stream=CellChannel.STDOUT,
+                    cell_id=self._stream.cell_id,
+                    data=data,
+                    mimetype=mimetype,
+                )
+            )
             self._stream.console_msg_cv.notify()
         return len(data)
 
@@ -338,8 +376,7 @@ def seekable(self) -> bool:
         return False
 
     def flush(self) -> None:
-        # TODO(akshayka): maybe force the buffered writer to write
-        return
+        self._stream.flush_console()
 
     def _write_with_mimetype(
         self, data: str, mimetype: ConsoleMimeType

marimo/_messaging/types.py

@@ -25,6 +25,10 @@ class Stream(abc.ABC):
     def write(self, data: KernelMessage) -> None:
         pass
 
+    def flush_console(self) -> None:
+        """Flush buffered console output, if any."""
+        return
+
     def stop(self) -> None:
         """Tear down resources, if any."""
         return

marimo/_runtime/runner/hooks_post_execution.py

@@ -494,6 +494,27 @@ def _reset_matplotlib_context(
         exec("__marimo__._output.mpl.close_figures()", ctx.glbls)
 
 
+@kernel_tracer.start_as_current_span("flush_console")
+def _flush_console(
+    cell: CellImpl,
+    ctx: PostExecutionHookContext,
+    run_result: cell_runner.RunResult,
+) -> None:
+    """Flush buffered console output before marking the cell idle.
+
+    Console messages (stdout/stderr) are batched by a background thread
+    for performance.  Without an explicit flush, the messages may arrive
+    at the frontend *after* the cell is marked idle and ``completed-run``
+    is sent.  A subsequent run would then clear the console (via
+    ``console=[]``) before the user sees the output.
+    """
+    del cell
+    del run_result
+    del ctx
+    stream = get_context().stream
+    stream.flush_console()
+
+
 POST_EXECUTION_HOOKS: list[PostExecutionHook] = [
     _set_imported_defs,
     _set_run_result_status,
@@ -506,6 +527,9 @@ def _reset_matplotlib_context(
     _broadcast_duckdb_datasource,
     _broadcast_outputs,
     _reset_matplotlib_context,
+    # Flush buffered console output so that stderr/stdout arrives at the
+    # frontend before the cell transitions to idle.
+    _flush_console,
     # set status to idle after all post-processing is done, in case the
     # other hooks take a long time (broadcast outputs can take a long time
     # if a formatter is slow).

tests/_messaging/test_console_output_worker.py

@@ -9,6 +9,7 @@
 from marimo._messaging.console_output_worker import (
     TIMEOUT_S,
     ConsoleMsg,
+    FlushMarker,
     _add_output_to_buffer,
     _can_merge_outputs,
     _write_console_output,
@@ -181,6 +182,74 @@ def test_buffered_writer_basic(self) -> None:
                 cv.notify()
             thread.join(timeout=1.0)
 
+    def test_buffered_writer_flush_marker_flushes_immediately(self) -> None:
+        # A FlushMarker should cause buffered output to be written out
+        # without waiting for the TIMEOUT_S batching window, and the
+        # marker's done event should be set once flushing is complete.
+        stream = MockStream()
+        msg_queue: deque[ConsoleMsg | FlushMarker | None] = deque()
+        cv = threading.Condition()
+
+        thread = threading.Thread(
+            target=buffered_writer, args=(msg_queue, stream, cv)
+        )
+        thread.daemon = True
+        thread.start()
+
+        try:
+            marker = FlushMarker()
+            with cv:
+                msg_queue.append(
+                    ConsoleMsg(
+                        stream=CellChannel.STDOUT,
+                        cell_id="cell1",
+                        data="Hello",
+                        mimetype="text/plain",
+                    )
+                )
+                msg_queue.append(marker)
+                cv.notify()
+
+            # The flush should complete promptly — well before a batching
+            # window would expire.  Give it generous headroom for CI jitter
+            # but assert it happens faster than a multi-batch wait.
+            assert marker.done.wait(timeout=1.0), (
+                "FlushMarker did not signal completion within timeout"
+            )
+            assert len(stream.operations) == 1
+            assert stream.operations[0]["console"]["data"] == "Hello"
+        finally:
+            with cv:
+                msg_queue.append(None)
+                cv.notify()
+            thread.join(timeout=1.0)
+
+    def test_buffered_writer_flush_marker_on_empty_queue(self) -> None:
+        # A FlushMarker on an empty buffer should still signal completion.
+        stream = MockStream()
+        msg_queue: deque[ConsoleMsg | FlushMarker | None] = deque()
+        cv = threading.Condition()
+
+        thread = threading.Thread(
+            target=buffered_writer, args=(msg_queue, stream, cv)
+        )
+        thread.daemon = True
+        thread.start()
+
+        try:
+            marker = FlushMarker()
+            with cv:
+                msg_queue.append(marker)
+                cv.notify()
+
+            assert marker.done.wait(timeout=1.0)
+            assert stream.operations == []
+        finally:
+            with cv:
+                msg_queue.append(None)
+                cv.notify()
+            thread.join(timeout=1.0)
+
     def test_buffered_writer_multiple_messages(self) -> None:
         # Test buffered writer with multiple messages
         stream = MockStream()