Fix flushing stdout/stderr after cell idle

Add a 50ms delay before finalizing cell execution to avoid dropping
trailing stdout/stderr events, since output flushes every 10ms.

The finalize step is deferred via a runtime fork and captures the
lastRunId. If a new execution starts before finalization, the runId
comparison will detect the change and skip stale finalization.

Author: manzt

extension/src/services/ExecutionRegistry.ts

@@ -1,4 +1,13 @@
-import { type Brand, Data, Effect, HashMap, Option, Ref } from "effect";
+import {
+  type Brand,
+  Data,
+  Effect,
+  FiberSet,
+  HashMap,
+  Option,
+  Ref,
+  String,
+} from "effect";
 import type * as vscode from "vscode";
 import { assert } from "../assert.ts";
 import {
@@ -17,6 +26,7 @@ export class ExecutionRegistry extends Effect.Service<ExecutionRegistry>()(
     scoped: Effect.gen(function* () {
       const code = yield* VsCode;
       const ref = yield* Ref.make(HashMap.empty<NotebookCellId, CellEntry>());
+
       yield* Effect.addFinalizer(() =>
         Ref.update(ref, (map) => {
           HashMap.forEach(map, (entry) =>
@@ -25,23 +35,20 @@ export class ExecutionRegistry extends Effect.Service<ExecutionRegistry>()(
           return HashMap.empty();
         }),
       );
+      const runFork = yield* FiberSet.makeRuntime();
       return {
-        handleInterrupted: Effect.fnUntraced(function* (
-          editor: vscode.NotebookEditor,
-        ) {
-          yield* Ref.update(ref, (map) =>
-            HashMap.map(map, (cell) => {
-              if (
-                cell.editor === editor &&
-                Option.isSome(cell.pendingExecution)
-              ) {
-                return CellEntry.interrupt(cell);
-              } else {
-                return cell;
-              }
-            }),
+        handleInterrupted(editor: vscode.NotebookEditor) {
+          return Ref.update(ref, (map) =>
+            HashMap.map(map, (cell) =>
+              Option.match(cell.pendingExecution, {
+                onSome: () => cell.editor === editor,
+                onNone: () => false,
+              })
+                ? CellEntry.interrupt(cell)
+                : cell,
+            ),
           );
-        }),
+        },
         handleCellOperation: Effect.fnUntraced(function* (
           msg: CellMessage,
           deps: {
@@ -106,13 +113,45 @@ export class ExecutionRegistry extends Effect.Service<ExecutionRegistry>()(
               }
               // MUST modify cell output before `ExecutionHandle.end`
               yield* CellEntry.maybeUpdateCellOutput(cell, code);
-              yield* Ref.update(ref, (map) =>
-                HashMap.set(
-                  map,
-                  cellId,
-                  CellEntry.end(cell, true, msg.timestamp),
-                ),
-              );
+
+              {
+                // FIXME: stdin/stdout are flushed every 10ms, so wait 50ms
+                // to ensure all related events arrive before finalizing.
+                //
+                // marimo doesn't set a run_id for idle messages, so we can't compare
+                // against the incoming message to detect if a new execution has started.
+                //
+                // Ref: https://github.com/marimo-team/marimo/blob/3644b6f/marimo/_messaging/ops.py#L148-L151
+                //
+                // Instead, we capture the `lastRunId` before the timeout and compare it
+                // when finalizing. If a new execution starts before the timeout fires,
+                // the `lastRunId` will have changed and we skip finalization.
+                const lastRunId = cell.lastRunId;
+                const finalize = Ref.update(ref, (map) => {
+                  const fresh = HashMap.get(map, cellId);
+
+                  if (Option.isNone(fresh)) {
+                    return map;
+                  }
+
+                  const isDifferentRun = !Option.getEquivalence(
+                    String.Equivalence,
+                  )(fresh.value.lastRunId, lastRunId);
+
+                  if (isDifferentRun) {
+                    return map;
+                  }
+
+                  return HashMap.set(
+                    map,
+                    cellId,
+                    CellEntry.end(fresh.value, true, msg.timestamp),
+                  );
+                });
+
+                setTimeout(() => runFork(finalize), 50);
+              }
+
               yield* Effect.logDebug("Cell execution completed").pipe(
                 Effect.annotateLogs({ cellId }),
               );
@@ -172,13 +211,15 @@ class CellEntry extends Data.TaggedClass("CellEntry")<{
   readonly state: CellRuntimeState;
   readonly editor: vscode.NotebookEditor;
   readonly pendingExecution: Option.Option<ExecutionHandle>;
+  readonly lastRunId: Option.Option<RunId>;
 }> {
   static make(id: NotebookCellId, editor: vscode.NotebookEditor) {
     return new CellEntry({
       id,
       editor,
       state: createCellRuntimeState(),
       pendingExecution: Option.none(),
+      lastRunId: Option.none(),
     });
   }
   static transition(cell: CellEntry, message: CellMessage) {
@@ -190,6 +231,7 @@ class CellEntry extends Data.TaggedClass("CellEntry")<{
   static withExecution(cell: CellEntry, execution: ExecutionHandle) {
     return new CellEntry({
       ...cell,
+      lastRunId: Option.some(execution.runId),
       pendingExecution: Option.some(execution),
     });
   }