deep_research/src/core.rs

@@ -1,12 +1,11 @@
 use crate::progress::Spinner;
 use crate::review::{self, Review};
+use crate::stream::write_text_stream;
 use crate::tools::{FetchPage, SearchWeb};
 use clap::Parser;
-use futures::StreamExt;
-use rig::agent::MultiTurnStreamItem;
 use rig::client::{AgentClientExt, Nothing};
 use rig::providers::ollama;
-use rig::streaming::{StreamedAssistantContent, StreamingPrompt};
+use rig::streaming::StreamingPrompt;
 use std::io::Write;
 
 /// 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
     // to the same terminal line, so the two must not race over stdout.
     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}"))
         .await;
     drop(spinner);
 
-    let mut report = String::new();
     // 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.
     let stdout = std::io::stdout();
     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)?;
 
     Ok(report)

deep_research/src/main.rs

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

deep_research/src/stream.rs

@@ -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");
+    }
+}