fix(google): Fix 400 Bad Request error with Gemini 3 thought signatures (#6035)

Signed-off-by: rabi <ramishra@redhat.com>
This commit is contained in:
Rabi Mishra
2025-12-17 19:08:46 +05:30
committed by GitHub
parent b401222b89
commit fc8a4b4adf
9 changed files with 273 additions and 57 deletions
+146 -24
View File
@@ -9,10 +9,25 @@ use rmcp::model::{
};
use std::borrow::Cow;
use crate::conversation::message::{Message, MessageContent};
use crate::conversation::message::{Message, MessageContent, ProviderMetadata};
use serde_json::{json, Map, Value};
use std::ops::Deref;
pub const THOUGHT_SIGNATURE_KEY: &str = "thoughtSignature";
pub fn metadata_with_signature(signature: &str) -> ProviderMetadata {
let mut map = ProviderMetadata::new();
map.insert(THOUGHT_SIGNATURE_KEY.to_string(), json!(signature));
map
}
pub fn get_thought_signature(metadata: &Option<ProviderMetadata>) -> Option<&str> {
metadata
.as_ref()
.and_then(|m| m.get(THOUGHT_SIGNATURE_KEY))
.and_then(|v| v.as_str())
}
/// Convert internal Message format to Google's API message specification
pub fn format_messages(messages: &[Message]) -> Vec<Value> {
messages
@@ -58,8 +73,8 @@ pub fn format_messages(messages: &[Message]) -> Vec<Value> {
let mut part = Map::new();
part.insert("functionCall".to_string(), json!(function_call_part));
if let Some(signature) = &request.thought_signature {
part.insert("thoughtSignature".to_string(), json!(signature));
if let Some(signature) = get_thought_signature(&request.metadata) {
part.insert(THOUGHT_SIGNATURE_KEY.to_string(), json!(signature));
}
parts.push(json!(part));
@@ -117,15 +132,44 @@ pub fn format_messages(messages: &[Message]) -> Vec<Value> {
if text.is_empty() {
text = "Tool call is done.".to_string();
}
parts.push(json!({
"functionResponse": {
"name": response.id,
"response": {"content": {"text": text}},
}}
));
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),
);
if let Some(signature) = get_thought_signature(&response.metadata) {
part.insert(
THOUGHT_SIGNATURE_KEY.to_string(),
json!(signature),
);
}
parts.push(json!(part));
}
Err(e) => {
parts.push(json!({"text":format!("Error: {}", 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),
);
if let Some(signature) = get_thought_signature(&response.metadata) {
part.insert(
THOUGHT_SIGNATURE_KEY.to_string(),
json!(signature),
);
}
parts.push(json!(part));
}
}
}
@@ -282,15 +326,30 @@ pub fn response_to_message(response: Value) -> Result<Message> {
.and_then(|parts| parts.as_array())
.unwrap_or(&binding);
// Track the last seen thought signature to use as fallback for function calls without one
// This handles cases where Google's API returns multiple function calls but only includes
// thoughtSignature on some of them
let mut last_signature: Option<String> = None;
let has_function_calls = parts.iter().any(|p| p.get("functionCall").is_some());
for part in parts {
let thought_signature = part
.get("thoughtSignature")
let signature = part
.get(THOUGHT_SIGNATURE_KEY)
.and_then(|v| v.as_str())
.map(|s| s.to_string());
if signature.is_some() {
last_signature = signature.clone();
}
if let Some(text) = part.get("text").and_then(|v| v.as_str()) {
if let Some(sig) = thought_signature {
content.push(MessageContent::thinking(text.to_string(), sig));
// Text is "thinking" only if:
// 1. It has a signature AND
// 2. The response also contains function calls (meaning this is reasoning before acting)
// If there are no function calls, this is the final response and should be shown
if let (Some(sig), true) = (&signature, has_function_calls) {
content.push(MessageContent::thinking(text.to_string(), sig.clone()));
} else {
content.push(MessageContent::text(text.to_string()));
}
@@ -316,16 +375,17 @@ pub fn response_to_message(response: Value) -> Result<Message> {
content.push(MessageContent::tool_request(id, Err(error)));
} else {
let parameters = function_call.get("args");
if let Some(params) = parameters {
content.push(MessageContent::tool_request_with_signature(
id,
Ok(CallToolRequestParam {
name: name.into(),
arguments: Some(object(params.clone())),
}),
thought_signature,
));
}
let arguments = parameters.map(|params| object(params.clone()));
let effective_signature = signature.as_deref().or(last_signature.as_deref());
let metadata = effective_signature.map(metadata_with_signature);
content.push(MessageContent::tool_request_with_metadata(
id,
Ok(CallToolRequestParam {
name: name.into(),
arguments,
}),
metadata.as_ref(),
));
}
}
}
@@ -919,4 +979,66 @@ mod tests {
assert_eq!(regular_field["type"], "number");
assert_eq!(regular_field["description"], "A regular number field");
}
fn google_response(parts: Vec<Value>) -> Value {
json!({"candidates": [{"content": {"role": "model", "parts": parts}}]})
}
fn tool_result(text: &str) -> CallToolResult {
CallToolResult {
content: vec![Content::text(text)],
structured_content: None,
is_error: Some(false),
meta: None,
}
}
#[test]
fn test_thought_signature_roundtrip() {
const SIG: &str = "thought_sig_abc";
let response_with_tools = google_response(vec![
json!({"text": "Let me think...", "thoughtSignature": SIG}),
json!({"functionCall": {"name": "shell", "args": {"cmd": "ls"}}, "thoughtSignature": SIG}),
json!({"functionCall": {"name": "read", "args": {}}}),
]);
let native = response_to_message(response_with_tools).unwrap();
assert_eq!(native.content.len(), 3, "Expected thinking + 2 tool calls");
let thinking = native.content[0]
.as_thinking()
.expect("Text with function calls should be Thinking");
assert_eq!(thinking.signature, SIG);
let req1 = native.content[1]
.as_tool_request()
.expect("Second part should be ToolRequest");
let req2 = native.content[2]
.as_tool_request()
.expect("Third part should be ToolRequest");
assert_eq!(get_thought_signature(&req1.metadata), Some(SIG));
assert_eq!(
get_thought_signature(&req2.metadata),
Some(SIG),
"Should inherit"
);
let tool_response = Message::user().with_tool_response_with_metadata(
req1.id.clone(),
Ok(tool_result("output")),
req1.metadata.as_ref(),
);
let google_out = format_messages(&[native, tool_response]);
assert_eq!(google_out[0]["parts"][0]["thoughtSignature"], SIG);
assert_eq!(google_out[1]["parts"][0]["thoughtSignature"], SIG);
let final_response =
google_response(vec![json!({"text": "Done!", "thoughtSignature": SIG})]);
let final_native = response_to_message(final_response).unwrap();
assert!(
final_native.content[0].as_text().is_some(),
"Text-only = final answer"
);
}
}