Ignore UI comm busy events by default in tests

Author: lionel-

crates/ark/tests/plots.rs

@@ -770,11 +770,11 @@ fn test_dev_hold_across_execute_requests() {
 fn test_positron_simple_plot() {
     let frontend = DummyArkFrontend::lock();
     frontend.open_ui_comm();
-    frontend.set_buffer_ui_events(true);
 
     frontend.send_execute_request("plot(1:10)", ExecuteRequestOptions::default());
     frontend.recv_iopub_busy();
     frontend.recv_iopub_execute_input();
+    frontend.recv_iopub_ui_prompt_state();
     frontend.recv_iopub_idle();
     frontend.recv_shell_execute_reply();
 
@@ -787,11 +787,11 @@ fn test_positron_simple_plot() {
 fn test_positron_multiple_plots() {
     let frontend = DummyArkFrontend::lock();
     frontend.open_ui_comm();
-    frontend.set_buffer_ui_events(true);
 
     frontend.send_execute_request("plot(1:10)\nplot(2:20)", ExecuteRequestOptions::default());
     frontend.recv_iopub_busy();
     frontend.recv_iopub_execute_input();
+    frontend.recv_iopub_ui_prompt_state();
     frontend.recv_iopub_idle();
     frontend.recv_shell_execute_reply();
 
@@ -810,16 +810,16 @@ fn test_positron_multiple_plots() {
 fn test_positron_par_multi_panel() {
     let frontend = DummyArkFrontend::lock();
     frontend.open_ui_comm();
-    frontend.set_buffer_ui_events(true);
 
     let code = "par(mfrow = c(2, 1))\nplot(1:10)\nplot(2:20)";
     frontend.send_execute_request(code, ExecuteRequestOptions::default());
     frontend.recv_iopub_busy();
     frontend.recv_iopub_execute_input();
 
-    // Panel update arrives directly on IOPub (before idle)
+    // Panel update arrives on IOPub (from on_did_execute_request)
     let update = frontend.recv_iopub_comm_msg();
 
+    frontend.recv_iopub_ui_prompt_state();
     frontend.recv_iopub_idle();
     frontend.recv_shell_execute_reply();
 
@@ -837,7 +837,6 @@ fn test_positron_par_multi_panel() {
 fn test_positron_layout_multi_plot() {
     let frontend = DummyArkFrontend::lock();
     frontend.open_ui_comm();
-    frontend.set_buffer_ui_events(true);
 
     let code = r#"
 plt2 = function() {
@@ -851,9 +850,10 @@ plt2()
     frontend.recv_iopub_busy();
     frontend.recv_iopub_execute_input();
 
-    // Second panel update (direct IOPub)
+    // Second panel update (from on_did_execute_request)
     let update = frontend.recv_iopub_comm_msg();
 
+    frontend.recv_iopub_ui_prompt_state();
     frontend.recv_iopub_idle();
     frontend.recv_shell_execute_reply();
 
@@ -870,7 +870,6 @@ plt2()
 fn test_positron_dev_hold_suppresses() {
     let frontend = DummyArkFrontend::lock();
     frontend.open_ui_comm();
-    frontend.set_buffer_ui_events(true);
 
     let code = r#"
 invisible(dev.hold())
@@ -881,6 +880,7 @@ invisible(dev.flush())
     frontend.send_execute_request(code, ExecuteRequestOptions::default());
     frontend.recv_iopub_busy();
     frontend.recv_iopub_execute_input();
+    frontend.recv_iopub_ui_prompt_state();
     frontend.recv_iopub_idle();
     frontend.recv_shell_execute_reply();
 
@@ -897,7 +897,6 @@ invisible(dev.flush())
 fn test_positron_dev_hold_across_requests() {
     let frontend = DummyArkFrontend::lock();
     frontend.open_ui_comm();
-    frontend.set_buffer_ui_events(true);
 
     // Hold and plot without flushing. No plot comm should open.
     frontend.send_execute_request(
@@ -906,13 +905,15 @@ fn test_positron_dev_hold_across_requests() {
     );
     frontend.recv_iopub_busy();
     frontend.recv_iopub_execute_input();
+    frontend.recv_iopub_ui_prompt_state();
     frontend.recv_iopub_idle();
     frontend.recv_shell_execute_reply();
 
     // Flush in a separate request. The held plot should now appear.
     frontend.send_execute_request("invisible(dev.flush())", ExecuteRequestOptions::default());
     frontend.recv_iopub_busy();
     frontend.recv_iopub_execute_input();
+    frontend.recv_iopub_ui_prompt_state();
     frontend.recv_iopub_idle();
     frontend.recv_shell_execute_reply();
 
@@ -930,13 +931,13 @@ fn test_positron_dev_hold_across_requests() {
 fn test_positron_sequential_plots() {
     let frontend = DummyArkFrontend::lock();
     frontend.open_ui_comm();
-    frontend.set_buffer_ui_events(true);
 
     for i in 1..=3 {
         let code = format!("plot({i}:10)");
         frontend.send_execute_request(&code, ExecuteRequestOptions::default());
         frontend.recv_iopub_busy();
         frontend.recv_iopub_execute_input();
+        frontend.recv_iopub_ui_prompt_state();
         frontend.recv_iopub_idle();
         frontend.recv_shell_execute_reply();
 
@@ -953,7 +954,6 @@ fn test_positron_sequential_plots() {
 fn test_positron_graphics_device_swap() {
     let frontend = DummyArkFrontend::lock();
     frontend.open_ui_comm();
-    frontend.set_buffer_ui_events(true);
 
     let code = r#"
 plot(1:10)
@@ -964,6 +964,7 @@ invisible(dev.off())
     frontend.send_execute_request(code, ExecuteRequestOptions::default());
     frontend.recv_iopub_busy();
     frontend.recv_iopub_execute_input();
+    frontend.recv_iopub_ui_prompt_state();
     frontend.recv_iopub_idle();
     frontend.recv_shell_execute_reply();
 
@@ -977,7 +978,6 @@ invisible(dev.off())
 fn test_positron_loop_plots() {
     let frontend = DummyArkFrontend::lock();
     frontend.open_ui_comm();
-    frontend.set_buffer_ui_events(true);
 
     let code = r#"
 for (i in 1:3) {
@@ -987,6 +987,7 @@ for (i in 1:3) {
     frontend.send_execute_request(code, ExecuteRequestOptions::default());
     frontend.recv_iopub_busy();
     frontend.recv_iopub_execute_input();
+    frontend.recv_iopub_ui_prompt_state();
     frontend.recv_iopub_idle();
     frontend.recv_shell_execute_reply();
 
@@ -1006,7 +1007,6 @@ for (i in 1:3) {
 fn test_positron_par_overflow_to_new_page() {
     let frontend = DummyArkFrontend::lock();
     frontend.open_ui_comm();
-    frontend.set_buffer_ui_events(true);
 
     let code = r#"
 par(mfrow = c(3, 1))
@@ -1020,10 +1020,11 @@ par(mfrow = c(1, 1))
     frontend.recv_iopub_busy();
     frontend.recv_iopub_execute_input();
 
-    // Panels 2 and 3 update the first page (direct IOPub, before idle)
+    // Panels 3 and 4 update the first page (from before.plot.new during R eval)
     let update1 = frontend.recv_iopub_comm_msg();
     let update2 = frontend.recv_iopub_comm_msg();
 
+    frontend.recv_iopub_ui_prompt_state();
     frontend.recv_iopub_idle();
     frontend.recv_shell_execute_reply();
 
@@ -1048,33 +1049,36 @@ par(mfrow = c(1, 1))
 fn test_positron_dev_hold_flush_interactive() {
     let frontend = DummyArkFrontend::lock();
     frontend.open_ui_comm();
-    frontend.set_buffer_ui_events(true);
 
     // Hold
     frontend.send_execute_request("invisible(dev.hold())", ExecuteRequestOptions::default());
     frontend.recv_iopub_busy();
     frontend.recv_iopub_execute_input();
+    frontend.recv_iopub_ui_prompt_state();
     frontend.recv_iopub_idle();
     frontend.recv_shell_execute_reply();
 
     // Draw first plot (held, no comm should open)
     frontend.send_execute_request("plot(1:10)", ExecuteRequestOptions::default());
     frontend.recv_iopub_busy();
     frontend.recv_iopub_execute_input();
+    frontend.recv_iopub_ui_prompt_state();
     frontend.recv_iopub_idle();
     frontend.recv_shell_execute_reply();
 
     // Draw over it (still held)
     frontend.send_execute_request("abline(1, 2)", ExecuteRequestOptions::default());
     frontend.recv_iopub_busy();
     frontend.recv_iopub_execute_input();
+    frontend.recv_iopub_ui_prompt_state();
     frontend.recv_iopub_idle();
     frontend.recv_shell_execute_reply();
 
     // Flush - the combined plot should now appear
     frontend.send_execute_request("invisible(dev.flush())", ExecuteRequestOptions::default());
     frontend.recv_iopub_busy();
     frontend.recv_iopub_execute_input();
+    frontend.recv_iopub_ui_prompt_state();
     frontend.recv_iopub_idle();
     frontend.recv_shell_execute_reply();
 

crates/ark/tests/ui-prompt-state.rs

@@ -14,6 +14,7 @@ use ark_test::DummyArkFrontend;
 fn test_prompt_state_after_execution() {
     let frontend = DummyArkFrontend::lock();
     let comm_id = frontend.open_ui_comm();
+    frontend.set_ignore_ui_busy(false);
 
     frontend.send_execute_request("1 + 1", ExecuteRequestOptions::default());
     frontend.recv_iopub_busy();
@@ -34,6 +35,7 @@ fn test_prompt_state_after_execution() {
 fn test_prompt_state_custom_prompt() {
     let frontend = DummyArkFrontend::lock();
     let comm_id = frontend.open_ui_comm();
+    frontend.set_ignore_ui_busy(false);
 
     // Change the prompt
     frontend.send_execute_request(
@@ -67,6 +69,7 @@ fn test_prompt_state_custom_prompt() {
 fn test_prompt_state_browser() {
     let frontend = DummyArkFrontend::lock();
     let comm_id = frontend.open_ui_comm();
+    frontend.set_ignore_ui_busy(false);
 
     // Enter the browser. The busy sequence differs from normal execution:
     // R briefly goes idle entering the browser's ReadConsole, then busy
@@ -102,6 +105,7 @@ fn test_prompt_state_browser() {
 fn test_prompt_state_custom_continuation() {
     let frontend = DummyArkFrontend::lock();
     let comm_id = frontend.open_ui_comm();
+    frontend.set_ignore_ui_busy(false);
 
     frontend.send_execute_request(
         "options(continue = '... ')",

crates/ark_test/src/dummy_frontend.rs

@@ -71,13 +71,15 @@ pub struct DummyArkFrontend {
     streams_handled: Cell<bool>,
     /// Whether we're currently in a debug context (between start_debug and stop_debug)
     in_debug: Cell<bool>,
-    /// Comm ID of the open UI comm, if any.
+    /// Comm ID of the open UI comm, if any. Set by `open_ui_comm()` so
+    /// that `recv_iopub_ui_busy()` and `recv_iopub_prompt_state()` can
+    /// identify UI comm messages.
     ui_comm_id: RefCell<Option<String>>,
-    /// When true, UI comm event messages (prompt_state, working_directory,
-    /// busy) are auto-buffered by `recv_iopub_next()` so they don't
-    /// interfere with plot and other assertions. Off by default so that
-    /// tests like `ui-prompt-state` can assert on these events explicitly.
-    buffer_ui_events: Cell<bool>,
+    /// When true (the default), UI comm `busy` events are silently
+    /// skipped by `recv_iopub_next()`. R fires busy(true)/busy(false)
+    /// per expression in multi-line requests, producing a variable
+    /// number of events that most tests don't care about.
+    ignore_ui_busy: Cell<bool>,
     /// Comm ID of the open variables comm, if any.
     variables_comm_id: RefCell<Option<String>>,
     /// Buffered variables comm events, auto-collected by `recv_iopub_next()`.
@@ -163,7 +165,7 @@ impl DummyArkFrontend {
             streams_handled: Cell::new(false),
             in_debug: Cell::new(false),
             ui_comm_id: RefCell::new(None),
-            buffer_ui_events: Cell::new(false),
+            ignore_ui_busy: Cell::new(true),
             variables_comm_id: RefCell::new(None),
             variables_events: RefCell::new(VecDeque::new()),
         }
@@ -256,14 +258,10 @@ impl DummyArkFrontend {
                 self.buffer_stream(&data.content);
                 true
             },
-            // When `buffer_ui_events` is enabled, silently consume UI comm
-            // event messages (prompt_state, working_directory, busy) so they
-            // don't interfere with plot and other assertions. RPC replies
-            // (no "method" field) always pass through.
             Message::CommMsg(ref data)
-                if self.buffer_ui_events.get() &&
+                if self.ignore_ui_busy.get() &&
                     self.is_ui_comm(&data.content.comm_id) &&
-                    data.content.data.get("method").is_some() =>
+                    data.content.data.get("method").and_then(|v| v.as_str()) == Some("busy") =>
             {
                 trace_iopub_msg(msg);
                 true
@@ -284,15 +282,47 @@ impl DummyArkFrontend {
             .is_some_and(|id| id == comm_id)
     }
 
-    /// Enable or disable automatic buffering of UI comm event messages.
+    /// Ignore UI comm `busy` events in `recv_iopub_next()`.
     ///
-    /// When enabled, prompt_state, working_directory, and busy events from
-    /// the UI comm are silently consumed by `recv_iopub_next()`. This is
-    /// useful for plot tests that don't care about these events. Tests that
-    /// need to assert on UI comm events (e.g. `ui-prompt-state`) should
-    /// leave this disabled (the default).
-    pub fn set_buffer_ui_events(&self, enabled: bool) {
-        self.buffer_ui_events.set(enabled);
+    /// R fires `busy(true)`/`busy(false)` for each top-level expression
+    /// in a multi-line request, so the number of events depends on how
+    /// many expressions the code contains. Enable this in tests that
+    /// don't care about busy transitions (e.g. plot tests).
+    pub fn set_ignore_ui_busy(&self, ignore: bool) {
+        self.ignore_ui_busy.set(ignore);
+    }
+
+    /// Receive from IOPub and assert a UI comm `busy` event.
+    /// Automatically skips any Stream messages.
+    #[track_caller]
+    pub fn recv_iopub_ui_busy(&self, expected: bool) {
+        let msg = self.recv_iopub_next();
+        match msg {
+            Message::CommMsg(data) if self.is_ui_comm(&data.content.comm_id) => {
+                assert_eq!(
+                    data.content.data.get("method").and_then(|v| v.as_str()),
+                    Some("busy")
+                );
+                assert_eq!(data.content.data["params"]["busy"], expected);
+            },
+            other => panic!("Expected UI busy={expected} CommMsg, got {other:?}"),
+        }
+    }
+
+    /// Receive from IOPub and assert a UI comm `prompt_state` event.
+    /// Automatically skips any Stream messages.
+    #[track_caller]
+    pub fn recv_iopub_ui_prompt_state(&self) {
+        let msg = self.recv_iopub_next();
+        match msg {
+            Message::CommMsg(data) if self.is_ui_comm(&data.content.comm_id) => {
+                assert_eq!(
+                    data.content.data.get("method").and_then(|v| v.as_str()),
+                    Some("prompt_state")
+                );
+            },
+            other => panic!("Expected prompt_state CommMsg on UI comm, got {other:?}"),
+        }
     }
 
     fn is_variables_comm(&self, comm_id: &str) -> bool {
@@ -1155,8 +1185,11 @@ impl DummyArkFrontend {
 
         self.recv_iopub_idle();
 
-        // Enable auto-buffering for subsequent UI comm events (prompt_state
-        // and working_directory refreshes sent after every execute request).
+        // Store the UI comm ID so `recv_iopub_ui_busy()` and
+        // `recv_iopub_prompt_state()` can identify UI comm messages.
+        // `ignore_ui_busy` is on by default, so UI busy events are
+        // auto-skipped. Tests that care (e.g. `ui-prompt-state`) can
+        // call `set_ignore_ui_busy(false)` to receive them inline.
         *self.ui_comm_id.borrow_mut() = Some(comm_id.clone());
 
         comm_id