From f89bbfc89a495a4498256c983d8270636e60776d Mon Sep 17 00:00:00 2001
From: Dhanji Prasanna <dhanji@squareup.com>
Date: Fri, 31 Oct 2025 14:48:36 +1100
Subject: [PATCH] fix final_output bug

---
 crates/g3-core/src/lib.rs            | 13 ++++++++++---
 crates/g3-providers/src/anthropic.rs | 18 +++++++++++++++---
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/crates/g3-core/src/lib.rs b/crates/g3-core/src/lib.rs
index c11084d..1b8b12e 100644
--- a/crates/g3-core/src/lib.rs
+++ b/crates/g3-core/src/lib.rs
@@ -2675,7 +2675,13 @@ Template:
                             ));
 
                             // Display tool execution result with proper indentation
-                            if tool_call.tool != "final_output" {
+                            if tool_call.tool == "final_output" {
+                                // For final_output, display the summary without truncation
+                                for line in tool_result.lines() {
+                                    self.ui_writer.update_tool_output_line(line);
+                                }
+                                self.ui_writer.println("");
+                            } else {
                                 let output_lines: Vec<&str> = tool_result.lines().collect();
 
                                 // Check if UI wants full output (machine mode) or truncated (human mode)
@@ -2723,8 +2729,9 @@ Template:
 
                             // Check if this was a final_output tool call
                             if tool_call.tool == "final_output" {
-                                // Don't return anything - the summary was already displayed during streaming
-                                // Returning it again would cause duplication
+                                // The summary was displayed above when we printed the tool result
+                                // Add it to full_response so it's included in the TaskResult
+                                full_response.push_str(&tool_result);
                                 self.ui_writer.println("");
                                 let _ttft =
                                     first_token_time.unwrap_or_else(|| stream_start.elapsed());
diff --git a/crates/g3-providers/src/anthropic.rs b/crates/g3-providers/src/anthropic.rs
index d3dfc52..fc2fc00 100644
--- a/crates/g3-providers/src/anthropic.rs
+++ b/crates/g3-providers/src/anthropic.rs
@@ -276,6 +276,7 @@ impl AnthropicProvider {
         let mut partial_tool_json = String::new(); // Accumulate partial JSON for tool calls
         let mut accumulated_usage: Option<Usage> = None;
         let mut byte_buffer = Vec::new(); // Buffer for incomplete UTF-8 sequences
+        let mut message_stopped = false; // Track if we've received message_stop
         
         while let Some(chunk_result) = stream.next().await {
             match chunk_result {
@@ -316,6 +317,12 @@ impl AnthropicProvider {
                             continue;
                         }
 
+                        // If we've already sent the final chunk, skip processing more events
+                        if message_stopped {
+                            debug!("Skipping event after message_stop: {}", line);
+                            continue;
+                        }
+
                         // Parse Server-Sent Events format
                         if let Some(data) = line.strip_prefix("data: ") {
                             if data == "[DONE]" {
@@ -451,6 +458,7 @@ impl AnthropicProvider {
                                         }
                                         "message_stop" => {
                                             debug!("Received message stop event");
+                                            message_stopped = true;
                                             let final_chunk = CompletionChunk {
                                                 content: String::new(),
                                                 finished: true,
@@ -460,7 +468,8 @@ impl AnthropicProvider {
                                             if tx.send(Ok(final_chunk)).await.is_err() {
                                                 debug!("Receiver dropped, stopping stream");
                                             }
-                                            return accumulated_usage;
+                                            // Don't return here - let the stream naturally exhaust
+                                            // This prevents dropping the sender prematurely
                                         }
                                         "error" => {
                                             if let Some(error) = event.error {
@@ -468,7 +477,7 @@ impl AnthropicProvider {
                                                 let _ = tx
                                                     .send(Err(anyhow!("Anthropic API error: {:?}", error)))
                                                     .await;
-                                                return accumulated_usage;
+                                                break; // Break to let stream exhaust naturally
                                             }
                                         }
                                         _ => {
@@ -487,7 +496,10 @@ impl AnthropicProvider {
                 Err(e) => {
                     error!("Stream error: {}", e);
                     let _ = tx.send(Err(anyhow!("Stream error: {}", e))).await;
-                    return accumulated_usage;
+                    // Don't return here either - let the stream exhaust naturally
+                    // The error has been sent to the receiver, so it will handle it
+                    // Breaking here ensures we clean up properly
+                    break;
                 }
             }
         }