diff --git a/deep_research/src/core.rs b/deep_research/src/core.rs index a96da55..0a55425 100644 --- a/deep_research/src/core.rs +++ b/deep_research/src/core.rs @@ -171,20 +171,19 @@ async fn write_report( ) .build(); - // Drop the spinner before streaming starts: report text is about to print - // to the same terminal line, so the two must not race over stdout. let spinner = Spinner::start(show_progress, format!("{REPORT_EMOJI} Writing report...")); let response_stream = writer .stream_prompt(format!("Topic: {topic}\n\nResearch notes:\n{findings}")) .await; - drop(spinner); // Locked once for the whole stream rather than per chunk (as print! // would do internally) — chunks arrive in a tight loop, so re-acquiring - // the lock on every one adds up. + // the lock on every one adds up. The spinner keeps running until the + // stream's first chunk arrives, so the terminal stays covered through + // the gap between sending the prompt and generation actually starting. let stdout = std::io::stdout(); let mut handle = stdout.lock(); - let report = write_text_stream(response_stream, &mut handle).await?; + let report = write_text_stream(response_stream, &mut handle, spinner).await?; writeln!(handle)?; Ok(report) diff --git a/deep_research/src/progress.rs b/deep_research/src/progress.rs index 7534f71..ed0312f 100644 --- a/deep_research/src/progress.rs +++ b/deep_research/src/progress.rs @@ -47,14 +47,23 @@ impl Spinner { Self(Some(bar)) } + + /// Clears the spinner immediately rather than waiting for drop — for + /// callers that need it gone at a precise moment (e.g. right as the + /// first chunk of a stream is about to print on the same line) rather + /// than whenever the value happens to go out of scope. Idempotent: a + /// spinner already stopped, or one that was never enabled, does nothing. + pub(crate) fn stop(&mut self) { + if let Some(bar) = self.0.take() { + bar.finish_and_clear(); + *active().lock().expect("spinner mutex poisoned") = None; + } + } } impl Drop for Spinner { fn drop(&mut self) { - if let Some(bar) = &self.0 { - bar.finish_and_clear(); - *active().lock().expect("spinner mutex poisoned") = None; - } + self.stop(); } } diff --git a/deep_research/src/stream.rs b/deep_research/src/stream.rs index 36d71cb..bb3622b 100644 --- a/deep_research/src/stream.rs +++ b/deep_research/src/stream.rs @@ -1,3 +1,4 @@ +use crate::progress::Spinner; use futures::{Stream, StreamExt}; use rig::agent::{MultiTurnStreamItem, StreamingError}; use rig::streaming::StreamedAssistantContent; @@ -12,9 +13,15 @@ use std::io::Write; /// caller controls the lock's lifetime: locking once around a whole report /// (as `write_report` does) avoids re-acquiring it on every chunk, the way /// `print!` would. +/// +/// `spinner` stays up until the stream actually produces its first item, +/// covering the gap between the prompt being sent and generation starting +/// (otherwise the terminal would go blank for however long that takes) +/// rather than being dropped by the caller before this is even called. pub(crate) async fn write_text_stream( mut stream: impl Stream, StreamingError>> + Unpin, writer: &mut impl Write, + mut spinner: Spinner, ) -> anyhow::Result where R: Clone, @@ -22,6 +29,8 @@ where let mut text = String::new(); while let Some(chunk) = stream.next().await { + spinner.stop(); + if let MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Text(chunk)) = chunk? { @@ -55,7 +64,7 @@ mod tests { let items = vec![text_item("Hello, "), text_item("world!")]; let mut written = Vec::new(); - let accumulated = write_text_stream(stream::iter(items), &mut written) + let accumulated = write_text_stream(stream::iter(items), &mut written, Spinner::start(false, "")) .await .unwrap(); @@ -72,7 +81,7 @@ mod tests { let items = vec![text_item("kept"), final_item]; let mut written = Vec::new(); - let accumulated = write_text_stream(stream::iter(items), &mut written) + let accumulated = write_text_stream(stream::iter(items), &mut written, Spinner::start(false, "")) .await .unwrap(); @@ -88,7 +97,7 @@ mod tests { let items = vec![text_item("kept"), Err(error)]; let mut written = Vec::new(); - let result = write_text_stream(stream::iter(items), &mut written).await; + let result = write_text_stream(stream::iter(items), &mut written, Spinner::start(false, "")).await; assert!(result.is_err()); assert_eq!(String::from_utf8(written).unwrap(), "kept");