Trying to get experimental UI working with watch mode

crates/turborepo-lib/src/run/builder.rs

@@ -26,7 +26,7 @@ use turborepo_telemetry::events::{
     repo::{RepoEventBuilder, RepoType},
     EventBuilder, TrackedErrors,
 };
-use turborepo_ui::{ColorSelector, UI};
+use turborepo_ui::{tui, tui::AppSender, ColorSelector, UI};
 #[cfg(feature = "daemon-package-discovery")]
 use {
     crate::run::package_discovery::DaemonPackageDiscovery,
@@ -59,6 +59,8 @@ pub struct RunBuilder {
     ui: UI,
     version: &'static str,
     experimental_ui: bool,
+    // If we have an existing sender, we can pass it here
+    experimental_ui_sender: Option<AppSender>,
     api_client: APIClient,
     // In watch mode, we can have a changed package that we want to serve as an entrypoint.
     // We will then prune away any tasks that do not depend on tasks inside
@@ -114,6 +116,7 @@ impl RunBuilder {
             ui,
             version,
             experimental_ui,
+            experimental_ui_sender: None,
             entrypoint_package: None,
             should_print_prelude_override: None,
         })
@@ -124,6 +127,15 @@ impl RunBuilder {
         self
     }
 
+    // If we've already started a UI thread, use that
+    pub fn with_existing_experimental_ui(
+        mut self,
+        experimental_ui_sender: Option<AppSender>,
+    ) -> Self {
+        self.experimental_ui_sender = experimental_ui_sender;
+        self
+    }
+
     pub fn hide_prelude(mut self) -> Self {
         self.should_print_prelude_override = Some(false);
         self
@@ -405,10 +417,27 @@ impl RunBuilder {
             self.opts.run_opts.dry_run.is_none() && self.opts.run_opts.graph.is_none()
         });
 
+        let (experimental_ui_sender, experimental_ui_handle) = match &self.experimental_ui_sender {
+            Some(sender) if self.experimental_ui => {
+                println!("1");
+                // We already have the UI thread, so use the sender, but don't keep a handle
+                (Some(sender.clone()), None)
+            }
+            None if self.experimental_ui => {
+                println!("2");
+                let task_names = engine.tasks_with_command(&pkg_dep_graph);
+                let (handle, receiver) = AppSender::new();
+                let app = tokio::task::spawn_blocking(move || tui::run_app(task_names, receiver));
+                (Some(handle), Some(app))
+            }
+            _ => (None, None),
+        };
+
         Ok(Run {
             version: self.version,
             ui: self.ui,
-            experimental_ui: self.experimental_ui,
+            experimental_ui_sender,
+            experimental_ui_handle,
             analytics_handle,
             start_at,
             processes: self.processes,

crates/turborepo-lib/src/run/mod.rs

@@ -16,7 +16,8 @@ use std::{collections::HashSet, io::Write, sync::Arc};
 
 pub use cache::{ConfigCache, RunCache, TaskCache};
 use chrono::{DateTime, Local};
-use tracing::debug;
+use tokio::task::JoinHandle;
+use tracing::{debug, error};
 use turbopath::AbsoluteSystemPathBuf;
 use turborepo_analytics::AnalyticsHandle;
 use turborepo_api_client::{APIAuth, APIClient};
@@ -25,7 +26,7 @@ use turborepo_env::EnvironmentVariableMap;
 use turborepo_repository::package_graph::{PackageGraph, PackageName};
 use turborepo_scm::SCM;
 use turborepo_telemetry::events::generic::GenericEventBuilder;
-use turborepo_ui::{cprint, cprintln, BOLD_GREY, GREY, UI};
+use turborepo_ui::{cprint, cprintln, tui, tui::AppSender, BOLD_GREY, GREY, UI};
 
 pub use crate::run::error::Error;
 use crate::{
@@ -43,7 +44,6 @@ use crate::{
 pub struct Run {
     version: &'static str,
     ui: UI,
-    experimental_ui: bool,
     start_at: DateTime<Local>,
     processes: ProcessManager,
     run_telemetry: GenericEventBuilder,
@@ -63,6 +63,8 @@ pub struct Run {
     task_access: TaskAccess,
     analytics_handle: Option<AnalyticsHandle>,
     should_print_prelude: bool,
+    experimental_ui_sender: Option<AppSender>,
+    experimental_ui_handle: Option<JoinHandle<Result<(), tui::Error>>>,
 }
 
 impl Run {
@@ -226,7 +228,7 @@ impl Run {
             self.processes.clone(),
             &self.repo_root,
             global_env,
-            self.experimental_ui,
+            self.experimental_ui_sender.clone(),
         );
 
         if self.opts.run_opts.dry_run.is_some() {
@@ -248,6 +250,19 @@ impl Run {
             // We hit some error, it shouldn't be exit code 0
             .unwrap_or(if errors.is_empty() { 0 } else { 1 });
 
+        // If we own the handle, we should stop the UI here.
+        // If we don't own the handle, then it's likely owned
+        // by the caller, i.e. the watch client
+        if let (Some(ui), Some(render_thread_handle)) = (
+            &self.experimental_ui_sender,
+            self.experimental_ui_handle.take(),
+        ) {
+            ui.stop();
+            if let Err(e) = render_thread_handle.await.expect("render thread panicked") {
+                error!("error encountered rendering tui: {e}");
+            }
+        }
+
         let error_prefix = if self.opts.run_opts.is_github_actions {
             "::error::"
         } else {

crates/turborepo-lib/src/run/watch.rs

@@ -6,6 +6,7 @@ use thiserror::Error;
 use tokio::{select, task::JoinHandle};
 use turborepo_repository::package_graph::PackageName;
 use turborepo_telemetry::events::command::CommandEventBuilder;
+use turborepo_ui::tui::AppSender;
 
 use crate::{
     cli::{Command, ExecutionArgs, RunArgs},
@@ -83,6 +84,12 @@ impl WatchClient {
 
         run.print_run_prelude();
 
+        println!(
+            "run sender is some: {}",
+            run.experimental_ui_sender.is_some()
+        );
+        let ui = &run.experimental_ui_sender;
+
         let has_persistent_tasks = run.has_persistent_tasks();
         let mut filtered_pkgs = run.filtered_pkgs;
 
@@ -110,6 +117,7 @@ impl WatchClient {
                     &handler,
                     &mut main_run_handle,
                     has_persistent_tasks,
+                    ui.clone(),
                 )
                 .await?;
             }
@@ -140,6 +148,7 @@ impl WatchClient {
         handler: &SignalHandler,
         main_run_handle: &mut Option<JoinHandle<Result<i32, run::error::Error>>>,
         has_persistent_tasks: bool,
+        ui_sender: Option<AppSender>,
     ) -> Result<(), Error> {
         // Should we recover here?
         match event {
@@ -184,6 +193,7 @@ impl WatchClient {
                     tokio::spawn(async move {
                         let mut run = RunBuilder::new(new_base)?
                             .with_entrypoint_package(package_name)
+                            .with_existing_experimental_ui(ui_sender)
                             .hide_prelude()
                             .build(&signal_handler, telemetry)
                             .await?;
@@ -220,9 +230,9 @@ impl WatchClient {
                 let handler = handler.clone();
                 let run_fut = async move {
                     let base = CommandBase::new(args, repo_root, get_version(), ui);
-
                     let mut run = RunBuilder::new(base)?
                         .hide_prelude()
+                        .with_existing_experimental_ui(ui_sender)
                         .build(&handler, telemetry)
                         .await?;
 

crates/turborepo-lib/src/task_graph/visitor.rs

@@ -23,7 +23,7 @@ use turborepo_telemetry::events::{
     generic::GenericEventBuilder, task::PackageTaskEventBuilder, EventBuilder, TrackedErrors,
 };
 use turborepo_ui::{
-    tui::{self, TuiTask},
+    tui::{self, AppSender, TuiTask},
     ColorSelector, OutputClient, OutputSink, OutputWriter, PrefixedUI, UI,
 };
 use which::which;
@@ -62,7 +62,7 @@ pub struct Visitor<'a> {
     sink: OutputSink<StdWriter>,
     task_hasher: TaskHasher<'a>,
     ui: UI,
-    experimental_ui: bool,
+    experimental_ui_sender: Option<AppSender>,
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -105,7 +105,7 @@ impl<'a> Visitor<'a> {
         manager: ProcessManager,
         repo_root: &'a AbsoluteSystemPath,
         global_env: EnvironmentVariableMap,
-        experimental_ui: bool,
+        experimental_ui_sender: Option<AppSender>,
     ) -> Self {
         let task_hasher = TaskHasher::new(
             package_inputs_hashes,
@@ -132,7 +132,7 @@ impl<'a> Visitor<'a> {
             task_hasher,
             ui,
             global_env,
-            experimental_ui,
+            experimental_ui_sender,
         }
     }
 
@@ -145,16 +145,6 @@ impl<'a> Visitor<'a> {
         let concurrency = self.run_opts.concurrency as usize;
         let (node_sender, mut node_stream) = mpsc::channel(concurrency);
 
-        let (ui, render_thread_handle) = if self.experimental_ui {
-            let task_names = engine.tasks_with_command(&self.package_graph);
-
-            let (handle, receiver) = tui::AppSender::new();
-            let app = tokio::task::spawn_blocking(move || tui::run_app(task_names, receiver));
-            (Some(handle), Some(app))
-        } else {
-            (None, None)
-        };
-
         let engine_handle = {
             let engine = engine.clone();
             tokio::spawn(engine.execute(ExecutionOptions::new(false, concurrency), node_sender))
@@ -282,7 +272,7 @@ impl<'a> Visitor<'a> {
                     let vendor_behavior =
                         Vendor::infer().and_then(|vendor| vendor.behavior.as_ref());
 
-                    let output_client = if let Some(handle) = &ui {
+                    let output_client = if let Some(handle) = &self.experimental_ui_sender {
                         TaskOutput::UI(handle.task(info.to_string()))
                     } else {
                         TaskOutput::Direct(self.output_client(&info, vendor_behavior))
@@ -315,16 +305,6 @@ impl<'a> Visitor<'a> {
             result.unwrap_or_else(|e| panic!("task executor panicked: {e}"));
         }
         drop(factory);
-        if let Some(handle) = ui {
-            handle.stop();
-            if let Err(e) = render_thread_handle
-                .unwrap()
-                .await
-                .expect("render thread panicked")
-            {
-                error!("error encountered rendering tui: {e}");
-            }
-        }
 
         // Write out the traced-config.json file if we have one
         self.task_access.save().await;
@@ -483,7 +463,7 @@ impl<'a> Visitor<'a> {
     pub fn dry_run(&mut self) {
         self.dry = true;
         // No need to start a TUI on dry run
-        self.experimental_ui = false;
+        self.experimental_ui_sender = None;
     }
 }
 
@@ -637,7 +617,7 @@ impl<'a> ExecContextFactory<'a> {
         ExecContext {
             engine: self.engine.clone(),
             ui: self.visitor.ui,
-            experimental_ui: self.visitor.experimental_ui,
+            experimental_ui: self.visitor.experimental_ui_sender.is_some(),
             is_github_actions: self.visitor.run_opts.is_github_actions,
             pretty_prefix: self
                 .visitor