Skip to main content

artifactize_tools/
result.rs

1use std::path::Path;
2
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5
6use super::image;
7
8/// Text stays readable in the next provider turn rather than filling its whole context.
9const TEXT_LIMIT: usize = 64 * 1024;
10/// Bound one dynamic JSON block before it reaches the provider.
11const JSON_LIMIT: usize = 512 * 1024;
12/// Cap the whole result, including base64 images and JSON envelope overhead.
13const RESULT_LIMIT: usize = 8 * 1024 * 1024;
14/// Bound multimodal block fan-out and per-block validation work even when each block is small.
15const MAX_CONTENT_BLOCKS: usize = 32;
16
17/// Self-contained tool content for persistence and multimodal Agent turns.
18/// Registry results contain validated images, never references to temporary files.
19#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
20#[serde(tag = "type", rename_all = "lowercase", deny_unknown_fields)]
21pub enum Content {
22    Text {
23        text: String,
24    },
25    Json {
26        data: Value,
27    },
28    /// Canonical base64 of at most 4 MiB; mimeType is image/png, image/jpeg or image/webp.
29    Image {
30        data: String,
31        #[serde(rename = "mimeType")]
32        mime_type: image::Mime,
33    },
34}
35
36/// A nonempty, bounded sequence of result blocks. Its contents cannot be mutated
37/// into an empty or oversized result after construction.
38#[derive(Debug, Clone, PartialEq)]
39pub struct ContentBlocks<T>(Vec<T>);
40
41impl<T> ContentBlocks<T> {
42    pub fn new(content: Vec<T>) -> Result<Self, InvalidResult> {
43        if !(1..=MAX_CONTENT_BLOCKS).contains(&content.len()) {
44            return Err(InvalidResult(
45                "tool result must contain 1..=32 content blocks",
46            ));
47        }
48        Ok(Self(content))
49    }
50
51    pub fn as_slice(&self) -> &[T] {
52        &self.0
53    }
54
55    pub fn into_vec(self) -> Vec<T> {
56        self.0
57    }
58}
59
60/// A single nonblank text block. Keeping the block itself lets callers borrow
61/// the same content slice for both success and error results.
62#[derive(Debug, Clone, PartialEq)]
63pub struct ErrorMessage(Content);
64
65impl ErrorMessage {
66    pub fn new(message: impl Into<String>) -> Result<Self, InvalidResult> {
67        let text = message.into();
68        if text.trim().is_empty() {
69            return Err(InvalidResult("tool error message must not be blank"));
70        }
71        Ok(Self(Content::Text { text }))
72    }
73
74    /// Replace a blank diagnostic with "Tool failed." without changing nonblank text.
75    pub fn diagnostic(message: impl Into<String>) -> Self {
76        Self::new(message)
77            .unwrap_or_else(|_| Self::new("Tool failed.").expect("fallback is nonblank"))
78    }
79
80    pub fn as_str(&self) -> &str {
81        let Content::Text { text } = &self.0 else {
82            unreachable!("error messages are constructed only from text")
83        };
84        text
85    }
86}
87
88/// A saved or constructed result contradicts the result type's invariants.
89#[derive(Debug, thiserror::Error)]
90#[error("{0}")]
91pub struct InvalidResult(&'static str);
92
93/// Successful multimodal content or one nonblank textual diagnostic.
94#[derive(Debug, Clone, PartialEq)]
95pub enum ToolResult {
96    Success(ContentBlocks<Content>),
97    Error(ErrorMessage),
98}
99
100impl ToolResult {
101    /// A single successful block is nonempty by construction.
102    pub fn success(content: Content) -> Self {
103        Self::Success(ContentBlocks(vec![content]))
104    }
105
106    pub fn try_success(content: Vec<Content>) -> Result<Self, InvalidResult> {
107        ContentBlocks::new(content).map(Self::Success)
108    }
109
110    /// Blank messages become "Tool failed." so even a missing diagnostic is a
111    /// valid error result. Nonblank messages are preserved byte-for-byte.
112    pub fn error(message: impl Into<String>) -> Self {
113        Self::Error(ErrorMessage::diagnostic(message))
114    }
115
116    pub fn is_error(&self) -> bool {
117        matches!(self, Self::Error(_))
118    }
119
120    pub fn content(&self) -> &[Content] {
121        match self {
122            Self::Success(content) => content.as_slice(),
123            Self::Error(message) => std::slice::from_ref(&message.0),
124        }
125    }
126
127    pub fn into_content(self) -> Vec<Content> {
128        match self {
129            Self::Success(content) => content.into_vec(),
130            Self::Error(message) => vec![message.0],
131        }
132    }
133
134    fn from_wire(content: Vec<Content>, is_error: bool) -> Result<Self, InvalidResult> {
135        if !is_error {
136            return Self::try_success(content);
137        }
138        let [Content::Text { text }] = content.as_slice() else {
139            return Err(InvalidResult(
140                "tool error must contain exactly one text block",
141            ));
142        };
143        ErrorMessage::new(text.clone()).map(Self::Error)
144    }
145}
146
147impl Serialize for ToolResult {
148    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
149        use serde::ser::SerializeStruct;
150        let mut result =
151            serializer.serialize_struct("ToolResult", if self.is_error() { 2 } else { 1 })?;
152        result.serialize_field("content", self.content())?;
153        if self.is_error() {
154            result.serialize_field("isError", &true)?;
155        }
156        result.end()
157    }
158}
159
160impl<'de> Deserialize<'de> for ToolResult {
161    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
162        // Saved images are already inline and validated. Do not read files or
163        // impose command-output byte limits when replaying an existing result.
164        #[derive(Deserialize)]
165        #[serde(rename_all = "camelCase", deny_unknown_fields)]
166        struct SavedResult {
167            content: Vec<Content>,
168            #[serde(default)]
169            is_error: bool,
170        }
171        let saved = SavedResult::deserialize(deserializer)?;
172        Self::from_wire(saved.content, saved.is_error).map_err(serde::de::Error::custom)
173    }
174}
175
176#[derive(Deserialize)]
177#[serde(rename_all = "camelCase", deny_unknown_fields)]
178struct WireResult {
179    content: Vec<WireContent>,
180    #[serde(default)]
181    is_error: bool,
182}
183
184#[derive(Deserialize)]
185#[serde(tag = "type", rename_all = "lowercase", deny_unknown_fields)]
186enum WireContent {
187    Text { text: String },
188    Json { data: Value },
189    Image(WireImage),
190}
191
192/// An image block carries its bytes inline or names a file in the output directory, never
193/// both and never neither.
194#[derive(Deserialize)]
195#[serde(untagged)]
196enum WireImage {
197    Data(InlineImage),
198    Path(ImageFile),
199}
200
201#[derive(Deserialize)]
202#[serde(rename_all = "camelCase", deny_unknown_fields)]
203struct InlineImage {
204    mime_type: image::Mime,
205    data: String,
206}
207
208#[derive(Deserialize)]
209#[serde(rename_all = "camelCase", deny_unknown_fields)]
210struct ImageFile {
211    mime_type: image::Mime,
212    path: String,
213}
214
215/// Malformed or oversized command output; callers expose their own execution diagnostic.
216#[derive(Debug, thiserror::Error)]
217#[error("Agent tool returned invalid output.")]
218pub struct InvalidOutput;
219
220pub fn parse(stdout: &[u8], output_dir: &Path) -> Result<ToolResult, InvalidOutput> {
221    let wire: WireResult = serde_json::from_slice(stdout).map_err(|_| InvalidOutput)?;
222    if !(1..=MAX_CONTENT_BLOCKS).contains(&wire.content.len())
223        || wire.is_error
224            && !matches!(
225                wire.content.as_slice(),
226                [WireContent::Text { text }] if !text.trim().is_empty()
227            )
228    {
229        return Err(InvalidOutput);
230    }
231    let mut content = Vec::with_capacity(wire.content.len());
232    let mut size = 0;
233    for block in wire.content {
234        let block = match block {
235            WireContent::Text { text } if text.len() <= TEXT_LIMIT => Content::Text { text },
236            WireContent::Json { data }
237                if serde_json::to_vec(&data).map_err(|_| InvalidOutput)?.len() <= JSON_LIMIT =>
238            {
239                Content::Json { data }
240            }
241            WireContent::Image(WireImage::Data(InlineImage { mime_type, data })) => {
242                image::from_base64(&data, mime_type.as_str()).map_err(|_| InvalidOutput)?
243            }
244            WireContent::Image(WireImage::Path(ImageFile { mime_type, path })) => {
245                image::from_output(output_dir, &path, mime_type.as_str())
246                    .map_err(|_| InvalidOutput)?
247            }
248            _ => return Err(InvalidOutput),
249        };
250        size += serde_json::to_vec(&block).map_err(|_| InvalidOutput)?.len();
251        if size > RESULT_LIMIT {
252            return Err(InvalidOutput);
253        }
254        content.push(block);
255    }
256    let result = ToolResult::from_wire(content, wire.is_error).map_err(|_| InvalidOutput)?;
257    if serde_json::to_vec(&result)
258        .map_err(|_| InvalidOutput)?
259        .len()
260        > RESULT_LIMIT
261    {
262        return Err(InvalidOutput);
263    }
264    Ok(result)
265}
266
267pub fn plain(stdout: &[u8], successful: bool, truncated: bool) -> ToolResult {
268    let clean = clean_output(stdout);
269    let text = String::from_utf8(clean).expect("clean output is UTF-8");
270    let failed = !successful;
271    let mut text = if failed {
272        format!("Agent tool exited unsuccessfully.\n{text}")
273    } else {
274        text
275    };
276    if text.len() > TEXT_LIMIT || truncated {
277        let suffix = "\n[output truncated]";
278        let mut end = (TEXT_LIMIT - suffix.len()).min(text.len());
279        while !text.is_char_boundary(end) {
280            end -= 1;
281        }
282        text.truncate(end);
283        text.push_str(suffix);
284    }
285    if failed {
286        ToolResult::error(text)
287    } else {
288        ToolResult::success(Content::Text { text })
289    }
290}
291
292/// ECMA-48 control sequence bytes: after `ESC [`, parameter bytes, then intermediate bytes,
293/// then one final byte end the sequence.
294const CSI_PARAMETER: std::ops::RangeInclusive<u8> = 0x30..=0x3f;
295const CSI_INTERMEDIATE: std::ops::RangeInclusive<u8> = 0x20..=0x2f;
296const CSI_FINAL: std::ops::RangeInclusive<u8> = 0x40..=0x7e;
297
298/// The two-byte prefix introducing an ECMA-48 control sequence.
299const CSI_PREFIX: &[u8] = b"\x1b[";
300/// C0 controls except tab, line feed and carriage return are not printable tool output.
301const C0_CONTROLS: &[std::ops::RangeInclusive<u8>] = &[0..=8, 11..=12, 14..=31];
302
303pub fn clean_output(bytes: &[u8]) -> Vec<u8> {
304    let text = String::from_utf8_lossy(bytes);
305    let mut bytes = text.as_bytes();
306    let mut clean = Vec::with_capacity(bytes.len());
307    while let Some((&byte, tail)) = bytes.split_first() {
308        if bytes.starts_with(CSI_PREFIX) {
309            let mut end = CSI_PREFIX.len();
310            while bytes
311                .get(end)
312                .is_some_and(|byte| CSI_PARAMETER.contains(byte))
313            {
314                end += 1;
315            }
316            while bytes
317                .get(end)
318                .is_some_and(|byte| CSI_INTERMEDIATE.contains(byte))
319            {
320                end += 1;
321            }
322            if bytes.get(end).is_some_and(|byte| CSI_FINAL.contains(byte)) {
323                bytes = &bytes[end + 1..];
324                continue;
325            }
326        }
327        if !C0_CONTROLS.iter().any(|range| range.contains(&byte)) {
328            clean.push(byte);
329        }
330        bytes = tail;
331    }
332    clean
333}