Fix 2265: close tui on success, custom tui subscriber (#2734)

packages/cli/src/assets.rs

@@ -1,4 +1,6 @@
-use crate::builder::{BuildMessage, MessageType, Stage, UpdateBuildProgress, UpdateStage};
+use crate::builder::{
+    BuildMessage, MessageSource, MessageType, Stage, UpdateBuildProgress, UpdateStage,
+};
 use crate::dioxus_crate::DioxusCrate;
 use crate::Result;
 use anyhow::Context;
@@ -67,7 +69,7 @@ pub(crate) fn process_assets(
                                 "Optimized static asset {}",
                                 file_asset
                             )),
-                            source: None,
+                            source: MessageSource::Build,
                         }),
                     });
                     assets_finished += 1;

packages/cli/src/builder/mod.rs

@@ -14,7 +14,9 @@ mod fullstack;
 mod prepare_html;
 mod progress;
 mod web;
-pub use progress::{BuildMessage, MessageType, Stage, UpdateBuildProgress, UpdateStage};
+pub use progress::{
+    BuildMessage, MessageSource, MessageType, Stage, UpdateBuildProgress, UpdateStage,
+};
 
 /// A request for a project to be built
 pub struct BuildRequest {

packages/cli/src/builder/prepare_html.rs

@@ -1,6 +1,7 @@
 //! Build the HTML file to load a web application. The index.html file may be created from scratch or modified from the `index.html` file in the crate root.
 
 use super::{BuildRequest, UpdateBuildProgress};
+use crate::builder::progress::MessageSource;
 use crate::builder::Stage;
 use crate::Result;
 use futures_channel::mpsc::UnboundedSender;
@@ -191,7 +192,7 @@ impl BuildRequest {
             update: super::UpdateStage::AddMessage(super::BuildMessage {
                 level: Level::WARN,
                 message: super::MessageType::Text(message),
-                source: None,
+                source: MessageSource::Build,
             }),
         });
     }

packages/cli/src/builder/progress.rs

@@ -3,6 +3,7 @@ use anyhow::Context;
 use cargo_metadata::{diagnostic::Diagnostic, Message};
 use futures_channel::mpsc::UnboundedSender;
 use serde::Deserialize;
+use std::fmt::Display;
 use std::ops::Deref;
 use std::path::PathBuf;
 use std::process::Stdio;
@@ -83,7 +84,7 @@ pub enum UpdateStage {
 pub struct BuildMessage {
     pub level: Level,
     pub message: MessageType,
-    pub source: Option<String>,
+    pub source: MessageSource,
 }
 
 #[derive(Debug, Clone, PartialEq)]
@@ -92,6 +93,34 @@ pub enum MessageType {
     Text(String),
 }
 
+/// Represents the source of where a message came from.
+///
+/// The CLI will render a prefix according to the message type
+/// but this prefix, [`MessageSource::to_string()`] shouldn't be used if a strict message source is required.
+#[derive(Debug, Clone, PartialEq)]
+pub enum MessageSource {
+    /// Represents any message from the running application. Renders `[app]`
+    App,
+    /// Represents any generic message from the CLI. Renders `[dev]`
+    ///
+    /// Usage of Tracing inside of the CLI will be routed to this type.
+    Dev,
+    /// Represents a message from the build process. Renders `[bld]`
+    ///
+    /// This is anything emitted from a build process such as cargo and optimizations.
+    Build,
+}
+
+impl Display for MessageSource {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::App => write!(f, "app"),
+            Self::Dev => write!(f, "dev"),
+            Self::Build => write!(f, "bld"),
+        }
+    }
+}
+
 impl From<Diagnostic> for BuildMessage {
     fn from(message: Diagnostic) -> Self {
         Self {
@@ -104,7 +133,7 @@ impl From<Diagnostic> for BuildMessage {
                 cargo_metadata::diagnostic::DiagnosticLevel::Help => Level::DEBUG,
                 _ => Level::DEBUG,
             },
-            source: Some("cargo".to_string()),
+            source: MessageSource::Build,
             message: MessageType::Cargo(message),
         }
     }
@@ -206,7 +235,7 @@ pub(crate) async fn build_cargo(
                     update: UpdateStage::AddMessage(BuildMessage {
                         level: Level::DEBUG,
                         message: MessageType::Text(line),
-                        source: None,
+                        source: MessageSource::Build,
                     }),
                 });
             }

packages/cli/src/cli/serve.rs

@@ -1,5 +1,6 @@
 use crate::{
     settings::{self},
+    tracer::CLILogControl,
     DioxusCrate,
 };
 use anyhow::Context;
@@ -105,13 +106,13 @@ impl Serve {
         Ok(())
     }
 
-    pub async fn serve(mut self) -> anyhow::Result<()> {
+    pub async fn serve(mut self, log_control: CLILogControl) -> anyhow::Result<()> {
         let mut dioxus_crate = DioxusCrate::new(&self.build_arguments.target_args)
             .context("Failed to load Dioxus workspace")?;
 
         self.resolve(&mut dioxus_crate)?;
 
-        crate::serve::serve_all(self, dioxus_crate).await?;
+        crate::serve::serve_all(self, dioxus_crate, log_control).await?;
         Ok(())
     }
 }

packages/cli/src/main.rs

@@ -6,6 +6,7 @@ pub mod assets;
 pub mod dx_build_info;
 pub mod serve;
 pub mod tools;
+pub mod tracer;
 
 pub mod cli;
 pub use cli::*;
@@ -23,21 +24,16 @@ pub(crate) use settings::*;
 
 pub(crate) mod metadata;
 
-use std::env;
-use tracing_subscriber::{prelude::*, EnvFilter, Layer};
-
 use anyhow::Context;
 use clap::Parser;
 
 use Commands::*;
 
-const LOG_ENV: &str = "DIOXUS_LOG";
-
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
     let args = Cli::parse();
 
-    build_tracing();
+    let log_control = tracer::build_tracing();
 
     match args.action {
         Translate(opts) => opts
@@ -79,7 +75,7 @@ async fn main() -> anyhow::Result<()> {
             .context(error_wrapper("Cleaning project failed")),
 
         Serve(opts) => opts
-            .serve()
+            .serve(log_control)
             .await
             .context(error_wrapper("Serving project failed")),
 
@@ -94,21 +90,3 @@ async fn main() -> anyhow::Result<()> {
 fn error_wrapper(message: &str) -> String {
     format!("🚫 {message}:")
 }
-
-fn build_tracing() {
-    // If {LOG_ENV} is set, default to env, otherwise filter to cli
-    // and manganis warnings and errors from other crates
-    let mut filter = EnvFilter::new("error,dx=info,dioxus-cli=info,manganis-cli-support=info");
-    if env::var(LOG_ENV).is_ok() {
-        filter = EnvFilter::from_env(LOG_ENV);
-    }
-
-    let sub =
-        tracing_subscriber::registry().with(tracing_subscriber::fmt::layer().with_filter(filter));
-
-    #[cfg(feature = "tokio-console")]
-    sub.with(console_subscriber::spawn()).init();
-
-    #[cfg(not(feature = "tokio-console"))]
-    sub.init();
-}

packages/cli/src/serve/builder.rs

@@ -6,9 +6,9 @@ use crate::serve::Serve;
 use crate::Result;
 use dioxus_cli_config::Platform;
 use futures_channel::mpsc::UnboundedReceiver;
-use futures_util::future::OptionFuture;
 use futures_util::stream::select_all;
 use futures_util::StreamExt;
+use futures_util::{future::OptionFuture, stream::FuturesUnordered};
 use std::process::Stdio;
 use tokio::{
     process::{Child, Command},
@@ -95,8 +95,20 @@ impl Builder {
                 .map(|(platform, rx)| rx.map(move |update| (*platform, update))),
         );
 
+        // The ongoing builds directly
         let results: OptionFuture<_> = self.build_results.as_mut().into();
 
+        // The process exits
+        let mut process_exited = self
+            .children
+            .iter_mut()
+            .map(|(_, child)| async move {
+                let status = child.wait().await.ok();
+
+                BuilderUpdate::ProcessExited { status }
+            })
+            .collect::<FuturesUnordered<_>>();
+
         // Wait for the next build result
         tokio::select! {
             Some(build_results) = results => {
@@ -111,6 +123,9 @@ impl Builder {
                 // If we have a build progress, send it to the screen
                 Ok(BuilderUpdate::Progress { platform, update })
             }
+            Some(exit_status) = process_exited.next() => {
+                Ok(exit_status)
+            }
             else => {
                 std::future::pending::<()>().await;
                 unreachable!("Pending cannot resolve")
@@ -164,4 +179,7 @@ pub enum BuilderUpdate {
     Ready {
         results: Vec<BuildResult>,
     },
+    ProcessExited {
+        status: Option<std::process::ExitStatus>,
+    },
 }

packages/cli/src/serve/mod.rs

@@ -1,6 +1,7 @@
 use crate::builder::{Stage, UpdateBuildProgress, UpdateStage};
 use crate::cli::serve::Serve;
 use crate::dioxus_crate::DioxusCrate;
+use crate::tracer::CLILogControl;
 use crate::Result;
 use dioxus_cli_config::Platform;
 use tokio::task::yield_now;
@@ -45,7 +46,11 @@ use watcher::*;
 /// - Consume logs from the wasm for web/fullstack
 /// - I want us to be able to detect a `server_fn` in the project and then upgrade from a static server
 ///   to a dynamic one on the fly.
-pub async fn serve_all(serve: Serve, dioxus_crate: DioxusCrate) -> Result<()> {
+pub async fn serve_all(
+    serve: Serve,
+    dioxus_crate: DioxusCrate,
+    log_control: CLILogControl,
+) -> Result<()> {
     let mut builder = Builder::new(&dioxus_crate, &serve);
 
     // Start the first build
@@ -53,7 +58,7 @@ pub async fn serve_all(serve: Serve, dioxus_crate: DioxusCrate) -> Result<()> {
 
     let mut server = Server::start(&serve, &dioxus_crate);
     let mut watcher = Watcher::start(&serve, &dioxus_crate);
-    let mut screen = Output::start(&serve).expect("Failed to open terminal logger");
+    let mut screen = Output::start(&serve, log_control).expect("Failed to open terminal logger");
 
     let is_hot_reload = serve.server_arguments.hot_reload.unwrap_or(true);
 
@@ -143,6 +148,15 @@ pub async fn serve_all(serve: Serve, dioxus_crate: DioxusCrate) -> Result<()> {
                         // And then finally tell the server to reload
                         server.send_reload_command().await;
                     },
+
+                    // If the process exited *cleanly*, we can exit
+                    Ok(BuilderUpdate::ProcessExited { status, ..}) => {
+                        if let Some(status) = status {
+                            if status.success() {
+                                break;
+                            }
+                        }
+                    }
                     Err(err) => {
                         server.send_build_error(err).await;
                     }

packages/cli/src/serve/output.rs

@@ -1,6 +1,7 @@
 use crate::{
-    builder::{BuildMessage, MessageType, Stage, UpdateBuildProgress},
+    builder::{BuildMessage, MessageSource, MessageType, Stage, UpdateBuildProgress},
     dioxus_crate::DioxusCrate,
+    tracer::CLILogControl,
 };
 use crate::{
     builder::{BuildResult, UpdateStage},
@@ -23,6 +24,7 @@ use std::{
     io::{self, stdout},
     pin::Pin,
     rc::Rc,
+    sync::atomic::Ordering,
     time::{Duration, Instant},
 };
 use tokio::{
@@ -55,6 +57,7 @@ impl BuildProgress {
 
 pub struct Output {
     term: Rc<RefCell<Option<TerminalBackend>>>,
+    log_control: CLILogControl,
 
     // optional since when there's no tty there's no eventstream to read from - just stdin
     events: Option<EventStream>,
@@ -88,12 +91,13 @@ enum Tab {
 type TerminalBackend = Terminal<CrosstermBackend<io::Stdout>>;
 
 impl Output {
-    pub fn start(cfg: &Serve) -> io::Result<Self> {
+    pub fn start(cfg: &Serve, log_control: CLILogControl) -> io::Result<Self> {
         let interactive = std::io::stdout().is_tty() && cfg.interactive.unwrap_or(true);
 
         let mut events = None;
 
         if interactive {
+            log_control.tui_enabled.store(true, Ordering::SeqCst);
             enable_raw_mode()?;
             stdout().execute(EnterAlternateScreen)?;
 
@@ -138,6 +142,7 @@ impl Output {
 
         Ok(Self {
             term: Rc::new(RefCell::new(term)),
+            log_control,
             events,
             _rustc_version,
             _rustc_nightly,
@@ -176,8 +181,8 @@ impl Output {
         let has_running_apps = !self.running_apps.is_empty();
         let next_stdout = self.running_apps.values_mut().map(|app| {
             let future = async move {
-                let (stdout, stderr) = match &mut app.stdout {
-                    Some(stdout) => (stdout.stdout.next_line(), stdout.stderr.next_line()),
+                let (stdout, stderr) = match &mut app.output {
+                    Some(out) => (out.stdout.next_line(), out.stderr.next_line()),
                     None => return futures_util::future::pending().await,
                 };
 
@@ -199,29 +204,39 @@ impl Output {
         };
 
         let animation_timeout = tokio::time::sleep(Duration::from_millis(300));
+        let tui_log_rx = &mut self.log_control.tui_rx;
 
         tokio::select! {
             (platform, stdout, stderr) = next_stdout => {
                 if let Some(stdout) = stdout {
-                    self.running_apps.get_mut(&platform).unwrap().stdout.as_mut().unwrap().stdout_line.push_str(&stdout);
+                    self.running_apps.get_mut(&platform).unwrap().output.as_mut().unwrap().stdout_line.push_str(&stdout);
                     self.push_log(platform, BuildMessage {
                         level: Level::INFO,
                         message: MessageType::Text(stdout),
-                        source: Some("app".to_string()),
+                        source: MessageSource::App,
                     })
                 }
                 if let Some(stderr) = stderr {
                     self.set_tab(Tab::BuildLog);
 
-                    self.running_apps.get_mut(&platform).unwrap().stdout.as_mut().unwrap().stderr_line.push_str(&stderr);
+                    self.running_apps.get_mut(&platform).unwrap().output.as_mut().unwrap().stderr_line.push_str(&stderr);
                     self.build_progress.build_logs.get_mut(&platform).unwrap().messages.push(BuildMessage {
                         level: Level::ERROR,
                         message: MessageType::Text(stderr),
-                        source: Some("app".to_string()),
+                        source: MessageSource::App,
                     });
                 }
             },
 
+            // Handle internal CLI tracing logs.
+            Some(log) = tui_log_rx.next() => {
+                self.push_log(self.platform, BuildMessage {
+                    level: Level::INFO,
+                    message: MessageType::Text(log),
+                    source: MessageSource::Dev,
+                });
+            }
+
             event = user_input => {
                 if self.handle_events(event.unwrap().unwrap()).await? {
                     return Ok(true)
@@ -238,6 +253,7 @@ impl Output {
     pub fn shutdown(&mut self) -> io::Result<()> {
         // if we're a tty then we need to disable the raw mode
         if self.interactive {
+            self.log_control.tui_enabled.store(false, Ordering::SeqCst);
             disable_raw_mode()?;
             stdout().execute(LeaveAlternateScreen)?;
             self.drain_print_logs();
@@ -366,7 +382,7 @@ impl Output {
                                 // we need to translate its styling into our own
                                 messages.first().unwrap_or(&String::new()).clone(),
                             ),
-                            source: Some("app".to_string()),
+                            source: MessageSource::App,
                         },
                     );
                 }
@@ -375,8 +391,8 @@ impl Output {
                         platform,
                         BuildMessage {
                             level: Level::ERROR,
-                            source: Some("app".to_string()),
-                            message: MessageType::Text(format!("Error parsing message: {err}")),
+                            source: MessageSource::Dev,
+                            message: MessageType::Text(format!("Error parsing app message: {err}")),
                         },
                     );
                 }
@@ -401,9 +417,12 @@ impl Output {
     pub fn push_log(&mut self, platform: Platform, message: BuildMessage) {
         let snapped = self.is_snapped(platform);
 
-        if let Some(build) = self.build_progress.build_logs.get_mut(&platform) {
-            build.stdout_logs.push(message);
-        }
+        self.build_progress
+            .build_logs
+            .entry(platform)
+            .or_default()
+            .stdout_logs
+            .push(message);
 
         if snapped {
             self.scroll_to_bottom();
@@ -453,7 +472,10 @@ impl Output {
                 stderr_line: String::new(),
             });
 
-            let app = RunningApp { result, stdout };
+            let app = RunningApp {
+                result,
+                output: stdout,
+            };
 
             self.running_apps.insert(platform, app);
 
@@ -639,7 +661,16 @@ impl Output {
                                 for line in line.lines() {
                                     let text = line.into_text().unwrap_or_default();
                                     for line in text.lines {
-                                        let mut out_line = vec![Span::from("[app] ").dark_gray()];
+                                        let source = format!("[{}] ", span.source);
+
+                                        let msg_span = Span::from(source);
+                                        let msg_span = match span.source {
+                                            MessageSource::App => msg_span.light_blue(),
+                                            MessageSource::Dev => msg_span.dark_gray(),
+                                            MessageSource::Build => msg_span.light_yellow(),
+                                        };
+
+                                        let mut out_line = vec![msg_span];
                                         for span in line.spans {
                                             out_line.push(span);
                                         }
@@ -855,7 +886,7 @@ async fn rustc_version() -> String {
 
 pub struct RunningApp {
     result: BuildResult,
-    stdout: Option<RunningAppOutput>,
+    output: Option<RunningAppOutput>,
 }
 
 struct RunningAppOutput {

packages/cli/src/tracer.rs

@@ -0,0 +1,96 @@
+use futures_channel::mpsc::{unbounded, UnboundedReceiver, UnboundedSender};
+use std::{
+    env, io,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc, Mutex,
+    },
+};
+use tracing_subscriber::{prelude::*, EnvFilter};
+
+const LOG_ENV: &str = "DIOXUS_LOG";
+
+/// Build tracing infrastructure.
+pub fn build_tracing() -> CLILogControl {
+    // If {LOG_ENV} is set, default to env, otherwise filter to cli
+    // and manganis warnings and errors from other crates
+    let mut filter = EnvFilter::new("error,dx=info,dioxus-cli=info,manganis-cli-support=info");
+    if env::var(LOG_ENV).is_ok() {
+        filter = EnvFilter::from_env(LOG_ENV);
+    }
+
+    // Create writer controller and custom writer.
+    let (tui_tx, tui_rx) = unbounded();
+    let tui_enabled = Arc::new(AtomicBool::new(false));
+
+    let writer_control = CLILogControl {
+        tui_rx,
+        tui_enabled: tui_enabled.clone(),
+    };
+    let cli_writer = Mutex::new(CLIWriter::new(tui_enabled, tui_tx));
+
+    // Build tracing
+    let fmt_layer = tracing_subscriber::fmt::layer()
+        .with_writer(cli_writer)
+        .with_filter(filter);
+    let sub = tracing_subscriber::registry().with(fmt_layer);
+
+    #[cfg(feature = "tokio-console")]
+    let sub = sub.with(console_subscriber::spawn());
+
+    sub.init();
+
+    writer_control
+}
+
+/// Contains the sync primitives to control the CLIWriter.
+pub struct CLILogControl {
+    pub tui_rx: UnboundedReceiver<String>,
+    pub tui_enabled: Arc<AtomicBool>,
+}
+
+/// Represents the CLI's custom tracing writer for conditionally writing logs between outputs.
+pub struct CLIWriter {
+    stdout: io::Stdout,
+    tui_tx: UnboundedSender<String>,
+    tui_enabled: Arc<AtomicBool>,
+}
+
+impl CLIWriter {
+    /// Create a new CLIWriter with required sync primitives for conditionally routing logs.
+    pub fn new(tui_enabled: Arc<AtomicBool>, tui_tx: UnboundedSender<String>) -> Self {
+        Self {
+            stdout: io::stdout(),
+            tui_tx,
+            tui_enabled,
+        }
+    }
+}
+
+// Implement a conditional writer so that logs are routed to the appropriate place.
+impl io::Write for CLIWriter {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        if self.tui_enabled.load(Ordering::SeqCst) {
+            let len = buf.len();
+
+            let as_string = String::from_utf8(buf.to_vec())
+                .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
+
+            self.tui_tx
+                .unbounded_send(as_string)
+                .map_err(|e| io::Error::new(io::ErrorKind::BrokenPipe, e))?;
+
+            Ok(len)
+        } else {
+            self.stdout.write(buf)
+        }
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        if !self.tui_enabled.load(Ordering::SeqCst) {
+            self.stdout.flush()
+        } else {
+            Ok(())
+        }
+    }
+}