feat: stream LLM responses

Cargo.toml

@@ -15,4 +15,4 @@ uninlined_format_args = "allow"
 
 # Patch for Windows cross-compilation issue with crunchy
 [patch.crates-io]
-crunchy = { git = "https://github.com/nmathewson/crunchy", branch = "cross-compilation-fix" }
+crunchy = { git = "https://github.com/nmathewson/crunchy", branch = "cross-compilation-fix" }

crates/goose-cli/src/session/mod.rs

@@ -10,6 +10,7 @@ pub use self::export::message_to_markdown;
 pub use builder::{build_session, SessionBuilderConfig, SessionSettings};
 use console::Color;
 use goose::agents::AgentEvent;
+use goose::message::push_message;
 use goose::permission::permission_confirmation::PrincipalType;
 use goose::permission::Permission;
 use goose::permission::PermissionConfirmation;
@@ -317,7 +318,7 @@ impl Session {
 
     /// Process a single message and get the response
     async fn process_message(&mut self, message: String) -> Result<()> {
-        self.messages.push(Message::user().with_text(&message));
+        self.push_message(Message::user().with_text(&message));
         // Get the provider from the agent for description generation
         let provider = self.agent.provider().await?;
 
@@ -417,7 +418,7 @@ impl Session {
                         RunMode::Normal => {
                             save_history(&mut editor);
 
-                            self.messages.push(Message::user().with_text(&content));
+                            self.push_message(Message::user().with_text(&content));
 
                             // Track the current directory and last instruction in projects.json
                             let session_id = self
@@ -740,7 +741,7 @@ impl Session {
                     self.messages.clear();
                     // add the plan response as a user message
                     let plan_message = Message::user().with_text(plan_response.as_concat_text());
-                    self.messages.push(plan_message);
+                    self.push_message(plan_message);
                     // act on the plan
                     output::show_thinking();
                     self.process_agent_response(true).await?;
@@ -755,13 +756,13 @@ impl Session {
                 } else {
                     // add the plan response (assistant message) & carry the conversation forward
                     // in the next round, the user might wanna slightly modify the plan
-                    self.messages.push(plan_response);
+                    self.push_message(plan_response);
                 }
             }
             PlannerResponseType::ClarifyingQuestions => {
                 // add the plan response (assistant message) & carry the conversation forward
                 // in the next round, the user will answer the clarifying questions
-                self.messages.push(plan_response);
+                self.push_message(plan_response);
             }
         }
 
@@ -833,7 +834,7 @@ impl Session {
                                         confirmation.id.clone(),
                                         Err(ToolError::ExecutionError("Tool call cancelled by user".to_string()))
                                     ));
-                                    self.messages.push(response_message);
+                                    push_message(&mut self.messages, response_message);
                                     if let Some(session_file) = &self.session_file {
                                         session::persist_messages_with_schedule_id(
                                             session_file,
@@ -930,7 +931,7 @@ impl Session {
                             }
                             // otherwise we have a model/tool to render
                             else {
-                                self.messages.push(message.clone());
+                                push_message(&mut self.messages, message.clone());
 
                                 // No need to update description on assistant messages
                                 if let Some(session_file) = &self.session_file {
@@ -946,7 +947,6 @@ impl Session {
                                 if interactive {output::hide_thinking()};
                                 let _ = progress_bars.hide();
                                 output::render_message(&message, self.debug);
-                                if interactive {output::show_thinking()};
                             }
                         }
                         Some(Ok(AgentEvent::McpNotification((_id, message)))) => {
@@ -1094,6 +1094,7 @@ impl Session {
                 }
             }
         }
+        println!();
 
         Ok(())
     }
@@ -1137,7 +1138,7 @@ impl Session {
                     Err(ToolError::ExecutionError(notification.clone())),
                 ));
             }
-            self.messages.push(response_message);
+            self.push_message(response_message);
 
             // No need for description update here
             if let Some(session_file) = &self.session_file {
@@ -1154,7 +1155,7 @@ impl Session {
                 "The existing call to {} was interrupted. How would you like to proceed?",
                 last_tool_name
             );
-            self.messages.push(Message::assistant().with_text(&prompt));
+            self.push_message(Message::assistant().with_text(&prompt));
 
             // No need for description update here
             if let Some(session_file) = &self.session_file {
@@ -1176,7 +1177,7 @@ impl Session {
                         Some(MessageContent::ToolResponse(_)) => {
                             // Interruption occurred after a tool had completed but not assistant reply
                             let prompt = "The tool calling loop was interrupted. How would you like to proceed?";
-                            self.messages.push(Message::assistant().with_text(prompt));
+                            self.push_message(Message::assistant().with_text(prompt));
 
                             // No need for description update here
                             if let Some(session_file) = &self.session_file {
@@ -1364,7 +1365,7 @@ impl Session {
                         if msg.role == mcp_core::Role::User {
                             output::render_message(&msg, self.debug);
                         }
-                        self.messages.push(msg);
+                        self.push_message(msg);
                     }
 
                     if valid {
@@ -1422,6 +1423,10 @@ impl Session {
 
         Ok(path)
     }
+
+    fn push_message(&mut self, message: Message) {
+        push_message(&mut self.messages, message);
+    }
 }
 
 fn get_reasoner() -> Result<Arc<dyn Provider>, anyhow::Error> {

crates/goose-cli/src/session/output.rs

@@ -8,7 +8,7 @@ use mcp_core::tool::ToolCall;
 use serde_json::Value;
 use std::cell::RefCell;
 use std::collections::HashMap;
-use std::io::Error;
+use std::io::{Error, Write};
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use std::time::Duration;
@@ -164,7 +164,8 @@ pub fn render_message(message: &Message, debug: bool) {
             }
         }
     }
-    println!();
+
+    let _ = std::io::stdout().flush();
 }
 
 pub fn render_text(text: &str, color: Option<Color>, dim: bool) {

crates/goose-server/src/routes/reply.rs

@@ -225,6 +225,7 @@ async fn handler(
                 return;
             }
         };
+        let saved_message_count = all_messages.len();
 
         loop {
             tokio::select! {
@@ -242,16 +243,6 @@ async fn handler(
                                 ).await;
                                 break;
                             }
-
-
-                            let session_path = session_path.clone();
-                            let messages = all_messages.clone();
-                            let provider = Arc::clone(provider.as_ref().unwrap());
-                            tokio::spawn(async move {
-                                if let Err(e) = session::persist_messages(&session_path, &messages, Some(provider)).await {
-                                    tracing::error!("Failed to store session history: {:?}", e);
-                                }
-                            });
                         }
                         Ok(Some(Ok(AgentEvent::ModelChange { model, mode }))) => {
                             if let Err(e) = stream_event(MessageEvent::ModelChange { model, mode }, &tx).await {
@@ -303,6 +294,17 @@ async fn handler(
             }
         }
 
+        if all_messages.len() > saved_message_count {
+            let provider = Arc::clone(provider.as_ref().unwrap());
+            tokio::spawn(async move {
+                if let Err(e) =
+                    session::persist_messages(&session_path, &all_messages, Some(provider)).await
+                {
+                    tracing::error!("Failed to store session history: {:?}", e);
+                }
+            });
+        }
+
         let _ = stream_event(
             MessageEvent::Finish {
                 reason: "stop".to_string(),

crates/goose/Cargo.toml

@@ -81,6 +81,7 @@ fs2 = "0.4.3"
 tokio-stream = "0.1.17"
 dashmap = "6.1"
 ahash = "0.8"
+tokio-util = "0.7.15"
 
 # Vector database for tool selection
 lancedb = "0.13"

crates/goose/examples/databricks_oauth.rs

@@ -2,8 +2,12 @@ use anyhow::Result;
 use dotenv::dotenv;
 use goose::{
     message::Message,
-    providers::{base::Provider, databricks::DatabricksProvider},
+    providers::{
+        base::{Provider, Usage},
+        databricks::DatabricksProvider,
+    },
 };
+use tokio_stream::StreamExt;
 
 #[tokio::main]
 async fn main() -> Result<()> {
@@ -20,21 +24,24 @@ async fn main() -> Result<()> {
     let message = Message::user().with_text("Tell me a short joke about programming.");
 
     // Get a response
-    let (response, usage) = provider
-        .complete("You are a helpful assistant.", &[message], &[])
+    let mut stream = provider
+        .stream("You are a helpful assistant.", &[message], &[])
         .await?;
 
-    // Print the response and usage statistics
     println!("\nResponse from AI:");
     println!("---------------");
-    for content in response.content {
-        dbg!(content);
+    let mut usage = Usage::default();
+    while let Some(Ok((msg, usage_part))) = stream.next().await {
+        dbg!(msg);
+        usage_part.map(|u| {
+            usage += u.usage;
+        });
     }
     println!("\nToken Usage:");
     println!("------------");
-    println!("Input tokens: {:?}", usage.usage.input_tokens);
-    println!("Output tokens: {:?}", usage.usage.output_tokens);
-    println!("Total tokens: {:?}", usage.usage.total_tokens);
+    println!("Input tokens: {:?}", usage.input_tokens);
+    println!("Output tokens: {:?}", usage.output_tokens);
+    println!("Total tokens: {:?}", usage.total_tokens);
 
     Ok(())
 }