Compare commits

..

No commits in common. "351cc79e573d8d4b00ee70a9826740e1d75b66dc" and "34c699df0fc39b91cbba94312d8f9ced41064273" have entirely different histories.

3 changed files with 12 additions and 29 deletions

View file

@ -171,19 +171,20 @@ async fn write_report(
) )
.build(); .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 spinner = Spinner::start(show_progress, format!("{REPORT_EMOJI} Writing report..."));
let response_stream = writer let response_stream = writer
.stream_prompt(format!("Topic: {topic}\n\nResearch notes:\n{findings}")) .stream_prompt(format!("Topic: {topic}\n\nResearch notes:\n{findings}"))
.await; .await;
drop(spinner);
// Locked once for the whole stream rather than per chunk (as print! // Locked once for the whole stream rather than per chunk (as print!
// would do internally) — chunks arrive in a tight loop, so re-acquiring // would do internally) — chunks arrive in a tight loop, so re-acquiring
// the lock on every one adds up. The spinner keeps running until the // the lock on every one adds up.
// 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 stdout = std::io::stdout();
let mut handle = stdout.lock(); let mut handle = stdout.lock();
let report = write_text_stream(response_stream, &mut handle, spinner).await?; let report = write_text_stream(response_stream, &mut handle).await?;
writeln!(handle)?; writeln!(handle)?;
Ok(report) Ok(report)

View file

@ -47,23 +47,14 @@ impl Spinner {
Self(Some(bar)) 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 { impl Drop for Spinner {
fn drop(&mut self) { fn drop(&mut self) {
self.stop(); if let Some(bar) = &self.0 {
bar.finish_and_clear();
*active().lock().expect("spinner mutex poisoned") = None;
}
} }
} }

View file

@ -1,4 +1,3 @@
use crate::progress::Spinner;
use futures::{Stream, StreamExt}; use futures::{Stream, StreamExt};
use rig::agent::{MultiTurnStreamItem, StreamingError}; use rig::agent::{MultiTurnStreamItem, StreamingError};
use rig::streaming::StreamedAssistantContent; use rig::streaming::StreamedAssistantContent;
@ -13,15 +12,9 @@ use std::io::Write;
/// caller controls the lock's lifetime: locking once around a whole report /// 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 /// (as `write_report` does) avoids re-acquiring it on every chunk, the way
/// `print!` would. /// `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<R>( pub(crate) async fn write_text_stream<R>(
mut stream: impl Stream<Item = Result<MultiTurnStreamItem<R>, StreamingError>> + Unpin, mut stream: impl Stream<Item = Result<MultiTurnStreamItem<R>, StreamingError>> + Unpin,
writer: &mut impl Write, writer: &mut impl Write,
mut spinner: Spinner,
) -> anyhow::Result<String> ) -> anyhow::Result<String>
where where
R: Clone, R: Clone,
@ -29,8 +22,6 @@ where
let mut text = String::new(); let mut text = String::new();
while let Some(chunk) = stream.next().await { while let Some(chunk) = stream.next().await {
spinner.stop();
if let MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Text(chunk)) = if let MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Text(chunk)) =
chunk? chunk?
{ {
@ -64,7 +55,7 @@ mod tests {
let items = vec![text_item("Hello, "), text_item("world!")]; let items = vec![text_item("Hello, "), text_item("world!")];
let mut written = Vec::new(); let mut written = Vec::new();
let accumulated = write_text_stream(stream::iter(items), &mut written, Spinner::start(false, "")) let accumulated = write_text_stream(stream::iter(items), &mut written)
.await .await
.unwrap(); .unwrap();
@ -81,7 +72,7 @@ mod tests {
let items = vec![text_item("kept"), final_item]; let items = vec![text_item("kept"), final_item];
let mut written = Vec::new(); let mut written = Vec::new();
let accumulated = write_text_stream(stream::iter(items), &mut written, Spinner::start(false, "")) let accumulated = write_text_stream(stream::iter(items), &mut written)
.await .await
.unwrap(); .unwrap();
@ -97,7 +88,7 @@ mod tests {
let items = vec![text_item("kept"), Err(error)]; let items = vec![text_item("kept"), Err(error)];
let mut written = Vec::new(); let mut written = Vec::new();
let result = write_text_stream(stream::iter(items), &mut written, Spinner::start(false, "")).await; let result = write_text_stream(stream::iter(items), &mut written).await;
assert!(result.is_err()); assert!(result.is_err());
assert_eq!(String::from_utf8(written).unwrap(), "kept"); assert_eq!(String::from_utf8(written).unwrap(), "kept");