In the provider compatibility layer for Gemini (and other providers requiring ), fully support the flow, round-trip, and placeholder fallback of thought_signature.

This commit is contained in:
sunmanbitch
2026-06-04 21:13:55 +08:00
parent 4619375c14
commit b119afcaca
11 changed files with 193 additions and 43 deletions

View File

@@ -5203,7 +5203,7 @@ async fn stream_with_provider(
) -> Result<Vec<AssistantEvent>, ApiError> {
let mut stream = client.stream_message(message_request).await?;
let mut events = Vec::new();
let mut pending_tools: BTreeMap<u32, (String, String, String)> = BTreeMap::new();
let mut pending_tools: BTreeMap<u32, (String, String, String, Option<String>)> = BTreeMap::new();
let mut pending_thinking: BTreeMap<u32, (String, Option<String>)> = BTreeMap::new();
let mut saw_stop = false;
@@ -5238,7 +5238,7 @@ async fn stream_with_provider(
}
}
ContentBlockDelta::InputJsonDelta { partial_json } => {
if let Some((_, _, input)) = pending_tools.get_mut(&delta.index) {
if let Some((_, _, input, _)) = pending_tools.get_mut(&delta.index) {
input.push_str(&partial_json);
}
}
@@ -5262,8 +5262,8 @@ async fn stream_with_provider(
signature,
});
}
if let Some((id, name, input)) = pending_tools.remove(&stop.index) {
events.push(AssistantEvent::ToolUse { id, name, input });
if let Some((id, name, input, thought_signature)) = pending_tools.remove(&stop.index) {
events.push(AssistantEvent::ToolUse { id, name, input, thought_signature });
}
}
ApiStreamEvent::MessageDelta(delta) => {
@@ -5371,11 +5371,12 @@ fn convert_messages(messages: &[ConversationMessage]) -> Vec<InputMessage> {
thinking: thinking.clone(),
signature: signature.clone(),
},
ContentBlock::ToolUse { id, name, input } => InputContentBlock::ToolUse {
ContentBlock::ToolUse { id, name, input, thought_signature } => InputContentBlock::ToolUse {
id: id.clone(),
name: name.clone(),
input: serde_json::from_str(input)
.unwrap_or_else(|_| serde_json::json!({ "raw": input })),
thought_signature: thought_signature.clone(),
},
ContentBlock::ToolResult {
tool_use_id,
@@ -5406,7 +5407,7 @@ fn push_output_block(
block: OutputContentBlock,
block_index: u32,
events: &mut Vec<AssistantEvent>,
pending_tools: &mut BTreeMap<u32, (String, String, String)>,
pending_tools: &mut BTreeMap<u32, (String, String, String, Option<String>)>,
pending_thinking: &mut BTreeMap<u32, (String, Option<String>)>,
streaming_tool_input: bool,
) {
@@ -5416,7 +5417,7 @@ fn push_output_block(
events.push(AssistantEvent::TextDelta(text));
}
}
OutputContentBlock::ToolUse { id, name, input } => {
OutputContentBlock::ToolUse { id, name, input, thought_signature } => {
let initial_input = if streaming_tool_input
&& input.is_object()
&& input.as_object().is_some_and(serde_json::Map::is_empty)
@@ -5425,7 +5426,7 @@ fn push_output_block(
} else {
input.to_string()
};
pending_tools.insert(block_index, (id, name, initial_input));
pending_tools.insert(block_index, (id, name, initial_input, thought_signature));
}
OutputContentBlock::Thinking {
thinking,
@@ -5459,8 +5460,8 @@ fn response_to_events(response: MessageResponse) -> Vec<AssistantEvent> {
&mut pending_thinking,
false,
);
if let Some((id, name, input)) = pending_tools.remove(&index) {
events.push(AssistantEvent::ToolUse { id, name, input });
if let Some((id, name, input, thought_signature)) = pending_tools.remove(&index) {
events.push(AssistantEvent::ToolUse { id, name, input, thought_signature });
}
}
@@ -8066,6 +8067,7 @@ mod tests {
id: "tool-1".to_string(),
name: "read_file".to_string(),
input: json!({}),
thought_signature: None,
},
1,
&mut events,
@@ -8078,6 +8080,7 @@ mod tests {
id: "tool-2".to_string(),
name: "grep_search".to_string(),
input: json!({}),
thought_signature: None,
},
2,
&mut events,
@@ -9358,6 +9361,7 @@ mod tests {
id: "tool-1".to_string(),
name: "read_file".to_string(),
input: json!({ "path": self.input_path }).to_string(),
thought_signature: None,
},
AssistantEvent::MessageStop,
])