refactor: extract the streamed-text draining loop into its own module

write_text_stream() in the new stream.rs doesn't touch anything specific
to write_report (topic, findings, the agent) — it just drains a
MultiTurnStreamItem stream, writes each text chunk to a caller-provided
writer, and returns the accumulated string. Pulling it out lets it be
covered by unit tests against a mocked stream and an in-memory writer,
independent of a live model.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Austin Schaefer 2026-08-17 12:37:11 +02:00
parent b12f156a1b
commit f53fcbd937
3 changed files with 101 additions and 22 deletions

View file

@ -1,12 +1,11 @@
use crate::progress::Spinner; use crate::progress::Spinner;
use crate::review::{self, Review}; use crate::review::{self, Review};
use crate::stream::write_text_stream;
use crate::tools::{FetchPage, SearchWeb}; use crate::tools::{FetchPage, SearchWeb};
use clap::Parser; use clap::Parser;
use futures::StreamExt;
use rig::agent::MultiTurnStreamItem;
use rig::client::{AgentClientExt, Nothing}; use rig::client::{AgentClientExt, Nothing};
use rig::providers::ollama; use rig::providers::ollama;
use rig::streaming::{StreamedAssistantContent, StreamingPrompt}; use rig::streaming::StreamingPrompt;
use std::io::Write; use std::io::Write;
/// The tool-calling research loop needs to reliably decide what to search /// The tool-calling research loop needs to reliably decide what to search
@ -171,34 +170,17 @@ async fn write_report(
// Drop the spinner before streaming starts: report text is about to print // 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. // to the same terminal line, so the two must not race over stdout.
let spinner = Spinner::start(show_progress, "Writing report..."); let spinner = Spinner::start(show_progress, "Writing report...");
let mut 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); drop(spinner);
let mut report = String::new();
// 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 lock on every one adds up.
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).await?;
while let Some(chunk) = response_stream.next().await {
match chunk? {
MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Text(text)) => {
report.push_str(&text.text);
// Terminal stdout is line-buffered, so a flush is needed here —
// otherwise a chunk without a trailing newline sits in the
// buffer instead of appearing as it streams in.
write!(handle, "{}", text.text)?;
handle.flush()?;
}
_ => continue,
}
}
// `handle` still holds stdout's lock here, so this must go through it
// rather than `println!` — reacquiring the same lock from this thread
// would deadlock.
writeln!(handle)?; writeln!(handle)?;
Ok(report) Ok(report)

View file

@ -3,6 +3,7 @@ use clap::Parser;
mod progress; mod progress;
mod review; mod review;
mod core; mod core;
mod stream;
mod tools; mod tools;
pub(crate) const DEFAULT_TOPIC: &str = pub(crate) const DEFAULT_TOPIC: &str =

View file

@ -0,0 +1,96 @@
use futures::{Stream, StreamExt};
use rig::agent::{MultiTurnStreamItem, StreamingError};
use rig::streaming::StreamedAssistantContent;
use std::io::Write;
/// Drains a multi-turn prompt stream, writing each text chunk to `writer` as
/// it arrives and returning the full accumulated text. Non-text items (tool
/// calls, reasoning, completion-call metadata, ...) are ignored — this is
/// only concerned with the assistant's prose.
///
/// `writer` is taken by reference rather than locked internally so the
/// 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.
pub(crate) async fn write_text_stream<R>(
mut stream: impl Stream<Item = Result<MultiTurnStreamItem<R>, StreamingError>> + Unpin,
writer: &mut impl Write,
) -> anyhow::Result<String>
where
R: Clone,
{
let mut text = String::new();
while let Some(chunk) = stream.next().await {
if let MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Text(chunk)) =
chunk?
{
write!(writer, "{}", chunk.text)?;
writer.flush()?;
text.push_str(&chunk.text);
}
}
Ok(text)
}
#[cfg(test)]
mod tests {
use super::*;
use futures::stream;
use rig::completion::CompletionError;
use rig::message::Text;
#[derive(Clone)]
struct DummyResponse;
fn text_item(text: &str) -> Result<MultiTurnStreamItem<DummyResponse>, StreamingError> {
Ok(MultiTurnStreamItem::StreamAssistantItem(
StreamedAssistantContent::Text(Text::new(text)),
))
}
#[tokio::test]
async fn writes_and_accumulates_text_chunks() {
let items = vec![text_item("Hello, "), text_item("world!")];
let mut written = Vec::new();
let accumulated = write_text_stream(stream::iter(items), &mut written)
.await
.unwrap();
assert_eq!(accumulated, "Hello, world!");
assert_eq!(String::from_utf8(written).unwrap(), "Hello, world!");
}
#[tokio::test]
async fn ignores_non_text_items() {
let final_item = Ok(MultiTurnStreamItem::final_response(
rig::OneOrMany::one(rig::message::AssistantContent::text("ignored")),
rig::completion::Usage::new(),
));
let items = vec![text_item("kept"), final_item];
let mut written = Vec::new();
let accumulated = write_text_stream(stream::iter(items), &mut written)
.await
.unwrap();
assert_eq!(accumulated, "kept");
assert_eq!(String::from_utf8(written).unwrap(), "kept");
}
#[tokio::test]
async fn propagates_stream_errors() {
let error = StreamingError::Completion(CompletionError::RequestError(Box::new(
std::io::Error::other("boom"),
)));
let items = vec![text_item("kept"), Err(error)];
let mut written = Vec::new();
let result = write_text_stream(stream::iter(items), &mut written).await;
assert!(result.is_err());
assert_eq!(String::from_utf8(written).unwrap(), "kept");
}
}