diff --git a/crates/goose/src/providers/formats/google.rs b/crates/goose/src/providers/formats/google.rs index cf5b89a2..5805917f 100644 --- a/crates/goose/src/providers/formats/google.rs +++ b/crates/goose/src/providers/formats/google.rs @@ -15,6 +15,8 @@ use serde_json::{json, Map, Value}; use std::ops::Deref; pub const THOUGHT_SIGNATURE_KEY: &str = "thoughtSignature"; +const SYNTHETIC_THOUGHT_SIGNATURE: &str = "skip_thought_signature_validator"; +const GEMINI25_DEFAULT_THINKING_BUDGET: i32 = 8192; pub fn metadata_with_signature(signature: &str) -> ProviderMetadata { let mut map = ProviderMetadata::new(); @@ -29,6 +31,36 @@ pub fn get_thought_signature(metadata: &Option) -> Option<&str .and_then(|v| v.as_str()) } +fn is_user_loop_boundary(message: &Message) -> bool { + message.role == Role::User + && message + .content + .iter() + .any(|content| !matches!(content, MessageContent::ToolResponse(_))) +} + +fn insert_thought_signature(part: &mut Map, signature: &str) { + part.insert(THOUGHT_SIGNATURE_KEY.to_string(), json!(signature)); +} + +fn maybe_insert_signature_from_metadata( + part: &mut Map, + metadata: &Option, +) { + if let Some(signature) = get_thought_signature(metadata) { + insert_thought_signature(part, signature); + } +} + +fn build_function_response_part(name: &str, text: String) -> Map { + let mut part = Map::new(); + let mut function_response = Map::new(); + function_response.insert("name".to_string(), json!(name)); + function_response.insert("response".to_string(), json!({"content": {"text": text}})); + part.insert("functionResponse".to_string(), json!(function_response)); + part +} + /// Convert internal Message format to Google's API message specification pub fn format_messages(messages: &[Message]) -> Vec { let filtered: Vec<_> = messages @@ -44,26 +76,27 @@ pub fn format_messages(messages: &[Message]) -> Vec { }) .collect(); - let last_assistant_idx = filtered + let active_loop_start_idx = filtered .iter() .enumerate() - .filter(|(_, m)| m.role != Role::User) - .map(|(i, _)| i) - .next_back(); + .rev() + .find(|(_, m)| is_user_loop_boundary(m)) + .map(|(i, _)| i); filtered .iter() .enumerate() - .map(|(idx, message)| { + .filter_map(|(idx, message)| { let role = if message.role == Role::User { "user" } else { "model" }; - let include_signature = match last_assistant_idx { - Some(last_idx) => idx >= last_idx, - None => false, - }; + let include_signature = active_loop_start_idx.is_none_or(|start_idx| idx >= start_idx); + // Only the first model tool call in a turn is guaranteed to carry + // a signature for loop continuity. + let mut needs_synthetic_for_first_model_tool_call = + include_signature && message.role != Role::User; let mut parts = Vec::new(); for message_content in message.content.iter() { match message_content { @@ -92,12 +125,15 @@ pub fn format_messages(messages: &[Message]) -> Vec { if include_signature { if let Some(signature) = get_thought_signature(&request.metadata) { - part.insert( - THOUGHT_SIGNATURE_KEY.to_string(), - json!(signature), + insert_thought_signature(&mut part, signature); + } else if needs_synthetic_for_first_model_tool_call { + insert_thought_signature( + &mut part, + SYNTHETIC_THOUGHT_SIGNATURE, ); } } + needs_synthetic_for_first_model_tool_call = false; parts.push(json!(part)); } @@ -138,50 +174,22 @@ pub fn format_messages(messages: &[Message]) -> Vec { if text.is_empty() { text = "Tool call is done.".to_string(); } - let mut part = Map::new(); - let mut function_response = Map::new(); - function_response.insert("name".to_string(), json!(response.id)); - function_response - .insert("response".to_string(), json!({"content": {"text": text}})); - part.insert("functionResponse".to_string(), json!(function_response)); + let mut part = build_function_response_part(&response.id, text); if include_signature { - if let Some(signature) = get_thought_signature(&response.metadata) { - part.insert( - THOUGHT_SIGNATURE_KEY.to_string(), - json!(signature), - ); - } + maybe_insert_signature_from_metadata(&mut part, &response.metadata); } parts.push(json!(part)); } Err(e) => { - let mut part = Map::new(); - let mut function_response = Map::new(); - function_response.insert("name".to_string(), json!(response.id)); - function_response.insert( - "response".to_string(), - json!({"content": {"text": format!("Error: {}", e)}}), - ); - part.insert("functionResponse".to_string(), json!(function_response)); + let mut part = + build_function_response_part(&response.id, format!("Error: {}", e)); if include_signature { - if let Some(signature) = get_thought_signature(&response.metadata) { - part.insert( - THOUGHT_SIGNATURE_KEY.to_string(), - json!(signature), - ); - } + maybe_insert_signature_from_metadata(&mut part, &response.metadata); } parts.push(json!(part)); } }, - MessageContent::Thinking(thinking) => { - let mut part = Map::new(); - part.insert("text".to_string(), json!(thinking.thinking)); - if include_signature { - part.insert("thoughtSignature".to_string(), json!(thinking.signature)); - } - parts.push(json!(part)); - } + MessageContent::Thinking(_) => {} MessageContent::Image(image) => { parts.push(json!({ "inline_data": { @@ -194,7 +202,11 @@ pub fn format_messages(messages: &[Message]) -> Vec { _ => {} } } - json!({"role": role, "parts": parts}) + if parts.is_empty() { + None + } else { + Some(json!({"role": role, "parts": parts})) + } }) .collect() } @@ -221,32 +233,15 @@ pub fn format_tools(tools: &[Tool]) -> Vec { .collect() } -#[derive(Clone, Copy)] -enum SignedTextHandling { - SignedTextAsThinking, - SignedTextAsRegularText, -} - -fn process_response_part_non_streaming( - part: &Value, - last_signature: &mut Option, - has_function_calls: bool, -) -> Option { - // For non-streaming: signed text is thinking only if there are function calls - let handling = if has_function_calls { - SignedTextHandling::SignedTextAsThinking - } else { - SignedTextHandling::SignedTextAsRegularText - }; - process_response_part_impl(part, last_signature, handling) -} - fn process_response_part_impl( part: &Value, last_signature: &mut Option, - signed_text_handling: SignedTextHandling, ) -> Option { let signature = part.get(THOUGHT_SIGNATURE_KEY).and_then(|v| v.as_str()); + let is_thought = part + .get("thought") + .and_then(|value| value.as_bool()) + .unwrap_or(false); if let Some(sig) = signature { *last_signature = Some(sig.to_string()); @@ -257,11 +252,13 @@ fn process_response_part_impl( if text.is_empty() { return None; } - match (signature, signed_text_handling) { - (Some(sig), SignedTextHandling::SignedTextAsThinking) => { - Some(MessageContent::thinking(text.to_string(), sig.to_string())) + if is_thought { + match signature { + Some(sig) => Some(MessageContent::thinking(text.to_string(), sig.to_string())), + None => Some(MessageContent::reasoning(text.to_string())), } - _ => Some(MessageContent::text(text.to_string())), + } else { + Some(MessageContent::text(text.to_string())) } } else if text_value.is_some() { tracing::warn!( @@ -323,15 +320,11 @@ pub fn response_to_message(response: Value) -> Result { return Ok(Message::new(role, created, Vec::new())); }; - let has_function_calls = parts.iter().any(|p| p.get("functionCall").is_some()); - let mut content = Vec::new(); let mut last_signature: Option = None; for part in parts { - if let Some(msg_content) = - process_response_part_non_streaming(part, &mut last_signature, has_function_calls) - { + if let Some(msg_content) = process_response_part_impl(part, &mut last_signature) { content.push(msg_content); } } @@ -467,9 +460,7 @@ where if let Some(parts) = parts { for part in parts { - // Always emit text as regular text during streaming — we can't - // know yet whether function calls will follow. - if let Some(content) = process_response_part_impl(part, &mut last_signature, SignedTextHandling::SignedTextAsRegularText) { + if let Some(content) = process_response_part_impl(part, &mut last_signature) { let message = Message::new( Role::Assistant, chrono::Utc::now().timestamp(), @@ -524,7 +515,11 @@ enum ThinkingLevel { #[derive(Serialize)] #[serde(rename_all = "camelCase")] struct ThinkingConfig { - thinking_level: ThinkingLevel, + #[serde(skip_serializing_if = "Option::is_none")] + thinking_level: Option, + #[serde(skip_serializing_if = "Option::is_none")] + thinking_budget: Option, + include_thoughts: bool, } #[derive(Serialize)] @@ -539,33 +534,59 @@ struct GoogleRequest<'a> { } fn get_thinking_config(model_config: &ModelConfig) -> Option { - if !model_config - .model_name - .to_lowercase() - .starts_with("gemini-3") - { + let model_name = model_config.model_name.to_lowercase(); + let is_gemini_3 = model_name.starts_with("gemini-3"); + let is_gemini_25 = model_name.starts_with("gemini-2.5"); + if !is_gemini_3 && !is_gemini_25 { return None; } - let thinking_level_str = model_config - .get_config_param::("thinking_level", "GEMINI3_THINKING_LEVEL") - .map(|s| s.to_lowercase()) - .unwrap_or_else(|| "low".to_string()); + if is_gemini_3 { + let thinking_level_str = model_config + .get_config_param::("thinking_level", "GEMINI3_THINKING_LEVEL") + .map(|s| s.to_lowercase()) + .unwrap_or_else(|| "low".to_string()); - let thinking_level = match thinking_level_str.as_str() { - "high" => ThinkingLevel::High, - "low" => ThinkingLevel::Low, - invalid => { - tracing::warn!( - "Invalid thinking level '{}' for model '{}'. Valid levels: low, high. Using 'low'.", - invalid, - model_config.model_name, - ); - ThinkingLevel::Low - } - }; + let thinking_level = match thinking_level_str.as_str() { + "high" => ThinkingLevel::High, + "low" => ThinkingLevel::Low, + invalid => { + tracing::warn!( + "Invalid thinking level '{}' for model '{}'. Valid levels: low, high. Using 'low'.", + invalid, + model_config.model_name, + ); + ThinkingLevel::Low + } + }; - Some(ThinkingConfig { thinking_level }) + Some(ThinkingConfig { + thinking_level: Some(thinking_level), + thinking_budget: None, + include_thoughts: true, + }) + } else { + let thinking_budget = match model_config + .get_config_param::("thinking_budget", "GEMINI25_THINKING_BUDGET") + { + Some(budget) if budget >= 0 => budget, + Some(budget) => { + tracing::warn!( + "Invalid thinking budget '{}' for model '{}'. Must be >= 0. Using '{}'.", + budget, + model_config.model_name, + GEMINI25_DEFAULT_THINKING_BUDGET, + ); + GEMINI25_DEFAULT_THINKING_BUDGET + } + None => GEMINI25_DEFAULT_THINKING_BUDGET, + }; + Some(ThinkingConfig { + thinking_level: None, + thinking_budget: Some(thinking_budget), + include_thoughts: true, + }) + } } pub fn create_request( @@ -609,6 +630,7 @@ mod tests { use rmcp::model::{CallToolRequestParams, CallToolResult}; use rmcp::{model::Content, object}; use serde_json::json; + use std::collections::HashMap; fn set_up_text_message(text: &str, role: Role) -> Message { Message::new(role, 0, vec![MessageContent::text(text.to_string())]) @@ -964,7 +986,7 @@ mod tests { const SIG: &str = "thought_sig_abc"; let response_with_tools = google_response(vec![ - json!({"text": "Let me think...", "thoughtSignature": SIG}), + json!({"text": "Let me think...", "thought": true, "thoughtSignature": SIG}), json!({"functionCall": {"name": "shell", "args": {"cmd": "ls"}}, "thoughtSignature": SIG}), json!({"functionCall": {"name": "read", "args": {}}}), ]); @@ -995,23 +1017,22 @@ mod tests { Ok(tool_result("output")), req1.metadata.as_ref(), ); - let google_out = format_messages(&[native.clone(), tool_response.clone()]); - assert_eq!(google_out[0]["parts"][0]["thoughtSignature"], SIG); + let user_prompt = set_up_text_message("List files", Role::User); + let google_out = + format_messages(&[user_prompt.clone(), native.clone(), tool_response.clone()]); assert_eq!(google_out[1]["parts"][0]["thoughtSignature"], SIG); + assert_eq!(google_out[2]["parts"][0]["thoughtSignature"], SIG); - let second_assistant = - Message::assistant().with_thinking("More thinking".to_string(), "sig_456".to_string()); - let google_multi = format_messages(&[native, tool_response, second_assistant]); - assert!(google_multi[0]["parts"][0] - .get("thoughtSignature") - .is_none()); - assert!(google_multi[1]["parts"][0] - .get("thoughtSignature") - .is_none()); - assert_eq!(google_multi[2]["parts"][0]["thoughtSignature"], "sig_456"); + let second_assistant = response_to_message(google_response(vec![json!({ + "functionCall": {"name": "echo", "args": {}}, + "thoughtSignature": "sig_456" + })])) + .unwrap(); + let google_multi = format_messages(&[user_prompt, native, tool_response, second_assistant]); + assert_eq!(google_multi[1]["parts"][0]["thoughtSignature"], SIG); + assert_eq!(google_multi[2]["parts"][0]["thoughtSignature"], SIG); + assert_eq!(google_multi[3]["parts"][0]["thoughtSignature"], "sig_456"); - // Text-only response WITH signature but WITHOUT function calls should be regular text - // (per original behavior: thinking is only when reasoning before tool calls) let final_response_with_sig = google_response(vec![json!({"text": "Done!", "thoughtSignature": SIG})]); let final_native_with_sig = response_to_message(final_response_with_sig).unwrap(); @@ -1028,6 +1049,49 @@ mod tests { ); } + #[test] + fn test_thought_without_signature_maps_to_reasoning() { + let response = google_response(vec![json!({ + "text": "Working through options...", + "thought": true + })]); + let native = response_to_message(response).unwrap(); + assert_eq!(native.content.len(), 1); + assert!(native.content[0].as_reasoning().is_some()); + } + + #[test] + fn test_format_messages_omits_messages_with_empty_parts() { + let user_prompt = set_up_text_message("hello", Role::User); + let thinking_only = + Message::assistant().with_thinking("internal".to_string(), "sig_123".to_string()); + let reasoning_only = response_to_message(google_response(vec![json!({ + "text": "deliberating", + "thought": true + })])) + .unwrap(); + + let formatted = format_messages(&[user_prompt, thinking_only, reasoning_only]); + assert_eq!(formatted.len(), 1); + assert_eq!(formatted[0]["role"], "user"); + assert_eq!(formatted[0]["parts"][0]["text"], "hello"); + } + + #[test] + fn test_active_loop_injects_synthetic_signature_for_first_model_tool_call() { + let user_prompt = set_up_text_message("Find a restaurant", Role::User); + let assistant_tool = response_to_message(google_response(vec![json!({ + "functionCall": {"name": "find_restaurant", "args": {"cuisine": "italian"}} + })])) + .unwrap(); + + let formatted = format_messages(&[user_prompt, assistant_tool]); + assert_eq!( + formatted[1]["parts"][0][THOUGHT_SIGNATURE_KEY], + SYNTHETIC_THOUGHT_SIGNATURE + ); + } + const GOOGLE_TEXT_STREAM: &str = concat!( r#"data: {"candidates": [{"content": {"role": "model", "#, r#""parts": [{"text": "Hello"}]}}]}"#, @@ -1082,7 +1146,6 @@ mod tests { assert_eq!(usage.usage.input_tokens, Some(10)); assert_eq!(usage.usage.output_tokens, Some(3)); - // Verify all streaming messages have consistent IDs for UI aggregation assert!( message_ids.iter().all(|id| id.is_some()), "All streaming messages should have an ID" @@ -1125,29 +1188,30 @@ mod tests { async fn test_streaming_with_thought_signature() { use futures::StreamExt; - async fn collect_streaming_text(raw: &str) -> (String, usize) { + async fn collect_streaming_text(raw: &str) -> (String, usize, usize) { let lines: Vec> = raw.lines().map(|l| Ok(l.to_string())).collect(); let stream = Box::pin(futures::stream::iter(lines)); let mut msg_stream = std::pin::pin!(response_to_streaming_message(stream)); let mut text = String::new(); let mut thinking = 0usize; + let mut reasoning = 0usize; while let Some(Ok((message, _))) = msg_stream.next().await { if let Some(msg) = message { for c in &msg.content { match c { MessageContent::Text(t) => text.push_str(&t.text), MessageContent::Thinking(_) => thinking += 1, + MessageContent::Reasoning(_) => reasoning += 1, _ => {} } } } } - (text, thinking) + (text, thinking, reasoning) } - // First chunk signed - let (text, thinking) = collect_streaming_text(concat!( + let (text, thinking, reasoning) = collect_streaming_text(concat!( r#"data: {"candidates": [{"content": {"role": "model", "#, r#""parts": [{"text": "Hello", "thoughtSignature": "sig1"}]}}], "#, r#""modelVersion": "gemini-3-flash-preview"}"#, @@ -1157,10 +1221,10 @@ mod tests { )) .await; assert_eq!(thinking, 0); + assert_eq!(reasoning, 0); assert_eq!(text, "Hello world"); - // Last chunk signed (the reported truncation bug) - let (text, thinking) = collect_streaming_text(concat!( + let (text, thinking, reasoning) = collect_streaming_text(concat!( r#"data: {"candidates": [{"content": {"role": "model", "#, r#""parts": [{"text": "SECURITY.md: Project"}]}}], "#, r#""modelVersion": "gemini-3-flash-preview"}"#, @@ -1171,10 +1235,10 @@ mod tests { )) .await; assert_eq!(thinking, 0); + assert_eq!(reasoning, 0); assert_eq!(text, "SECURITY.md: Project policies.\n\nRead it?"); - // Intermediate chunk signed - let (text, thinking) = collect_streaming_text(concat!( + let (text, thinking, reasoning) = collect_streaming_text(concat!( r#"data: {"candidates": [{"content": {"role": "model", "#, r#""parts": [{"text": "one "}]}}], "modelVersion": "gemini-3-flash-preview"}"#, "\n", @@ -1186,7 +1250,20 @@ mod tests { )) .await; assert_eq!(thinking, 0); + assert_eq!(reasoning, 0); assert_eq!(text, "one two three"); + + let (text, thinking, reasoning) = collect_streaming_text(concat!( + r#"data: {"candidates": [{"content": {"role": "model", "#, + r#""parts": [{"text": "internal chain", "thought": true, "thoughtSignature": "sig4"}]}}]}"#, + "\n", + r#"data: {"candidates": [{"content": {"role": "model", "#, + r#""parts": [{"text": "visible"}]}}]}"# + )) + .await; + assert_eq!(thinking, 1); + assert_eq!(reasoning, 0); + assert_eq!(text, "visible"); } #[tokio::test] @@ -1215,7 +1292,6 @@ mod tests { async fn test_streaming_with_sse_event_lines() { use futures::StreamExt; - // SSE format can include event: lines which should be skipped let sse_stream = r#"event: message data: {"candidates": [{"content": {"role": "model", "parts": [{"text": "Hello"}]}}]} @@ -1272,7 +1348,6 @@ data: [DONE]"#; } } - // Only "Complete" should be captured, stream should stop at [DONE] assert_eq!(text_parts, vec!["Complete"]); } @@ -1308,19 +1383,56 @@ data: [DONE]"#; fn test_get_thinking_config() { use crate::model::ModelConfig; - // Test 1: Gemini 3 model defaults to low thinking level let config = ModelConfig::new("gemini-3-pro").unwrap(); let result = get_thinking_config(&config); assert!(result.is_some()); let thinking_config = result.unwrap(); - assert!(matches!(thinking_config.thinking_level, ThinkingLevel::Low)); + assert!(thinking_config.thinking_level.is_some()); + assert!(thinking_config.thinking_budget.is_none()); + assert!(thinking_config.include_thoughts); - // Test 2: Case-insensitive model detection let config = ModelConfig::new("Gemini-3-Flash").unwrap(); let result = get_thinking_config(&config); assert!(result.is_some()); - // Test 3: Non-Gemini 3 model returns None + let config = ModelConfig::new("gemini-2.5-flash").unwrap(); + let result = get_thinking_config(&config); + assert!(result.is_some()); + let thinking_config = result.unwrap(); + assert!(thinking_config.include_thoughts); + assert!(thinking_config.thinking_level.is_none()); + assert_eq!( + thinking_config.thinking_budget, + Some(GEMINI25_DEFAULT_THINKING_BUDGET) + ); + + let mut params = HashMap::new(); + params.insert("thinking_budget".to_string(), json!(4096)); + let config = ModelConfig::new("gemini-2.5-flash") + .unwrap() + .with_request_params(Some(params)); + let result = get_thinking_config(&config); + assert!(result.is_some()); + let thinking_config = result.unwrap(); + assert_eq!(thinking_config.thinking_budget, Some(4096)); + + let mut params = HashMap::new(); + params.insert("thinking_budget".to_string(), json!(-1)); + let config = ModelConfig::new("gemini-2.5-flash") + .unwrap() + .with_request_params(Some(params)); + let result = get_thinking_config(&config); + assert!(result.is_some()); + let thinking_config = result.unwrap(); + assert_eq!( + thinking_config.thinking_budget, + Some(GEMINI25_DEFAULT_THINKING_BUDGET) + ); + + let config = ModelConfig::new("gemini-2.0-flash").unwrap(); + let result = get_thinking_config(&config); + assert!(result.is_none()); + let config = ModelConfig::new("gpt-4o").unwrap(); let result = get_thinking_config(&config); assert!(result.is_none());