1use super::PromptAttachment;
2use agent_client_protocol::schema as acp;
3use base64::Engine;
4use base64::engine::general_purpose::STANDARD as BASE64;
5use std::path::Path;
6use tokio::io::AsyncReadExt;
7use url::Url;
8
9const MAX_EMBED_TEXT_BYTES: usize = 1024 * 1024;
10const MAX_MEDIA_BYTES: usize = 10 * 1024 * 1024;
11
12const IMAGE_MIME_TYPES: &[&str] = &["image/png", "image/jpeg", "image/gif", "image/webp"];
13const AUDIO_MIME_TYPES: &[&str] = &["audio/wav", "audio/mpeg", "audio/mp3", "audio/ogg"];
14
15#[derive(Debug, thiserror::Error)]
16pub enum AttachmentError {
17 #[error("Failed to read {name}: {source}")]
18 Read {
19 name: String,
20 #[source]
21 source: std::io::Error,
22 },
23 #[error("Skipped {name}: file too large ({size} bytes, max {max})")]
24 TooLarge { name: String, size: u64, max: usize },
25 #[error("Skipped binary or non-UTF8 file: {0}")]
26 NotUtf8(String),
27 #[error("Failed to build file URI for {0}")]
28 FileUri(String),
29}
30
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub enum AttachmentKind {
33 Text,
34 Image,
35 Audio,
36 Unsupported,
37}
38
39pub fn classify_attachment(path: &Path) -> AttachmentKind {
40 let mime = mime_guess::from_path(path).first_or_octet_stream().to_string();
41
42 if IMAGE_MIME_TYPES.contains(&mime.as_str()) {
43 AttachmentKind::Image
44 } else if AUDIO_MIME_TYPES.contains(&mime.as_str()) {
45 AttachmentKind::Audio
46 } else if mime.starts_with("text/") {
47 AttachmentKind::Text
48 } else {
49 AttachmentKind::Unsupported
50 }
51}
52
53#[derive(Debug, Default)]
54pub struct AttachmentBuildOutcome {
55 pub blocks: Vec<acp::ContentBlock>,
56 pub transcript_placeholders: Vec<String>,
57 pub warnings: Vec<String>,
58}
59
60pub async fn build_attachment_blocks(attachments: &[PromptAttachment]) -> AttachmentBuildOutcome {
61 let mut outcome = AttachmentBuildOutcome::default();
62
63 for attachment in attachments {
64 match try_build_attachment_block(&attachment.path, &attachment.display_name).await {
65 Ok(result) => {
66 outcome.blocks.push(result.block);
67 if let Some(placeholder) = result.transcript_placeholder {
68 outcome.transcript_placeholders.push(placeholder);
69 }
70 if let Some(warning) = result.warning {
71 outcome.warnings.push(warning);
72 }
73 }
74 Err(warning) => outcome.warnings.push(warning.to_string()),
75 }
76 }
77
78 outcome
79}
80
81struct AttachmentBlockResult {
82 block: acp::ContentBlock,
83 transcript_placeholder: Option<String>,
84 warning: Option<String>,
85}
86
87async fn try_build_attachment_block(path: &Path, display_name: &str) -> Result<AttachmentBlockResult, AttachmentError> {
88 let kind = classify_attachment(path);
89 let mime_type = mime_guess::from_path(path).first_or_octet_stream().to_string();
90
91 match kind {
92 AttachmentKind::Image | AttachmentKind::Audio => {
93 let bytes = read_media_bytes(path, display_name).await?;
94 let data = BASE64.encode(&bytes);
95 let (block, placeholder) = match kind {
96 AttachmentKind::Image => (
97 acp::ContentBlock::Image(acp::ImageContent::new(data, &mime_type)),
98 format!("[image attachment: {display_name}]"),
99 ),
100 _ => (
101 acp::ContentBlock::Audio(acp::AudioContent::new(data, &mime_type)),
102 format!("[audio attachment: {display_name}]"),
103 ),
104 };
105 Ok(AttachmentBlockResult { block, transcript_placeholder: Some(placeholder), warning: None })
106 }
107 _ => build_text_resource_block(path, display_name, &mime_type).await,
108 }
109}
110
111async fn read_media_bytes(path: &Path, display_name: &str) -> Result<Vec<u8>, AttachmentError> {
112 let metadata = tokio::fs::metadata(path)
113 .await
114 .map_err(|e| AttachmentError::Read { name: display_name.to_string(), source: e })?;
115
116 if metadata.len() > MAX_MEDIA_BYTES as u64 {
117 return Err(AttachmentError::TooLarge {
118 name: display_name.to_string(),
119 size: metadata.len(),
120 max: MAX_MEDIA_BYTES,
121 });
122 }
123
124 tokio::fs::read(path).await.map_err(|e| AttachmentError::Read { name: display_name.to_string(), source: e })
125}
126
127async fn build_text_resource_block(
128 path: &Path,
129 display_name: &str,
130 mime_type: &str,
131) -> Result<AttachmentBlockResult, AttachmentError> {
132 let file = tokio::fs::File::open(path)
133 .await
134 .map_err(|e| AttachmentError::Read { name: display_name.to_string(), source: e })?;
135
136 let mut bytes = Vec::new();
137 file.take((MAX_EMBED_TEXT_BYTES + 1) as u64)
138 .read_to_end(&mut bytes)
139 .await
140 .map_err(|e| AttachmentError::Read { name: display_name.to_string(), source: e })?;
141
142 let truncated = bytes.len() > MAX_EMBED_TEXT_BYTES;
143 if truncated {
144 bytes.truncate(MAX_EMBED_TEXT_BYTES);
145 }
146 let text_bytes = bytes.as_slice();
147
148 let text = match std::str::from_utf8(text_bytes) {
149 Ok(text) => text.to_string(),
150 Err(error) if truncated && error.valid_up_to() > 0 => {
151 let valid_bytes = &text_bytes[..error.valid_up_to()];
152 std::str::from_utf8(valid_bytes).expect("valid_up_to must point at a utf8 boundary").to_string()
153 }
154 Err(_) => return Err(AttachmentError::NotUtf8(display_name.to_string())),
155 };
156
157 let file_uri = build_attachment_file_uri(path, display_name).await?;
158 let warning = truncated.then(|| format!("Truncated {display_name} to {MAX_EMBED_TEXT_BYTES} bytes"));
159
160 let block =
161 acp::ContentBlock::Resource(acp::EmbeddedResource::new(acp::EmbeddedResourceResource::TextResourceContents(
162 acp::TextResourceContents::new(text, file_uri).mime_type(mime_type),
163 )));
164
165 Ok(AttachmentBlockResult { block, transcript_placeholder: None, warning })
166}
167
168async fn build_attachment_file_uri(path: &Path, display_name: &str) -> Result<String, AttachmentError> {
169 let canonical_path = tokio::fs::canonicalize(path).await.ok();
170 let uri_path = canonical_path.as_deref().unwrap_or(path);
171 Url::from_file_path(uri_path)
172 .map_err(|()| AttachmentError::FileUri(display_name.to_string()))
173 .map(|url| url.to_string())
174}
175
176#[cfg(test)]
177mod tests {
178 use super::*;
179 use tempfile::TempDir;
180
181 #[tokio::test]
182 async fn build_attachment_blocks_truncates_large_file_with_warning() {
183 let tmp = TempDir::new().unwrap();
184 let path = tmp.path().join("large.txt");
185 let display_name = "large.txt".to_string();
186 std::fs::write(&path, "x".repeat(MAX_EMBED_TEXT_BYTES + 64)).unwrap();
187
188 let attachments = vec![PromptAttachment { path, display_name: display_name.clone() }];
189 let outcome = build_attachment_blocks(&attachments).await;
190
191 assert_eq!(outcome.blocks.len(), 1);
192 assert_eq!(outcome.warnings.len(), 1);
193 assert!(outcome.warnings[0].contains(&format!("Truncated {display_name} to {MAX_EMBED_TEXT_BYTES} bytes")));
194 }
195
196 #[tokio::test]
197 async fn build_attachment_blocks_skips_non_utf8_files() {
198 let tmp = TempDir::new().unwrap();
199 let path = tmp.path().join("binary.bin");
200 let display_name = "binary.bin".to_string();
201 std::fs::write(&path, [0xff, 0xfe, 0xfd]).unwrap();
202
203 let attachments = vec![PromptAttachment { path, display_name: display_name.clone() }];
204 let outcome = build_attachment_blocks(&attachments).await;
205
206 assert!(outcome.blocks.is_empty());
207 assert_eq!(outcome.warnings.len(), 1);
208 assert!(outcome.warnings[0].contains(&format!("Skipped binary or non-UTF8 file: {display_name}")));
209 }
210
211 #[tokio::test]
212 async fn build_attachment_file_uri_falls_back_when_canonicalize_fails() {
213 let tmp = TempDir::new().unwrap();
214 let path = tmp.path().join("missing.txt");
215 let expected = Url::from_file_path(&path).unwrap().to_string();
216
217 let uri = build_attachment_file_uri(&path, "missing.txt")
218 .await
219 .expect("URI should be built from original absolute path");
220
221 assert_eq!(uri, expected);
222 }
223
224 #[tokio::test]
225 async fn png_file_produces_image_content_block() {
226 let tmp = TempDir::new().unwrap();
227 let path = tmp.path().join("test.png");
228 std::fs::write(&path, b"fake png data").unwrap();
229
230 let attachments = vec![PromptAttachment { path, display_name: "test.png".to_string() }];
231 let outcome = build_attachment_blocks(&attachments).await;
232
233 assert_eq!(outcome.blocks.len(), 1);
234 assert!(outcome.warnings.is_empty());
235 assert_eq!(outcome.transcript_placeholders, vec!["[image attachment: test.png]"]);
236 assert!(matches!(outcome.blocks[0], acp::ContentBlock::Image(_)));
237 }
238
239 #[tokio::test]
240 async fn wav_file_produces_audio_content_block() {
241 let tmp = TempDir::new().unwrap();
242 let path = tmp.path().join("test.wav");
243 std::fs::write(&path, b"fake wav data").unwrap();
244
245 let attachments = vec![PromptAttachment { path, display_name: "test.wav".to_string() }];
246 let outcome = build_attachment_blocks(&attachments).await;
247
248 assert_eq!(outcome.blocks.len(), 1);
249 assert!(outcome.warnings.is_empty());
250 assert_eq!(outcome.transcript_placeholders, vec!["[audio attachment: test.wav]"]);
251 assert!(matches!(outcome.blocks[0], acp::ContentBlock::Audio(_)));
252 }
253
254 #[test]
255 fn classify_attachment_detects_images() {
256 assert_eq!(classify_attachment(Path::new("photo.png")), AttachmentKind::Image);
257 assert_eq!(classify_attachment(Path::new("photo.jpg")), AttachmentKind::Image);
258 assert_eq!(classify_attachment(Path::new("photo.gif")), AttachmentKind::Image);
259 assert_eq!(classify_attachment(Path::new("photo.webp")), AttachmentKind::Image);
260 }
261
262 #[test]
263 fn classify_attachment_detects_audio() {
264 assert_eq!(classify_attachment(Path::new("note.wav")), AttachmentKind::Audio);
265 assert_eq!(classify_attachment(Path::new("note.mp3")), AttachmentKind::Audio);
266 assert_eq!(classify_attachment(Path::new("note.ogg")), AttachmentKind::Audio);
267 }
268
269 #[test]
270 fn classify_attachment_detects_text() {
271 assert_eq!(classify_attachment(Path::new("readme.txt")), AttachmentKind::Text);
272 }
273
274 #[test]
275 fn classify_attachment_unknown_extension_is_unsupported() {
276 assert_eq!(classify_attachment(Path::new("data.xyz")), AttachmentKind::Unsupported);
277 }
278}