claude: Include all mentions and images in user message (#36539)

Agus Zubiaga created

User messages sent to Claude Code will now include the content of all
mentions, and any images included.

Release Notes:

- N/A

Change summary

crates/agent_servers/src/claude.rs | 242 ++++++++++++++++++++++++++++---
1 file changed, 218 insertions(+), 24 deletions(-)

Detailed changes

crates/agent_servers/src/claude.rs 🔗

@@ -32,7 +32,7 @@ use util::{ResultExt, debug_panic};
 use crate::claude::mcp_server::{ClaudeZedMcpServer, McpConfig};
 use crate::claude::tools::ClaudeTool;
 use crate::{AgentServer, AgentServerCommand, AllAgentServersSettings};
-use acp_thread::{AcpThread, AgentConnection, AuthRequired};
+use acp_thread::{AcpThread, AgentConnection, AuthRequired, MentionUri};
 
 #[derive(Clone)]
 pub struct ClaudeCode;
@@ -267,27 +267,12 @@ impl AgentConnection for ClaudeAgentConnection {
         let (end_tx, end_rx) = oneshot::channel();
         session.turn_state.replace(TurnState::InProgress { end_tx });
 
-        let mut content = String::new();
-        for chunk in params.prompt {
-            match chunk {
-                acp::ContentBlock::Text(text_content) => {
-                    content.push_str(&text_content.text);
-                }
-                acp::ContentBlock::ResourceLink(resource_link) => {
-                    content.push_str(&format!("@{}", resource_link.uri));
-                }
-                acp::ContentBlock::Audio(_)
-                | acp::ContentBlock::Image(_)
-                | acp::ContentBlock::Resource(_) => {
-                    // TODO
-                }
-            }
-        }
+        let content = acp_content_to_claude(params.prompt);
 
         if let Err(err) = session.outgoing_tx.unbounded_send(SdkMessage::User {
             message: Message {
                 role: Role::User,
-                content: Content::UntaggedText(content),
+                content: Content::Chunks(content),
                 id: None,
                 model: None,
                 stop_reason: None,
@@ -513,10 +498,17 @@ impl ClaudeAgentSession {
                                 chunk
                             );
                         }
+                        ContentChunk::Image { source } => {
+                            if !turn_state.borrow().is_canceled() {
+                                thread
+                                    .update(cx, |thread, cx| {
+                                        thread.push_user_content_block(None, source.into(), cx)
+                                    })
+                                    .log_err();
+                            }
+                        }
 
-                        ContentChunk::Image
-                        | ContentChunk::Document
-                        | ContentChunk::WebSearchToolResult => {
+                        ContentChunk::Document | ContentChunk::WebSearchToolResult => {
                             thread
                                 .update(cx, |thread, cx| {
                                     thread.push_assistant_content_block(
@@ -602,7 +594,14 @@ impl ClaudeAgentSession {
                                 "Should not get tool results with role: assistant. should we handle this?"
                             );
                         }
-                        ContentChunk::Image | ContentChunk::Document => {
+                        ContentChunk::Image { source } => {
+                            thread
+                                .update(cx, |thread, cx| {
+                                    thread.push_assistant_content_block(source.into(), false, cx)
+                                })
+                                .log_err();
+                        }
+                        ContentChunk::Document => {
                             thread
                                 .update(cx, |thread, cx| {
                                     thread.push_assistant_content_block(
@@ -768,14 +767,44 @@ enum ContentChunk {
         thinking: String,
     },
     RedactedThinking,
+    Image {
+        source: ImageSource,
+    },
     // TODO
-    Image,
     Document,
     WebSearchToolResult,
     #[serde(untagged)]
     UntaggedText(String),
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+enum ImageSource {
+    Base64 { data: String, media_type: String },
+    Url { url: String },
+}
+
+impl Into<acp::ContentBlock> for ImageSource {
+    fn into(self) -> acp::ContentBlock {
+        match self {
+            ImageSource::Base64 { data, media_type } => {
+                acp::ContentBlock::Image(acp::ImageContent {
+                    annotations: None,
+                    data,
+                    mime_type: media_type,
+                    uri: None,
+                })
+            }
+            ImageSource::Url { url } => acp::ContentBlock::Image(acp::ImageContent {
+                annotations: None,
+                data: "".to_string(),
+                mime_type: "".to_string(),
+                uri: Some(url),
+            }),
+        }
+    }
+}
+
 impl Display for ContentChunk {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
@@ -784,7 +813,7 @@ impl Display for ContentChunk {
             ContentChunk::RedactedThinking => write!(f, "Thinking: [REDACTED]"),
             ContentChunk::UntaggedText(text) => write!(f, "{}", text),
             ContentChunk::ToolResult { content, .. } => write!(f, "{}", content),
-            ContentChunk::Image
+            ContentChunk::Image { .. }
             | ContentChunk::Document
             | ContentChunk::ToolUse { .. }
             | ContentChunk::WebSearchToolResult => {
@@ -896,6 +925,75 @@ impl Display for ResultErrorType {
     }
 }
 
+fn acp_content_to_claude(prompt: Vec<acp::ContentBlock>) -> Vec<ContentChunk> {
+    let mut content = Vec::with_capacity(prompt.len());
+    let mut context = Vec::with_capacity(prompt.len());
+
+    for chunk in prompt {
+        match chunk {
+            acp::ContentBlock::Text(text_content) => {
+                content.push(ContentChunk::Text {
+                    text: text_content.text,
+                });
+            }
+            acp::ContentBlock::ResourceLink(resource_link) => {
+                match MentionUri::parse(&resource_link.uri) {
+                    Ok(uri) => {
+                        content.push(ContentChunk::Text {
+                            text: format!("{}", uri.as_link()),
+                        });
+                    }
+                    Err(_) => {
+                        content.push(ContentChunk::Text {
+                            text: resource_link.uri,
+                        });
+                    }
+                }
+            }
+            acp::ContentBlock::Resource(resource) => match resource.resource {
+                acp::EmbeddedResourceResource::TextResourceContents(resource) => {
+                    match MentionUri::parse(&resource.uri) {
+                        Ok(uri) => {
+                            content.push(ContentChunk::Text {
+                                text: format!("{}", uri.as_link()),
+                            });
+                        }
+                        Err(_) => {
+                            content.push(ContentChunk::Text {
+                                text: resource.uri.clone(),
+                            });
+                        }
+                    }
+
+                    context.push(ContentChunk::Text {
+                        text: format!(
+                            "\n<context ref=\"{}\">\n{}\n</context>",
+                            resource.uri, resource.text
+                        ),
+                    });
+                }
+                acp::EmbeddedResourceResource::BlobResourceContents(_) => {
+                    // Unsupported by SDK
+                }
+            },
+            acp::ContentBlock::Image(acp::ImageContent {
+                data, mime_type, ..
+            }) => content.push(ContentChunk::Image {
+                source: ImageSource::Base64 {
+                    data,
+                    media_type: mime_type,
+                },
+            }),
+            acp::ContentBlock::Audio(_) => {
+                // Unsupported by SDK
+            }
+        }
+    }
+
+    content.extend(context);
+    content
+}
+
 fn new_request_id() -> String {
     use rand::Rng;
     // In the Claude Code TS SDK they just generate a random 12 character string,
@@ -1112,4 +1210,100 @@ pub(crate) mod tests {
             _ => panic!("Expected ToolResult variant"),
         }
     }
+
+    #[test]
+    fn test_acp_content_to_claude() {
+        let acp_content = vec![
+            acp::ContentBlock::Text(acp::TextContent {
+                text: "Hello world".to_string(),
+                annotations: None,
+            }),
+            acp::ContentBlock::Image(acp::ImageContent {
+                data: "base64data".to_string(),
+                mime_type: "image/png".to_string(),
+                annotations: None,
+                uri: None,
+            }),
+            acp::ContentBlock::ResourceLink(acp::ResourceLink {
+                uri: "file:///path/to/example.rs".to_string(),
+                name: "example.rs".to_string(),
+                annotations: None,
+                description: None,
+                mime_type: None,
+                size: None,
+                title: None,
+            }),
+            acp::ContentBlock::Resource(acp::EmbeddedResource {
+                annotations: None,
+                resource: acp::EmbeddedResourceResource::TextResourceContents(
+                    acp::TextResourceContents {
+                        mime_type: None,
+                        text: "fn main() { println!(\"Hello!\"); }".to_string(),
+                        uri: "file:///path/to/code.rs".to_string(),
+                    },
+                ),
+            }),
+            acp::ContentBlock::ResourceLink(acp::ResourceLink {
+                uri: "invalid_uri_format".to_string(),
+                name: "invalid.txt".to_string(),
+                annotations: None,
+                description: None,
+                mime_type: None,
+                size: None,
+                title: None,
+            }),
+        ];
+
+        let claude_content = acp_content_to_claude(acp_content);
+
+        assert_eq!(claude_content.len(), 6);
+
+        match &claude_content[0] {
+            ContentChunk::Text { text } => assert_eq!(text, "Hello world"),
+            _ => panic!("Expected Text chunk"),
+        }
+
+        match &claude_content[1] {
+            ContentChunk::Image { source } => match source {
+                ImageSource::Base64 { data, media_type } => {
+                    assert_eq!(data, "base64data");
+                    assert_eq!(media_type, "image/png");
+                }
+                _ => panic!("Expected Base64 image source"),
+            },
+            _ => panic!("Expected Image chunk"),
+        }
+
+        match &claude_content[2] {
+            ContentChunk::Text { text } => {
+                assert!(text.contains("example.rs"));
+                assert!(text.contains("file:///path/to/example.rs"));
+            }
+            _ => panic!("Expected Text chunk for ResourceLink"),
+        }
+
+        match &claude_content[3] {
+            ContentChunk::Text { text } => {
+                assert!(text.contains("code.rs"));
+                assert!(text.contains("file:///path/to/code.rs"));
+            }
+            _ => panic!("Expected Text chunk for Resource"),
+        }
+
+        match &claude_content[4] {
+            ContentChunk::Text { text } => {
+                assert_eq!(text, "invalid_uri_format");
+            }
+            _ => panic!("Expected Text chunk for invalid URI"),
+        }
+
+        match &claude_content[5] {
+            ContentChunk::Text { text } => {
+                assert!(text.contains("<context ref=\"file:///path/to/code.rs\">"));
+                assert!(text.contains("fn main() { println!(\"Hello!\"); }"));
+                assert!(text.contains("</context>"));
+            }
+            _ => panic!("Expected Text chunk for context"),
+        }
+    }
 }