From fe6064188735cb3da9df785020073e25aa701f49 Mon Sep 17 00:00:00 2001
From: centdix <40307056+centdix@users.noreply.github.com>
Date: Mon, 15 Dec 2025 23:53:13 +0100
Subject: [PATCH] fix(aiagent): fix gemini-3.0 usage (#7382)

* fix gemini for ai agent

* no clone
---
 .../windmill-worker/src/ai/providers/bedrock.rs |  2 ++
 .../windmill-worker/src/ai/providers/openai.rs  | 17 +++++++++++++++++
 backend/windmill-worker/src/ai/sse.rs           | 15 ++++++++++++---
 3 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/backend/windmill-worker/src/ai/providers/bedrock.rs b/backend/windmill-worker/src/ai/providers/bedrock.rs
index 18d08d0a41..57f44a8ab7 100644
--- a/backend/windmill-worker/src/ai/providers/bedrock.rs
+++ b/backend/windmill-worker/src/ai/providers/bedrock.rs
@@ -464,6 +464,7 @@ pub fn bedrock_response_to_openai(
                                 name: tool_use.name().to_string(),
                                 arguments,
                             },
+                            extra_content: None, // Bedrock doesn't use thought signatures
                         });
                     }
                     _ => {}
@@ -544,6 +545,7 @@ pub fn streaming_tool_calls_to_openai(tool_calls: Vec<StreamingToolCall>) -> Vec
             id: tc.id,
             function: OpenAIFunction { name: tc.name, arguments: tc.arguments },
             r#type: FUNCTION_TYPE.to_string(),
+            extra_content: None, // Bedrock doesn't use thought signatures
         })
         .collect()
 }
diff --git a/backend/windmill-worker/src/ai/providers/openai.rs b/backend/windmill-worker/src/ai/providers/openai.rs
index b3205b529a..178f3b23a4 100644
--- a/backend/windmill-worker/src/ai/providers/openai.rs
+++ b/backend/windmill-worker/src/ai/providers/openai.rs
@@ -17,11 +17,28 @@ pub struct OpenAIFunction {
     pub arguments: String,
 }
 
+/// Google-specific extra content for thought signatures (Gemini 3 Pro / 2.5)
+#[derive(Deserialize, Serialize, Clone, Debug, Default)]
+pub struct GoogleExtraContent {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub thought_signature: Option<String>,
+}
+
+/// Extra content for provider-specific metadata (e.g., Google thought signatures)
+#[derive(Deserialize, Serialize, Clone, Debug, Default)]
+pub struct ExtraContent {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub google: Option<GoogleExtraContent>,
+}
+
 #[derive(Deserialize, Serialize, Clone, Debug)]
 pub struct OpenAIToolCall {
     pub id: String,
     pub function: OpenAIFunction,
     pub r#type: String,
+    /// Extra content for provider-specific metadata (e.g., Google Gemini thought signatures)
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub extra_content: Option<ExtraContent>,
 }
 
 #[derive(Deserialize)]
diff --git a/backend/windmill-worker/src/ai/sse.rs b/backend/windmill-worker/src/ai/sse.rs
index 16cc418d52..91dcd6e94b 100644
--- a/backend/windmill-worker/src/ai/sse.rs
+++ b/backend/windmill-worker/src/ai/sse.rs
@@ -8,7 +8,7 @@ use tokio_stream::StreamExt;
 use windmill_common::{error::Error, utils::rd_string};
 
 use crate::ai::{
-    providers::openai::{OpenAIFunction, OpenAIToolCall},
+    providers::openai::{ExtraContent, OpenAIFunction, OpenAIToolCall},
     query_builder::StreamEventProcessor,
     types::StreamingEvent,
 };
@@ -24,6 +24,8 @@ pub struct OpenAIChoiceDeltaToolCall {
     pub index: Option<i64>,
     pub id: Option<String>,
     pub function: Option<OpenAIChoiceDeltaToolCallFunction>,
+    /// Extra content for provider-specific metadata (e.g., Google Gemini thought signatures)
+    pub extra_content: Option<ExtraContent>,
 }
 
 #[derive(Deserialize)]
@@ -119,9 +121,15 @@ impl SSEParser for OpenAISSEParser {
                             let idx = tool_call.index.unwrap_or_else(|| idx as i64);
 
                             if let Some(function) = tool_call.function {
-                                if let Some(tool_call) = self.accumulated_tool_calls.get_mut(&idx) {
+                                if let Some(existing_tool_call) =
+                                    self.accumulated_tool_calls.get_mut(&idx)
+                                {
                                     if let Some(arguments) = function.arguments {
-                                        tool_call.function.arguments += &arguments;
+                                        existing_tool_call.function.arguments += &arguments;
+                                    }
+                                    // Update extra_content if provided in this delta (for thought signatures)
+                                    if let Some(extra) = tool_call.extra_content {
+                                        existing_tool_call.extra_content = Some(extra);
                                     }
                                 } else {
                                     let fun_name = function.name.unwrap_or_default();
@@ -142,6 +150,7 @@ impl SSEParser for OpenAISSEParser {
                                                 arguments: function.arguments.unwrap_or_default(),
                                             },
                                             r#type: "function".to_string(),
+                                            extra_content: tool_call.extra_content,
                                         },
                                     );
                                 }