1use std::{
16 fs, io,
17 path::{Path, PathBuf},
18};
19
20use basis::{
21 async_trait, media,
22 tools::{
23 RuntimeToolDescriptor, ToolContext, ToolDefinition, ToolDurability, ToolExecutor,
24 ToolResult, ToolSideEffectLevel,
25 },
26};
27use serde::{Deserialize, Serialize};
28use serde_json::{Value, json};
29
30use crate::data_dir::write_private_atomic;
31
32pub const ATTACH_FILE: &str = "attach_file";
34
35pub const MAX_ANSWER_FILES_BYTES: usize = 6 * 1024 * 1024;
38
39#[derive(Debug, Clone, PartialEq, Eq)]
41#[non_exhaustive]
42pub struct AnswerFile {
43 pub name: String,
45 pub media_type: String,
46 pub data: Vec<u8>,
47}
48
49#[derive(Debug, Clone, Serialize, Deserialize)]
51struct Entry {
52 name: String,
53 media_type: String,
54 stored: String,
56 bytes: usize,
57}
58
59fn index(dir: &Path) -> PathBuf {
60 dir.join("files.json")
61}
62
63fn load(dir: &Path) -> io::Result<Vec<Entry>> {
64 match fs::read(index(dir)) {
65 Ok(bytes) => serde_json::from_slice(&bytes)
66 .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)),
67 Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(Vec::new()),
68 Err(error) => Err(error),
69 }
70}
71
72pub(crate) fn read(dir: &Path) -> io::Result<Vec<AnswerFile>> {
75 load(dir)?
76 .into_iter()
77 .map(|entry| {
78 Ok(AnswerFile {
79 data: fs::read(dir.join("files").join(&entry.stored))?,
80 name: entry.name,
81 media_type: entry.media_type,
82 })
83 })
84 .collect()
85}
86
87pub(crate) struct AttachFile {
89 dir: PathBuf,
90}
91
92impl AttachFile {
93 pub(crate) fn new(dir: PathBuf) -> Self {
94 Self { dir }
95 }
96
97 fn attach(&self, name: &str, data: &[u8]) -> Result<Entry, String> {
99 let mut entries = load(&self.dir).map_err(|error| format!("read files.json: {error}"))?;
100 let total: usize = entries.iter().map(|entry| entry.bytes).sum();
101 if total.saturating_add(data.len()) > MAX_ANSWER_FILES_BYTES {
102 return Err(format!(
103 "an answer carries at most {} MiB of files, and this one would take it past \
104 that; attach fewer, or smaller ones",
105 MAX_ANSWER_FILES_BYTES / (1024 * 1024)
106 ));
107 }
108 let stored = format!("{}-{name}", entries.len() + 1);
109 let files = self.dir.join("files");
110 fs::create_dir_all(&files).map_err(|error| format!("create files/: {error}"))?;
111 write_private_atomic(&files.join(&stored), data)
112 .map_err(|error| format!("copy {name}: {error}"))?;
113 let entry = Entry {
114 name: name.to_string(),
115 media_type: media::media_type(name, data).to_string(),
116 stored,
117 bytes: data.len(),
118 };
119 entries.push(entry.clone());
120 let json = serde_json::to_vec_pretty(&entries).map_err(|error| error.to_string())?;
121 write_private_atomic(&index(&self.dir), &json)
122 .map_err(|error| format!("write files.json: {error}"))?;
123 Ok(entry)
124 }
125}
126
127impl ToolDefinition for AttachFile {
128 fn descriptor(&self) -> RuntimeToolDescriptor {
129 RuntimeToolDescriptor::builder(ATTACH_FILE)
130 .description(
131 "Attach a workspace file to your answer, so whoever asked receives the file \
132 itself as well as your words — a screenshot, a report, a patch. The file is \
133 copied as it is now. Up to 6 MiB of files in all.",
134 )
135 .input_schema(json!({
136 "type": "object",
137 "properties": {
138 "path": {
139 "type": "string",
140 "description": "The workspace path of the file to attach.",
141 },
142 },
143 "required": ["path"],
144 }))
145 .side_effect_level(ToolSideEffectLevel::None)
150 .durability(ToolDurability::ReplaySafe)
151 .build()
152 }
153}
154
155#[async_trait]
156impl ToolExecutor for AttachFile {
157 async fn execute_mut(&self, ctx: ToolContext<'_>, input: Value) -> ToolResult {
158 let raw = input
159 .get("path")
160 .and_then(Value::as_str)
161 .filter(|path| !path.trim().is_empty())
162 .ok_or("`path` is the workspace path of the file to attach")?;
163 let (path, data) = ctx.read_file_bytes(raw)?;
164 if path.is_dir() {
165 return Err(format!(
166 "`{raw}` is a directory; attach the files in it by name"
167 ));
168 }
169 let name = path
170 .file_name()
171 .and_then(|name| media::safe_file_name(&name.to_string_lossy()))
172 .unwrap_or_else(|| "file".to_string());
173 let entry = self.attach(&name, &data)?;
174 Ok(format!(
175 "attached {} ({}, {} bytes); it goes back with your answer",
176 entry.name, entry.media_type, entry.bytes
177 ))
178 }
179}
180
181#[cfg(test)]
182mod tests {
183 use super::*;
184
185 #[test]
186 fn attached_files_are_read_back_in_order_with_their_types() {
187 let dir = tempfile::tempdir().unwrap();
188 let tool = AttachFile::new(dir.path().to_path_buf());
189 tool.attach("shot.png", b"\x89PNG\r\n\x1a\nrest").unwrap();
190 tool.attach("notes.md", b"# done").unwrap();
191
192 let files = read(dir.path()).unwrap();
193 assert_eq!(
194 files,
195 vec![
196 AnswerFile {
197 name: "shot.png".into(),
198 media_type: "image/png".into(),
199 data: b"\x89PNG\r\n\x1a\nrest".to_vec(),
200 },
201 AnswerFile {
202 name: "notes.md".into(),
203 media_type: "text/markdown".into(),
204 data: b"# done".to_vec(),
205 },
206 ]
207 );
208 }
209
210 #[test]
211 fn a_task_that_attached_nothing_has_no_files() {
212 let dir = tempfile::tempdir().unwrap();
213 assert!(read(dir.path()).unwrap().is_empty());
214 }
215
216 #[test]
217 fn the_answer_is_bounded_in_total() {
218 let dir = tempfile::tempdir().unwrap();
219 let tool = AttachFile::new(dir.path().to_path_buf());
220 tool.attach("a.bin", &vec![0; MAX_ANSWER_FILES_BYTES - 1])
221 .unwrap();
222 let refused = tool.attach("b.bin", &[0, 0]).unwrap_err();
223 assert!(refused.contains("at most 6 MiB"), "{refused}");
224 assert_eq!(
225 read(dir.path()).unwrap().len(),
226 1,
227 "the refused one is not recorded"
228 );
229 }
230}