fix(databricks): add tool_call_id attribution to deferred image messages

crates/goose/src/providers/formats/databricks.rs

@@ -1,18 +1,18 @@
 use crate::conversation::message::{Message, MessageContent};
 use crate::model::ModelConfig;
-use crate::providers::formats::anthropic::{thinking_effort, thinking_type, ThinkingType};
+use crate::providers::formats::anthropic::{ThinkingType, thinking_effort, thinking_type};
 use crate::providers::utils::{
-    convert_image, detect_image_path, extract_reasoning_effort, is_openai_responses_model,
-    is_valid_function_name, load_image_file, safely_parse_json, sanitize_function_name,
-    ImageFormat,
+    ImageFormat, convert_image, detect_image_path, extract_reasoning_effort,
+    is_openai_responses_model, is_valid_function_name, load_image_file, safely_parse_json,
+    sanitize_function_name,
 };
-use anyhow::{anyhow, Error};
+use anyhow::{Error, anyhow};
 use rmcp::model::{
-    object, AnnotateAble, CallToolRequestParams, Content, ErrorCode, ErrorData, RawContent,
-    ResourceContents, Role, Tool,
+    AnnotateAble, CallToolRequestParams, Content, ErrorCode, ErrorData, RawContent,
+    ResourceContents, Role, Tool, object,
 };
 use serde::Serialize;
-use serde_json::{json, Value};
+use serde_json::{Value, json};
 use std::borrow::Cow;
 
 #[derive(Serialize)]
@@ -55,11 +55,14 @@ fn format_tool_response(
                 match content {
                     RawContent::Image(image) => {
                         tool_content.push(Content::text(
-                            "This tool result included an image that is uploaded in the next message.",
+                            &format!("This tool result included an image, provided in a following user message (tool_call_id: {}).", response.id),
                         ));
                         image_messages.push(DatabricksMessage {
                             role: "user".to_string(),
-                            content: [convert_image(&image.no_annotation(), image_format)].into(),
+                            content: json!([
+                                {"type": "text", "text": format!("Image result for tool call {}:", response.id)},
+                                convert_image(&image.no_annotation(), image_format)
+                            ]),
                             tool_calls: None,
                             tool_call_id: None,
                         });
@@ -75,11 +78,13 @@ fn format_tool_response(
                 }
             }
 
-            let tool_response_content: Value = json!(tool_content
-                .iter()
-                .filter_map(|c| c.as_text().map(|t| t.text.clone()))
-                .collect::<Vec<String>>()
-                .join(" "));
+            let tool_response_content: Value = json!(
+                tool_content
+                    .iter()
+                    .filter_map(|c| c.as_text().map(|t| t.text.clone()))
+                    .collect::<Vec<String>>()
+                    .join(" ")
+            );
 
             result.push(DatabricksMessage {
                 content: tool_response_content,
@@ -793,10 +798,14 @@ mod tests {
 
     #[test]
     fn test_format_messages_multiple_content() -> anyhow::Result<()> {
-        let mut messages = vec![Message::assistant().with_tool_request(
-            "tool1",
-            Ok(CallToolRequestParams::new("example").with_arguments(object!({"param1": "value1"}))),
-        )];
+        let mut messages =
+            vec![
+                Message::assistant().with_tool_request(
+                    "tool1",
+                    Ok(CallToolRequestParams::new("example")
+                        .with_arguments(object!({"param1": "value1"}))),
+                ),
+            ];
 
         let tool_id = if let MessageContent::ToolRequest(request) = &messages[0].content[0] {
             &request.id
@@ -857,10 +866,12 @@ mod tests {
 
         let result = format_tools(&[tool1, tool2], "gpt-4o");
         assert!(result.is_err());
-        assert!(result
-            .unwrap_err()
-            .to_string()
-            .contains("Duplicate tool name"));
+        assert!(
+            result
+                .unwrap_err()
+                .to_string()
+                .contains("Duplicate tool name")
+        );
 
         Ok(())
     }
@@ -892,10 +903,12 @@ mod tests {
         assert_eq!(content[0]["type"], "text");
         assert!(content[0]["text"].as_str().unwrap().contains(png_path_str));
         assert_eq!(content[1]["type"], "image_url");
-        assert!(content[1]["image_url"]["url"]
-            .as_str()
-            .unwrap()
-            .starts_with("data:image/png;base64,"));
+        assert!(
+            content[1]["image_url"]["url"]
+                .as_str()
+                .unwrap()
+                .starts_with("data:image/png;base64,")
+        );
 
         Ok(())
     }
@@ -1692,4 +1705,51 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_image_tool_response_carries_attribution() -> anyhow::Result<()> {
+        // When a tool response contains an image, both the placeholder in the tool result
+        // and the deferred user message must reference the tool_call_id so the model can
+        // correlate the image back to the correct tool even when multiple tools ran in parallel.
+        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::text("text result")])),
+                ),
+        ];
+
+        let as_value =
+            serde_json::to_value(format_messages(&messages, &ImageFormat::OpenAi)).unwrap();
+        let spec = as_value.as_array().unwrap();
+
+        // tool result for id1 must reference its own tool_call_id in the placeholder
+        let tool1_content = spec[1]["content"].as_str().unwrap();
+        assert!(
+            tool1_content.contains("id1"),
+            "placeholder should reference tool_call_id id1, got: {tool1_content}"
+        );
+
+        // deferred user image message must identify which tool it belongs to
+        let image_msg = &spec[3];
+        assert_eq!(image_msg["role"], "user");
+        let image_content = image_msg["content"].as_array().unwrap();
+        let label = image_content[0]["text"].as_str().unwrap();
+        assert!(
+            label.contains("id1"),
+            "image message label should reference id1, got: {label}"
+        );
+
+        Ok(())
+    }
 }