Skip to main content

basis_tasks/
answer.rs

1//! What a task answering a caller hands back beside its text (ADR-0041 §7).
2//!
3//! A task's result is the text of its last turn, and a model cannot put a
4//! file in text. So a task spawned to answer someone —
5//! [`RunSpec::answering`](crate::RunSpec::answering) — is offered one tool,
6//! [`ATTACH_FILE`], that adds a workspace file to its answer. The file is
7//! copied into the task's own directory at the moment of the call: the answer
8//! is what the model chose then, not whatever the workspace file becomes by
9//! the time the caller reads it.
10//!
11//! The copies sit under `files/`, indexed by `files.json` — written
12//! atomically after each copy, so a reader sees every attached file or the
13//! index from before it, never a half-written entry.
14
15use 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
32/// The tool a task answering a caller attaches a file with.
33pub const ATTACH_FILE: &str = "attach_file";
34
35/// The most one task's answer carries in files: 6 MiB, which A2A's base64
36/// makes 8 MiB — mentra's default bound on a reply a caller reads.
37pub const MAX_ANSWER_FILES_BYTES: usize = 6 * 1024 * 1024;
38
39/// One file a task's answer carries.
40#[derive(Debug, Clone, PartialEq, Eq)]
41#[non_exhaustive]
42pub struct AnswerFile {
43    /// The name it goes by: the workspace file's own.
44    pub name: String,
45    pub media_type: String,
46    pub data: Vec<u8>,
47}
48
49/// One entry of `files.json`: what was attached, and the copy holding it.
50#[derive(Debug, Clone, Serialize, Deserialize)]
51struct Entry {
52    name: String,
53    media_type: String,
54    /// The copy's file name under `files/`.
55    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
72/// The files the task in `dir` attached to its answer, in the order it
73/// attached them. Empty for a task that attached none.
74pub(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
87/// [`ATTACH_FILE`] for the task whose directory is `dir`.
88pub(crate) struct AttachFile {
89    dir: PathBuf,
90}
91
92impl AttachFile {
93    pub(crate) fn new(dir: PathBuf) -> Self {
94        Self { dir }
95    }
96
97    /// Copies `data` in as `name` and records it, refusing past the bound.
98    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            // Reads like `read` does and writes only the task's own record,
146            // so it is never a question for whoever approves this task's
147            // consequential calls — a task told to refuse them can still
148            // answer with a file.
149            .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}