doubleo7/deep_research/src/stream.rs

use crate::progress::Spinner;
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.
///
/// `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,
{
    let mut text = String::new();

    while let Some(chunk) = stream.next().await {
        spinner.stop();

        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, Spinner::start(false, ""))
            .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, Spinner::start(false, ""))
            .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, Spinner::start(false, "")).await;

        assert!(result.is_err());
        assert_eq!(String::from_utf8(written).unwrap(), "kept");
    }
}