From e4938c69e0198acb6d027ffb7f26af4394b6526c Mon Sep 17 00:00:00 2001
From: Rabi Mishra <ramishra@redhat.com>
Date: Wed, 28 Jan 2026 22:47:08 +0530
Subject: [PATCH] fix(google): preserve thought signatures in streaming
 responses (#6708)

Signed-off-by: rabi <ramishra@redhat.com>
---
 crates/goose/src/providers/formats/google.rs | 28 +++++++++++++-------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/crates/goose/src/providers/formats/google.rs b/crates/goose/src/providers/formats/google.rs
index 61c4fa5f3d..32f7545421 100644
--- a/crates/goose/src/providers/formats/google.rs
+++ b/crates/goose/src/providers/formats/google.rs
@@ -252,12 +252,13 @@ pub fn process_response_part(
     part: &Value,
     last_signature: &mut Option<String>,
 ) -> Option<MessageContent> {
-    // Gemini 2.5 models include thoughtSignature on the first streaming chunk
-    process_response_part_impl(
-        part,
-        last_signature,
-        SignedTextHandling::SignedTextAsRegularText,
-    )
+    let has_signature = part.get(THOUGHT_SIGNATURE_KEY).is_some();
+    let handling = if has_signature {
+        SignedTextHandling::SignedTextAsThinking
+    } else {
+        SignedTextHandling::SignedTextAsRegularText
+    };
+    process_response_part_impl(part, last_signature, handling)
 }
 
 fn process_response_part_non_streaming(
@@ -1116,17 +1117,26 @@ mod tests {
         let mut message_stream = std::pin::pin!(response_to_streaming_message(stream));
 
         let mut text_parts = Vec::new();
+        let mut thinking_parts = Vec::new();
 
         while let Some(result) = message_stream.next().await {
             let (message, _usage) = result.unwrap();
             if let Some(msg) = message {
-                if let Some(MessageContent::Text(text)) = msg.content.first() {
-                    text_parts.push(text.text.clone());
+                match msg.content.first() {
+                    Some(MessageContent::Text(text)) => {
+                        text_parts.push(text.text.clone());
+                    }
+                    Some(MessageContent::Thinking(thinking)) => {
+                        thinking_parts.push(thinking.thinking.clone());
+                        assert_eq!(thinking.signature, "sig123");
+                    }
+                    _ => {}
                 }
             }
         }
 
-        assert_eq!(text_parts, vec!["Begin", " middle", " end"]);
+        assert_eq!(thinking_parts, vec!["Begin"]);
+        assert_eq!(text_parts, vec![" middle", " end"]);
     }
 
     #[tokio::test]