| 1 | //! Converting ACP prompt content into model content: text, embedded and linked files |
| 2 | //! (inlined, a selection as just its lines), images, and audio saved to a scratch file. |
| 3 | |
| 4 | use std::path::{Path, PathBuf}; |
| 5 | |
| 6 | use agent_client_protocol::schema::v1::{ |
| 7 | ClientCapabilities, ContentBlock, EmbeddedResource, EmbeddedResourceResource, |
| 8 | ReadTextFileRequest, ResourceLink, SessionId, TextResourceContents, |
| 9 | }; |
| 10 | use agent_client_protocol::{Client, ConnectionTo}; |
| 11 | use anyhow::{Context, Result}; |
| 12 | use base64::Engine; |
| 13 | use serde_json::{Value, json}; |
| 14 | |
| 15 | fn embedded_text(res: &EmbeddedResource) -> Option<String> { |
| 16 | match &res.resource { |
| 17 | EmbeddedResourceResource::TextResourceContents(r) => { |
| 18 | Some(format!("<file uri=\"{}\">\n{}\n</file>", r.uri, r.text)) |
| 19 | } |
| 20 | // Describe binary resources rather than dropping them, so the model knows they exist. |
| 21 | EmbeddedResourceResource::BlobResourceContents(b) => Some(format!( |
| 22 | "[Attached binary resource: {} ({})]", |
| 23 | b.uri, |
| 24 | b.mime_type.as_deref().unwrap_or("unknown type") |
| 25 | )), |
| 26 | _ => None, |
| 27 | } |
| 28 | } |
| 29 | |
| 30 | /// How audio reaches the model. A product with an audio tool points the model at it here. |
| 31 | #[derive(Clone, Copy)] |
| 32 | pub struct AudioHints { |
| 33 | /// The text for a `file://` link to an audio file, given its path. |
| 34 | pub link: fn(&Path) -> String, |
| 35 | /// The text for an attached clip, given its MIME type and the path it was saved to. |
| 36 | pub attachment: fn(&str, &Path) -> String, |
| 37 | } |
| 38 | |
| 39 | impl Default for AudioHints { |
| 40 | fn default() -> Self { |
| 41 | Self { |
| 42 | link: |path| format!("[Referenced audio file: {}]", path.display()), |
| 43 | attachment: |mime, path| { |
| 44 | format!("[Attached audio ({mime}) saved at {}]", path.display()) |
| 45 | }, |
| 46 | } |
| 47 | } |
| 48 | } |
| 49 | |
| 50 | /// How a resource link is shown to the model. Audio files are never read as text; they get |
| 51 | /// the product's [`AudioHints::link`] text instead. |
| 52 | fn link_text(link: &ResourceLink, hints: &AudioHints) -> String { |
| 53 | match parse_file_link(&link.uri) { |
| 54 | Some(FileLink { path, .. }) if is_audio(&path, link.mime_type.as_deref()) => { |
| 55 | (hints.link)(&path) |
| 56 | } |
| 57 | _ => format!("[Referenced: {}]", link.uri), |
| 58 | } |
| 59 | } |
| 60 | |
| 61 | /// Common audio file extensions; a link to one is never read as text. |
| 62 | const AUDIO_EXTENSIONS: &[&str] = &[ |
| 63 | "wav", "wave", "mp3", "flac", "m4a", "aac", "mp4", "ogg", "oga", "opus", "aif", "aiff", "aifc", |
| 64 | "caf", "webm", "mkv", |
| 65 | ]; |
| 66 | |
| 67 | fn is_audio(path: &Path, mime: Option<&str>) -> bool { |
| 68 | mime.is_some_and(|m| m.starts_with("audio/")) |
| 69 | || path |
| 70 | .extension() |
| 71 | .and_then(|e| e.to_str()) |
| 72 | .is_some_and(|e| AUDIO_EXTENSIONS.contains(&e.to_ascii_lowercase().as_str())) |
| 73 | } |
| 74 | |
| 75 | /// Largest file a `file://` resource link is inlined for; bigger ones stay as references. |
| 76 | const MAX_LINKED_FILE_BYTES: usize = 256 * 1024; |
| 77 | |
| 78 | /// Replace `file://` resource links to text files with their embedded contents, so the model |
| 79 | /// sees the file. Reads through the client when it supports `fs/read_text_file` (unsaved |
| 80 | /// buffers), else from disk. A link to a selection (`#L10:20`) is inlined as just those |
| 81 | /// lines. Audio files, links that can't be read and oversized files stay as references. |
| 82 | pub async fn resolve_resource_links( |
| 83 | blocks: &[ContentBlock], |
| 84 | caps: &ClientCapabilities, |
| 85 | connection: &ConnectionTo<Client>, |
| 86 | session_id: &SessionId, |
| 87 | ) -> Vec<ContentBlock> { |
| 88 | let mut out = Vec::with_capacity(blocks.len()); |
| 89 | for block in blocks { |
| 90 | if let ContentBlock::ResourceLink(link) = block |
| 91 | && let Some(FileLink { path, lines }) = parse_file_link(&link.uri) |
| 92 | && !is_audio(&path, link.mime_type.as_deref()) |
| 93 | { |
| 94 | let text = if caps.fs.read_text_file { |
| 95 | connection |
| 96 | .send_request(ReadTextFileRequest::new(session_id.clone(), path.clone())) |
| 97 | .block_task() |
| 98 | .await |
| 99 | .map(|r| r.content) |
| 100 | .map_err(|e| e.to_string()) |
| 101 | } else { |
| 102 | tokio::fs::read_to_string(&path) |
| 103 | .await |
| 104 | .map_err(|e| e.to_string()) |
| 105 | }; |
| 106 | let text = match lines { |
| 107 | Some(range) => text.map(|t| slice_lines(&t, range)), |
| 108 | None => text, |
| 109 | }; |
| 110 | match text { |
| 111 | Ok(text) if lines.is_some() && text.is_empty() => { |
| 112 | tracing::debug!("{} selects no lines", link.uri) |
| 113 | } |
| 114 | Ok(text) if text.len() <= MAX_LINKED_FILE_BYTES => { |
| 115 | out.push(ContentBlock::Resource(EmbeddedResource::new( |
| 116 | EmbeddedResourceResource::TextResourceContents( |
| 117 | TextResourceContents::new(text, link.uri.clone()) |
| 118 | .mime_type(link.mime_type.clone()), |
| 119 | ), |
| 120 | ))); |
| 121 | continue; |
| 122 | } |
| 123 | Ok(_) => tracing::debug!("{} too large to inline", link.uri), |
| 124 | Err(e) => tracing::debug!("could not read {}: {e}", link.uri), |
| 125 | } |
| 126 | } |
| 127 | out.push(block.clone()); |
| 128 | } |
| 129 | out |
| 130 | } |
| 131 | |
| 132 | /// A `file://` resource link: the absolute path, and the lines it selects, if any. |
| 133 | #[derive(Debug, PartialEq)] |
| 134 | struct FileLink { |
| 135 | path: PathBuf, |
| 136 | /// 1-based, inclusive. |
| 137 | lines: Option<(usize, usize)>, |
| 138 | } |
| 139 | |
| 140 | /// Parse a `file://` URI, percent-decoding the path. Editors put more than the path in it: |
| 141 | /// Zed links a selection as `file:///a.md?column=5#L10:20` and a symbol as |
| 142 | /// `file:///a.rs?symbol=main#L3:9`. The query is dropped and the fragment read as a line |
| 143 | /// range; a literal `?` or `#` in a file name arrives percent-encoded, so splitting first is |
| 144 | /// safe. |
| 145 | fn parse_file_link(uri: &str) -> Option<FileLink> { |
| 146 | let rest = uri.strip_prefix("file://")?; |
| 147 | let (rest, fragment) = match rest.split_once('#') { |
| 148 | Some((rest, fragment)) => (rest, Some(fragment)), |
| 149 | None => (rest, None), |
| 150 | }; |
| 151 | let rest = rest.split_once('?').map_or(rest, |(rest, _)| rest); |
| 152 | // Allow an authority of "" or "localhost"; reject other hosts. |
| 153 | let path = if rest.starts_with('/') { |
| 154 | rest |
| 155 | } else { |
| 156 | rest.strip_prefix("localhost")? |
| 157 | }; |
| 158 | let bytes = path.as_bytes(); |
| 159 | let mut decoded = Vec::with_capacity(bytes.len()); |
| 160 | let mut i = 0; |
| 161 | while i < bytes.len() { |
| 162 | if bytes[i] == b'%' |
| 163 | && i + 2 < bytes.len() |
| 164 | && let Ok(b) = u8::from_str_radix(std::str::from_utf8(&bytes[i + 1..i + 3]).ok()?, 16) |
| 165 | { |
| 166 | decoded.push(b); |
| 167 | i += 3; |
| 168 | } else { |
| 169 | decoded.push(bytes[i]); |
| 170 | i += 1; |
| 171 | } |
| 172 | } |
| 173 | let path = PathBuf::from(String::from_utf8(decoded).ok()?); |
| 174 | path.is_absolute().then(|| FileLink { |
| 175 | path, |
| 176 | lines: fragment.and_then(parse_line_range), |
| 177 | }) |
| 178 | } |
| 179 | |
| 180 | /// A `#L10:20` fragment as a 1-based inclusive range. Accepts `L10:20`, `L10-20`, `L10-L20` |
| 181 | /// and a single line `L10`. |
| 182 | fn parse_line_range(fragment: &str) -> Option<(usize, usize)> { |
| 183 | let range = fragment.strip_prefix('L')?; |
| 184 | let (start, end) = range |
| 185 | .split_once(':') |
| 186 | .or_else(|| range.split_once('-')) |
| 187 | .unwrap_or((range, range)); |
| 188 | let end = end.strip_prefix('L').unwrap_or(end); |
| 189 | let (start, end) = (start.parse::<usize>().ok()?, end.parse::<usize>().ok()?); |
| 190 | (start >= 1 && end >= start).then_some((start, end)) |
| 191 | } |
| 192 | |
| 193 | /// Lines `start..=end` (1-based) of `text`, keeping their line endings. |
| 194 | fn slice_lines(text: &str, (start, end): (usize, usize)) -> String { |
| 195 | text.split_inclusive('\n') |
| 196 | .skip(start - 1) |
| 197 | .take(end - start + 1) |
| 198 | .collect() |
| 199 | } |
| 200 | |
| 201 | /// Flatten text-like prompt blocks into one string (used for titles and slash commands). |
| 202 | pub fn prompt_to_text(blocks: &[ContentBlock], hints: &AudioHints) -> String { |
| 203 | let mut parts = Vec::new(); |
| 204 | for block in blocks { |
| 205 | match block { |
| 206 | ContentBlock::Text(t) => parts.push(t.text.clone()), |
| 207 | ContentBlock::ResourceLink(link) => parts.push(link_text(link, hints)), |
| 208 | ContentBlock::Resource(res) => parts.extend(embedded_text(res)), |
| 209 | _ => {} |
| 210 | } |
| 211 | } |
| 212 | parts.join("\n\n") |
| 213 | } |
| 214 | |
| 215 | fn audio_extension(mime: &str) -> &'static str { |
| 216 | match mime.split(';').next().unwrap_or("").trim() { |
| 217 | "audio/wav" | "audio/x-wav" | "audio/wave" | "audio/vnd.wave" => "wav", |
| 218 | "audio/mpeg" | "audio/mp3" => "mp3", |
| 219 | "audio/flac" | "audio/x-flac" => "flac", |
| 220 | "audio/ogg" | "audio/vorbis" | "application/ogg" => "ogg", |
| 221 | "audio/mp4" | "audio/m4a" | "audio/x-m4a" | "audio/aac" => "m4a", |
| 222 | "audio/aiff" | "audio/x-aiff" => "aiff", |
| 223 | "audio/webm" | "video/webm" => "webm", |
| 224 | _ => "audio", |
| 225 | } |
| 226 | } |
| 227 | |
| 228 | /// Write each audio block under `dir`; returns the saved path per audio block, in order. |
| 229 | pub async fn save_audio_blocks(blocks: &[ContentBlock], dir: &Path) -> Result<Vec<PathBuf>> { |
| 230 | let mut saved = Vec::new(); |
| 231 | for block in blocks { |
| 232 | let ContentBlock::Audio(audio) = block else { |
| 233 | continue; |
| 234 | }; |
| 235 | let bytes = base64::engine::general_purpose::STANDARD |
| 236 | .decode(audio.data.trim()) |
| 237 | .context("audio block is not valid base64")?; |
| 238 | tokio::fs::create_dir_all(dir).await?; |
| 239 | let path = dir.join(format!( |
| 240 | "audio-{}.{}", |
| 241 | uuid::Uuid::new_v4().simple(), |
| 242 | audio_extension(&audio.mime_type) |
| 243 | )); |
| 244 | tokio::fs::write(&path, bytes).await?; |
| 245 | saved.push(path); |
| 246 | } |
| 247 | Ok(saved) |
| 248 | } |
| 249 | |
| 250 | /// Convert prompt blocks to an OpenAI `content` value: a string for text only, or parts when |
| 251 | /// images are present. Audio is referenced by the path it was saved to, worded by |
| 252 | /// [`AudioHints::attachment`]. `slash` replaces the first text block (an expanded slash |
| 253 | /// command). |
| 254 | pub fn prompt_to_content( |
| 255 | blocks: &[ContentBlock], |
| 256 | audio_paths: &[PathBuf], |
| 257 | slash: Option<&str>, |
| 258 | hints: &AudioHints, |
| 259 | ) -> Value { |
| 260 | let mut audio = audio_paths.iter(); |
| 261 | let mut parts: Vec<Value> = Vec::new(); |
| 262 | let mut slash = slash; |
| 263 | let mut has_image = false; |
| 264 | let text = |t: String| json!({ "type": "text", "text": t }); |
| 265 | for block in blocks { |
| 266 | match block { |
| 267 | ContentBlock::Text(t) => match slash.take() { |
| 268 | Some(expanded) => parts.push(text(expanded.to_string())), |
| 269 | None => parts.push(text(t.text.clone())), |
| 270 | }, |
| 271 | ContentBlock::ResourceLink(link) => parts.push(text(link_text(link, hints))), |
| 272 | ContentBlock::Resource(res) => { |
| 273 | if let Some(t) = embedded_text(res) { |
| 274 | parts.push(text(t)); |
| 275 | } |
| 276 | } |
| 277 | ContentBlock::Image(img) => { |
| 278 | has_image = true; |
| 279 | let url = format!("data:{};base64,{}", img.mime_type, img.data); |
| 280 | parts.push(json!({ "type": "image_url", "image_url": { "url": url } })); |
| 281 | } |
| 282 | ContentBlock::Audio(a) => match audio.next() { |
| 283 | Some(path) => parts.push(text((hints.attachment)(&a.mime_type, path))), |
| 284 | None => parts.push(text("[Attached audio could not be saved]".into())), |
| 285 | }, |
| 286 | _ => {} |
| 287 | } |
| 288 | } |
| 289 | if has_image { |
| 290 | return Value::Array(parts); |
| 291 | } |
| 292 | let joined = parts |
| 293 | .iter() |
| 294 | .filter_map(|p| p.get("text").and_then(Value::as_str)) |
| 295 | .collect::<Vec<_>>() |
| 296 | .join("\n\n"); |
| 297 | Value::String(joined) |
| 298 | } |
| 299 | |
| 300 | #[cfg(test)] |
| 301 | mod tests { |
| 302 | use super::*; |
| 303 | use agent_client_protocol::schema::v1::{AudioContent, ImageContent, TextContent}; |
| 304 | |
| 305 | fn link(uri: &str) -> Option<FileLink> { |
| 306 | parse_file_link(uri) |
| 307 | } |
| 308 | |
| 309 | #[test] |
| 310 | fn plain_paths_decode_and_must_be_absolute() { |
| 311 | assert_eq!( |
| 312 | link("file:///a%20b/c.md"), |
| 313 | Some(FileLink { |
| 314 | path: "/a b/c.md".into(), |
| 315 | lines: None |
| 316 | }) |
| 317 | ); |
| 318 | assert_eq!( |
| 319 | link("file://localhost/x.md").unwrap().path, |
| 320 | Path::new("/x.md") |
| 321 | ); |
| 322 | assert_eq!(link("file://other-host/x.md"), None); |
| 323 | assert_eq!(link("https://example.com/x.md"), None); |
| 324 | } |
| 325 | |
| 326 | #[test] |
| 327 | fn zed_selection_and_symbol_links() { |
| 328 | let sel = link("file:///notes/song.md?column=5#L10:20").unwrap(); |
| 329 | assert_eq!(sel.path, Path::new("/notes/song.md")); |
| 330 | assert_eq!(sel.lines, Some((10, 20))); |
| 331 | assert_eq!( |
| 332 | link("file:///a.rs?symbol=main#L3:9").unwrap().lines, |
| 333 | Some((3, 9)) |
| 334 | ); |
| 335 | } |
| 336 | |
| 337 | #[test] |
| 338 | fn line_range_forms() { |
| 339 | assert_eq!(parse_line_range("L10:20"), Some((10, 20))); |
| 340 | assert_eq!(parse_line_range("L10-20"), Some((10, 20))); |
| 341 | assert_eq!(parse_line_range("L10-L20"), Some((10, 20))); |
| 342 | assert_eq!(parse_line_range("L7"), Some((7, 7))); |
| 343 | assert_eq!(parse_line_range("L0"), None); |
| 344 | assert_eq!(parse_line_range("L5:2"), None); |
| 345 | assert_eq!(parse_line_range("x"), None); |
| 346 | } |
| 347 | |
| 348 | #[test] |
| 349 | fn slices_inclusive_lines() { |
| 350 | assert_eq!(slice_lines("a\nb\nc\nd\n", (2, 3)), "b\nc\n"); |
| 351 | assert_eq!(slice_lines("a\nb", (2, 9)), "b"); |
| 352 | assert_eq!(slice_lines("a\n", (5, 6)), ""); |
| 353 | } |
| 354 | |
| 355 | #[test] |
| 356 | fn audio_links_use_the_product_hint_and_are_never_read() { |
| 357 | let hints = AudioHints { |
| 358 | link: |p| format!("analyze {}", p.display()), |
| 359 | ..AudioHints::default() |
| 360 | }; |
| 361 | let wav = ResourceLink::new("mix.wav", "file:///music/mix%20v2.WAV"); |
| 362 | assert_eq!(link_text(&wav, &hints), "analyze /music/mix v2.WAV"); |
| 363 | let by_mime = |
| 364 | ResourceLink::new("take", "file:///t/take").mime_type("audio/flac".to_string()); |
| 365 | assert_eq!( |
| 366 | link_text(&by_mime, &AudioHints::default()), |
| 367 | "[Referenced audio file: /t/take]" |
| 368 | ); |
| 369 | let md = ResourceLink::new("notes", "file:///t/notes.md"); |
| 370 | assert_eq!(link_text(&md, &hints), "[Referenced: file:///t/notes.md]"); |
| 371 | } |
| 372 | |
| 373 | #[tokio::test] |
| 374 | async fn audio_is_saved_and_referenced() { |
| 375 | let dir = std::env::temp_dir().join(format!("sf-prompt-{}", uuid::Uuid::new_v4())); |
| 376 | let data = base64::engine::general_purpose::STANDARD.encode(b"RIFFfake"); |
| 377 | let blocks = vec![ |
| 378 | ContentBlock::Text(TextContent::new("what key?")), |
| 379 | ContentBlock::Audio(AudioContent::new(data, "audio/wav")), |
| 380 | ]; |
| 381 | let saved = save_audio_blocks(&blocks, &dir).await.unwrap(); |
| 382 | assert_eq!(saved.len(), 1); |
| 383 | assert_eq!(saved[0].extension().unwrap(), "wav"); |
| 384 | assert_eq!(std::fs::read(&saved[0]).unwrap(), b"RIFFfake"); |
| 385 | let content = prompt_to_content(&blocks, &saved, None, &AudioHints::default()); |
| 386 | let s = content.as_str().unwrap(); |
| 387 | assert!(s.contains("what key?") && s.contains(&saved[0].display().to_string())); |
| 388 | std::fs::remove_dir_all(dir).ok(); |
| 389 | |
| 390 | let bad = vec![ContentBlock::Audio(AudioContent::new("!!!", "audio/wav"))]; |
| 391 | assert!( |
| 392 | save_audio_blocks(&bad, &std::env::temp_dir()) |
| 393 | .await |
| 394 | .is_err() |
| 395 | ); |
| 396 | } |
| 397 | |
| 398 | #[test] |
| 399 | fn images_use_parts_and_slash_replaces_text() { |
| 400 | let blocks = vec![ |
| 401 | ContentBlock::Text(TextContent::new("/chords Am F")), |
| 402 | ContentBlock::Image(ImageContent::new("AAAA", "image/png")), |
| 403 | ]; |
| 404 | let v = prompt_to_content(&blocks, &[], Some("EXPANDED"), &AudioHints::default()); |
| 405 | let parts = v.as_array().unwrap(); |
| 406 | assert_eq!(parts[0]["text"], "EXPANDED"); |
| 407 | assert_eq!(parts[1]["type"], "image_url"); |
| 408 | let plain = prompt_to_content(&blocks[..1], &[], None, &AudioHints::default()); |
| 409 | assert_eq!(plain, json!("/chords Am F")); |
| 410 | } |
| 411 | } |