feat: add streaming support to Tetrate Agent Router Service provider

crates/goose/src/providers/tetrate.rs

@@ -1,14 +1,22 @@
 use anyhow::Result;
+use async_stream::try_stream;
 use async_trait::async_trait;
-use serde_json::Value;
+use futures::TryStreamExt;
+use serde_json::{json, Value};
+use std::io;
+use tokio::pin;
+use tokio_stream::StreamExt;
+use tokio_util::codec::{FramedRead, LinesCodec};
+use tokio_util::io::StreamReader;
 
 use super::api_client::{ApiClient, AuthMethod};
-use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage};
+use super::base::{ConfigKey, MessageStream, Provider, ProviderMetadata, ProviderUsage, Usage};
 use super::errors::ProviderError;
+use super::formats::openai::response_to_streaming_message;
 use super::retry::ProviderRetry;
 use super::utils::{
     emit_debug_trace, get_model, handle_response_google_compat, handle_response_openai_compat,
-    is_google_model,
+    handle_status_openai_compat, is_google_model,
 };
 use crate::config::signup_tetrate::TETRATE_DEFAULT_MODEL;
 use crate::conversation::message::Message;
@@ -38,6 +46,7 @@ pub struct TetrateProvider {
     #[serde(skip)]
     api_client: ApiClient,
     model: ModelConfig,
+    supports_streaming: bool,
 }
 
 impl_provider_default!(TetrateProvider);
@@ -56,7 +65,11 @@ impl TetrateProvider {
             .with_header("HTTP-Referer", "https://block.github.io/goose")?
             .with_header("X-Title", "Goose")?;
 
-        Ok(Self { api_client, model })
+        Ok(Self {
+            api_client,
+            model,
+            supports_streaming: true,
+        })
     }
 
     async fn post(&self, payload: &Value) -> Result<Value, ProviderError> {
@@ -178,6 +191,49 @@ impl Provider for TetrateProvider {
         Ok((message, ProviderUsage::new(model, usage)))
     }
 
+    async fn stream(
+        &self,
+        system: &str,
+        messages: &[Message],
+        tools: &[Tool],
+    ) -> Result<MessageStream, ProviderError> {
+        let mut payload = create_request(
+            &self.model,
+            system,
+            messages,
+            tools,
+            &super::utils::ImageFormat::OpenAi,
+        )?;
+
+        // Enable streaming
+        payload["stream"] = json!(true);
+        payload["stream_options"] = json!({
+            "include_usage": true,
+        });
+
+        let response = self
+            .api_client
+            .response_post("v1/chat/completions", &payload)
+            .await?;
+
+        let response = handle_status_openai_compat(response).await?;
+        let stream = response.bytes_stream().map_err(io::Error::other);
+        let model_config = self.model.clone();
+
+        Ok(Box::pin(try_stream! {
+            let stream_reader = StreamReader::new(stream);
+            let framed = FramedRead::new(stream_reader, LinesCodec::new()).map_err(anyhow::Error::from);
+
+            let message_stream = response_to_streaming_message(framed);
+            pin!(message_stream);
+            while let Some(message) = message_stream.next().await {
+                let (message, usage) = message.map_err(|e| ProviderError::RequestFailed(format!("Stream decode error: {}", e)))?;
+                emit_debug_trace(&model_config, &payload, &message, &usage.as_ref().map(|f| f.usage).unwrap_or_default());
+                yield (message, usage);
+            }
+        }))
+    }
+
     /// Fetch supported models from Tetrate Agent Router Service API (only models with tool support)
     async fn fetch_supported_models(&self) -> Result<Option<Vec<String>>, ProviderError> {
         // Use the existing api_client which already has authentication configured
@@ -251,4 +307,8 @@ impl Provider for TetrateProvider {
         models.sort();
         Ok(Some(models))
     }
+
+    fn supports_streaming(&self) -> bool {
+        self.supports_streaming
+    }
 }