fix(databricks): ensure parallel tool image responses don't interleave tool results (#9241)

Signed-off-by: Steve Marshall <steve.marshall@fasthosts.com>
Signed-off-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: Douwe Osinga <douwe@squareup.com>
This commit is contained in:
Steve Marshall
2026-05-20 12:44:10 +01:00
committed by GitHub
parent 8f9270e295
commit 955288db36
@@ -124,6 +124,8 @@ fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<Data
let mut content_array = Vec::new();
let mut has_tool_calls = false;
let mut has_multiple_content = false;
// Deferred so all tool-role messages stay consecutive (required by Claude via Databricks).
let mut pending_image_messages: Vec<DatabricksMessage> = Vec::new();
for content in &message.content {
match content {
@@ -190,7 +192,13 @@ fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<Data
}
}
MessageContent::ToolResponse(response) => {
result.extend(format_tool_response(response, image_format));
for msg in format_tool_response(response, image_format) {
if msg.role == "user" {
pending_image_messages.push(msg);
} else {
result.push(msg);
}
}
}
MessageContent::Image(image) => {
content_array.push(convert_image(image, image_format));
@@ -212,6 +220,8 @@ fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<Data
}
}
result.extend(pending_image_messages);
if !content_array.is_empty() {
converted.content = if content_array.len() == 1
&& !has_multiple_content
@@ -1683,4 +1693,70 @@ mod tests {
Ok(())
}
#[test]
fn test_parallel_tool_responses_with_images_are_consecutive() -> anyhow::Result<()> {
// Regression: #7449 — parallel tool calls with images must keep tool messages consecutive.
let messages = vec![
Message::assistant()
.with_tool_request("id1", Ok(CallToolRequestParams::new("tool_a")))
.with_tool_request("id2", Ok(CallToolRequestParams::new("tool_b"))),
Message::user()
.with_tool_response(
"id1",
Ok(CallToolResult::success(vec![Content::image(
"base64data1".to_string(),
"image/png".to_string(),
)])),
)
.with_tool_response(
"id2",
Ok(CallToolResult::success(vec![Content::image(
"base64data2".to_string(),
"image/png".to_string(),
)])),
),
];
let as_value =
serde_json::to_value(format_messages(&messages, &ImageFormat::OpenAi)).unwrap();
let spec = as_value.as_array().unwrap();
let roles: Vec<&str> = spec.iter().map(|m| m["role"].as_str().unwrap()).collect();
// Without the fix this was ["assistant", "tool", "user", "tool", "user"].
assert_eq!(roles, vec!["assistant", "tool", "tool", "user", "user"]);
Ok(())
}
#[test]
fn test_mixed_tool_responses_image_and_text_ordering() -> anyhow::Result<()> {
// Mixed case: only one tool response has an image.
let messages = vec![
Message::assistant()
.with_tool_request("id1", Ok(CallToolRequestParams::new("tool_a")))
.with_tool_request("id2", Ok(CallToolRequestParams::new("tool_b"))),
Message::user()
.with_tool_response(
"id1",
Ok(CallToolResult::success(vec![Content::text("text result")])),
)
.with_tool_response(
"id2",
Ok(CallToolResult::success(vec![Content::image(
"base64data".to_string(),
"image/png".to_string(),
)])),
),
];
let as_value =
serde_json::to_value(format_messages(&messages, &ImageFormat::OpenAi)).unwrap();
let spec = as_value.as_array().unwrap();
let roles: Vec<&str> = spec.iter().map(|m| m["role"].as_str().unwrap()).collect();
assert_eq!(roles, vec!["assistant", "tool", "tool", "user"]);
Ok(())
}
}