artifactize_tools/
result.rs1use std::path::Path;
2
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5
6use super::image;
7
8const TEXT_LIMIT: usize = 64 * 1024;
10const JSON_LIMIT: usize = 512 * 1024;
12const RESULT_LIMIT: usize = 8 * 1024 * 1024;
14const MAX_CONTENT_BLOCKS: usize = 32;
16
17#[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 Image {
30 data: String,
31 #[serde(rename = "mimeType")]
32 mime_type: image::Mime,
33 },
34}
35
36#[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#[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 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#[derive(Debug, thiserror::Error)]
90#[error("{0}")]
91pub struct InvalidResult(&'static str);
92
93#[derive(Debug, Clone, PartialEq)]
95pub enum ToolResult {
96 Success(ContentBlocks<Content>),
97 Error(ErrorMessage),
98}
99
100impl ToolResult {
101 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 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 #[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#[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#[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
292const 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
298const CSI_PREFIX: &[u8] = b"\x1b[";
300const 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}