Everything is streaming (#7247)

Co-authored-by: Douwe Osinga <douwe@squareup.com> Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
2026-02-17 19:53:38 +01:00
parent 65d3a1550f
commit 9a91fccfb5
42 changed files with 801 additions and 1718 deletions
@@ -9,7 +9,7 @@ use super::errors::ProviderError;
 use super::retry::RetryConfig;
 use crate::config::base::ConfigValue;
 use crate::config::ExtensionConfig;
-use crate::conversation::message::Message;
+use crate::conversation::message::{Message, MessageContent};
 use crate::conversation::Conversation;
 use crate::model::ModelConfig;
 use crate::utils::safe_truncate;
@@ -379,34 +379,32 @@ pub trait Provider: Send + Sync {
    /// Get the name of this provider instance
    fn get_name(&self) -> &str;

-    // Internal implementation of complete, used by complete_fast and complete
-    // Providers should override this to implement their actual completion logic
-    //
-    /// # Parameters
-    /// - `session_id`: Use `None` only for configuration or pre-session tasks.
-    async fn complete_with_model(
+    /// Primary streaming method that all providers must implement.
+    async fn stream(
        &self,
-        session_id: Option<&str>,
        model_config: &ModelConfig,
+        session_id: &str,
        system: &str,
        messages: &[Message],
        tools: &[Tool],
-    ) -> Result<(Message, ProviderUsage), ProviderError>;
+    ) -> Result<MessageStream, ProviderError>;

-    // Default implementation: use the provider's configured model
+    /// Complete with a specific model config.
    async fn complete(
        &self,
+        model_config: &ModelConfig,
        session_id: &str,
        system: &str,
        messages: &[Message],
        tools: &[Tool],
    ) -> Result<(Message, ProviderUsage), ProviderError> {
-        let model_config = self.get_model_config();
-        self.complete_with_model(Some(session_id), &model_config, system, messages, tools)
-            .await
+        let stream = self
+            .stream(model_config, session_id, system, messages, tools)
+            .await?;
+        collect_stream(stream).await
    }

-    // Check if a fast model is configured, otherwise fall back to regular model
+    /// Try fast model first, fall back to regular model on failure.
    async fn complete_fast(
        &self,
        session_id: &str,
@@ -417,11 +415,12 @@ pub trait Provider: Send + Sync {
        let model_config = self.get_model_config();
        let fast_config = model_config.use_fast_model();

-        match self
-            .complete_with_model(Some(session_id), &fast_config, system, messages, tools)
-            .await
-        {
-            Ok(result) => Ok(result),
+        let result = self
+            .complete(&fast_config, session_id, system, messages, tools)
+            .await;
+
+        match result {
+            Ok(response) => Ok(response),
            Err(e) => {
                if fast_config.model_name != model_config.model_name {
                    tracing::warn!(
@@ -430,14 +429,8 @@ pub trait Provider: Send + Sync {
                        e,
                        model_config.model_name
                    );
-                    self.complete_with_model(
-                        Some(session_id),
-                        &model_config,
-                        system,
-                        messages,
-                        tools,
-                    )
-                    .await
+                    self.complete(&model_config, session_id, system, messages, tools)
+                        .await
                } else {
                    Err(e)
                }
@@ -553,22 +546,6 @@ pub trait Provider: Send + Sync {
        None
    }

-    async fn stream(
-        &self,
-        _session_id: &str,
-        _system: &str,
-        _messages: &[Message],
-        _tools: &[Tool],
-    ) -> Result<MessageStream, ProviderError> {
-        Err(ProviderError::NotImplemented(
-            "streaming not implemented".to_string(),
-        ))
-    }
-
-    fn supports_streaming(&self) -> bool {
-        false
-    }
-
    /// Get the currently active model name
    /// For regular providers, this returns the configured model
    /// For LeadWorkerProvider, this returns the currently active model (lead or worker)
@@ -664,18 +641,149 @@ pub fn stream_from_single_message(message: Message, usage: ProviderUsage) -> Mes
    Box::pin(stream)
 }

+/// Collect all chunks from a MessageStream into a single Message and ProviderUsage
+pub async fn collect_stream(
+    mut stream: MessageStream,
+) -> Result<(Message, ProviderUsage), ProviderError> {
+    use futures::StreamExt;
+
+    let mut final_message: Option<Message> = None;
+    let mut final_usage: Option<ProviderUsage> = None;
+
+    while let Some(result) = stream.next().await {
+        let (msg_opt, usage_opt) = result?;
+
+        if let Some(msg) = msg_opt {
+            final_message = Some(match final_message {
+                Some(mut prev) => {
+                    for new_content in msg.content {
+                        match (&mut prev.content.last_mut(), &new_content) {
+                            // Coalesce consecutive text blocks
+                            (
+                                Some(MessageContent::Text(last_text)),
+                                MessageContent::Text(new_text),
+                            ) => {
+                                last_text.text.push_str(&new_text.text);
+                            }
+                            _ => {
+                                prev.content.push(new_content);
+                            }
+                        }
+                    }
+                    prev
+                }
+                None => msg,
+            });
+        }
+
+        if let Some(usage) = usage_opt {
+            final_usage = Some(usage);
+        }
+    }
+
+    match final_message {
+        Some(msg) => {
+            let usage = final_usage
+                .unwrap_or_else(|| ProviderUsage::new("unknown".to_string(), Usage::default()));
+            Ok((msg, usage))
+        }
+        None => Err(ProviderError::ExecutionError(
+            "Stream yielded no message".to_string(),
+        )),
+    }
+}
+
 #[cfg(test)]
 mod tests {
    use super::*;
    use std::collections::HashMap;
+    use test_case::test_case;

    use serde_json::json;
-    #[test]
-    fn test_usage_creation() {
-        let usage = Usage::new(Some(10), Some(20), Some(30));
-        assert_eq!(usage.input_tokens, Some(10));
-        assert_eq!(usage.output_tokens, Some(20));
-        assert_eq!(usage.total_tokens, Some(30));
+    fn content_from_str(s: String) -> MessageContent {
+        if let Some(img_data) = s.strip_prefix("*img:") {
+            MessageContent::image(format!("http://example.com/{}", img_data), "image/png")
+        } else if let Some(tool_name) = s.strip_prefix("*tool:") {
+            let tool_call = Ok(rmcp::model::CallToolRequestParams {
+                meta: None,
+                task: None,
+                name: tool_name.to_string().into(),
+                arguments: Some(serde_json::Map::new()),
+            });
+            MessageContent::tool_request(format!("tool_{}", tool_name), tool_call)
+        } else {
+            MessageContent::text(s)
+        }
+    }
+
+    fn create_test_stream(
+        items: Vec<String>,
+    ) -> impl Stream<Item = Result<(Option<Message>, Option<ProviderUsage>), ProviderError>> {
+        use futures::stream;
+        stream::iter(items.into_iter().map(|item| {
+            let content = content_from_str(item);
+            let message = Message::new(
+                rmcp::model::Role::Assistant,
+                chrono::Utc::now().timestamp(),
+                vec![content],
+            );
+            Ok((Some(message), None))
+        }))
+    }
+
+    fn content_to_strings(msg: &Message) -> Vec<String> {
+        msg.content
+            .iter()
+            .map(|c| match c {
+                MessageContent::Text(t) => t.text.clone(),
+                MessageContent::Image(_) => "*img".to_string(),
+                MessageContent::ToolRequest(tr) => {
+                    if let Ok(call) = &tr.tool_call {
+                        format!("*tool:{}", call.name)
+                    } else {
+                        "*tool:error".to_string()
+                    }
+                }
+                _ => "*other".to_string(),
+            })
+            .collect()
+    }
+
+    #[test_case(
+        vec!["Hello", " ", "world"],
+        vec!["Hello world"]
+        ; "consecutive text coalesces"
+    )]
+    #[test_case(
+        vec!["Hello", "*img:pic1", "world"],
+        vec!["Hello", "*img", "world"]
+        ; "non-text breaks coalescing"
+    )]
+    #[test_case(
+        vec!["A", "B", "*img:pic1", "C", "D", "*tool:read", "E", "F"],
+        vec!["AB", "*img", "CD", "*tool:read", "EF"]
+        ; "multiple text groups"
+    )]
+    #[test_case(
+        vec!["Text1", "*img:pic", "Text2"],
+        vec!["Text1", "*img", "Text2"]
+        ; "mixed content in chunk"
+    )]
+    #[tokio::test]
+    async fn test_collect_stream_coalescing(input_items: Vec<&str>, expected: Vec<&str>) {
+        let items: Vec<String> = input_items.into_iter().map(|s| s.to_string()).collect();
+        let stream = create_test_stream(items);
+        let (msg, _) = collect_stream(Box::pin(stream)).await.unwrap();
+        assert_eq!(content_to_strings(&msg), expected);
+    }
+
+    #[tokio::test]
+    async fn test_collect_stream_defaults_usage() {
+        // Should not error when usage is missing
+        let stream = create_test_stream(vec!["Hello".to_string()]);
+        let (msg, usage) = collect_stream(Box::pin(stream)).await.unwrap();
+        assert_eq!(content_to_strings(&msg), vec!["Hello"]);
+        assert_eq!(usage.model, "unknown"); // Default usage
    }

    #[test]