Compare commits
No commits in common. "351cc79e573d8d4b00ee70a9826740e1d75b66dc" and "34c699df0fc39b91cbba94312d8f9ced41064273" have entirely different histories.
351cc79e57
...
34c699df0f
3 changed files with 12 additions and 29 deletions
|
|
@ -171,19 +171,20 @@ 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 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.
|
||||
// the lock on every one adds up.
|
||||
let stdout = std::io::stdout();
|
||||
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)?;
|
||||
|
||||
Ok(report)
|
||||
|
|
|
|||
|
|
@ -47,23 +47,14 @@ 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) {
|
||||
self.stop();
|
||||
if let Some(bar) = &self.0 {
|
||||
bar.finish_and_clear();
|
||||
*active().lock().expect("spinner mutex poisoned") = None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
use crate::progress::Spinner;
|
||||
use futures::{Stream, StreamExt};
|
||||
use rig::agent::{MultiTurnStreamItem, StreamingError};
|
||||
use rig::streaming::StreamedAssistantContent;
|
||||
|
|
@ -13,15 +12,9 @@ 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<R>(
|
||||
mut stream: impl Stream<Item = Result<MultiTurnStreamItem<R>, StreamingError>> + Unpin,
|
||||
writer: &mut impl Write,
|
||||
mut spinner: Spinner,
|
||||
) -> anyhow::Result<String>
|
||||
where
|
||||
R: Clone,
|
||||
|
|
@ -29,8 +22,6 @@ where
|
|||
let mut text = String::new();
|
||||
|
||||
while let Some(chunk) = stream.next().await {
|
||||
spinner.stop();
|
||||
|
||||
if let MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Text(chunk)) =
|
||||
chunk?
|
||||
{
|
||||
|
|
@ -64,7 +55,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, Spinner::start(false, ""))
|
||||
let accumulated = write_text_stream(stream::iter(items), &mut written)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
|
|
@ -81,7 +72,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, Spinner::start(false, ""))
|
||||
let accumulated = write_text_stream(stream::iter(items), &mut written)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
|
|
@ -97,7 +88,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, Spinner::start(false, "")).await;
|
||||
let result = write_text_stream(stream::iter(items), &mut written).await;
|
||||
|
||||
assert!(result.is_err());
|
||||
assert_eq!(String::from_utf8(written).unwrap(), "kept");
|
||||
|
|
|
|||
Loading…
Reference in a new issue