megacode_agent_sdk/
tools.rs

1//! Client tools (spec slice-ar R8.4): a definition with a typed handler, the
2//! context a call runs with, its output, and the `client_tool_result` line
3//! the SDK sends for it.
4
5use std::fmt::{self, Display};
6use std::future::Future;
7use std::pin::Pin;
8use std::sync::Arc;
9use std::time::Duration;
10
11use serde::de::DeserializeOwned;
12use serde_json::{Map, Value, json};
13use tokio::sync::watch;
14
15use crate::wire::{self, LINE_LIMIT, REASON_LIMIT};
16
17/// A tool's default timeout (R3.1.4).
18const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5 * 60);
19
20/// The longest result line the SDK sends: harness's line limit less room
21/// for the envelope (R8.4.6).
22pub(crate) const RESULT_LIMIT: usize = LINE_LIMIT - 64 * 1024;
23
24type BoxFuture<T> = Pin<Box<dyn Future<Output = T> + Send + 'static>>;
25
26/// A handler with its argument type erased: the arguments are deserialized
27/// inside the call, so a failure there is the call's failure.
28type Handler =
29    Arc<dyn Fn(Value, ToolContext) -> BoxFuture<Result<ToolOutput, String>> + Send + Sync>;
30
31/// A tool the program runs (R8.4.1): a name, a description, an input
32/// schema sent to the model as given, whether it asks before every call,
33/// how long harness waits for it, and an async handler.
34#[derive(Clone)]
35pub struct Tool {
36    pub(crate) name: String,
37    pub(crate) description: String,
38    pub(crate) schema: Value,
39    pub(crate) ask: bool,
40    pub(crate) timeout: Duration,
41    pub(crate) handler: Handler,
42}
43
44impl fmt::Debug for Tool {
45    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
46        f.debug_struct("Tool")
47            .field("name", &self.name)
48            .field("description", &self.description)
49            .field("schema", &self.schema)
50            .field("ask", &self.ask)
51            .field("timeout", &self.timeout)
52            .finish_non_exhaustive()
53    }
54}
55
56impl Tool {
57    /// A tool whose calls run `handler` with the model's arguments as a `T`.
58    ///
59    /// Before the handler runs, the arguments are checked against the
60    /// schema's top-level `required`, each property's `type` and
61    /// `additionalProperties: false`, then deserialized; a call that fails
62    /// either, an `Err`, or a panic answers `is_error` with the reason.
63    pub fn new<T, E, F, Fut>(
64        name: impl Into<String>,
65        description: impl Into<String>,
66        schema: Value,
67        handler: F,
68    ) -> Self
69    where
70        T: DeserializeOwned + Send + 'static,
71        E: Display + 'static,
72        F: Fn(T, ToolContext) -> Fut + Send + Sync + 'static,
73        Fut: Future<Output = Result<ToolOutput, E>> + Send + 'static,
74    {
75        let handler = Arc::new(handler);
76        let checked = schema.clone();
77        let erased: Handler = Arc::new(move |args: Value, ctx: ToolContext| {
78            if let Err(why) = check_arguments(&checked, &args) {
79                return Box::pin(std::future::ready(Err(format!("invalid arguments: {why}"))))
80                    as BoxFuture<_>;
81            }
82            match serde_json::from_value::<T>(args) {
83                Err(e) => Box::pin(std::future::ready(Err(format!(
84                    "the arguments do not fit the tool's input: {e}"
85                )))),
86                Ok(args) => {
87                    let running = handler(args, ctx);
88                    Box::pin(async move { running.await.map_err(|e| e.to_string()) })
89                }
90            }
91        });
92        Self {
93            name: name.into(),
94            description: description.into(),
95            schema,
96            ask: false,
97            timeout: DEFAULT_TIMEOUT,
98            handler: erased,
99        }
100    }
101
102    /// Ask before every call (R3.5.1); default `false`.
103    pub fn ask(mut self, ask: bool) -> Self {
104        self.ask = ask;
105        self
106    }
107
108    /// How long harness waits for a call before cancelling it; default 5
109    /// minutes, at most 24 hours. Sent in whole seconds, rounded up.
110    pub fn timeout(mut self, timeout: Duration) -> Self {
111        self.timeout = timeout;
112        self
113    }
114
115    /// The tool's name.
116    pub fn name(&self) -> &str {
117        &self.name
118    }
119
120    /// The definition as the options file's `client_tools` takes it (R3.1).
121    pub(crate) fn definition(&self) -> Value {
122        json!({
123            "name": self.name,
124            "description": self.description,
125            "input_schema": self.schema,
126            "ask": self.ask,
127            "timeout": format!("{}s", whole_seconds(self.timeout)),
128        })
129    }
130}
131
132/// `timeout` in whole seconds, rounded up, at least one.
133pub(crate) fn whole_seconds(timeout: Duration) -> u64 {
134    let up = u64::from(timeout.subsec_nanos() > 0);
135    timeout.as_secs().saturating_add(up).max(1)
136}
137
138/// What a call runs with (R8.4.4).
139#[derive(Clone, Debug)]
140#[non_exhaustive]
141pub struct ToolContext {
142    /// The call's id: the `tool_use` event's `id`.
143    pub call_id: String,
144    /// Fires when harness cancels the call, the session's process exits or
145    /// its connection fails, or the session is closed. A result returned
146    /// after it fired is dropped.
147    pub cancelled: CancelSignal,
148}
149
150/// A call's cancellation signal.
151#[derive(Clone, Debug)]
152pub struct CancelSignal(watch::Receiver<bool>);
153
154impl CancelSignal {
155    /// Whether the signal fired.
156    pub fn is_fired(&self) -> bool {
157        *self.0.borrow()
158    }
159
160    /// Waits until the signal fires.
161    pub async fn fired(&self) {
162        let mut rx = self.0.clone();
163        // A closed sender is a session gone: that is a cancellation too.
164        let _ = rx.wait_for(|fired| *fired).await;
165    }
166}
167
168/// A new signal and its trigger.
169pub(crate) fn signal() -> (watch::Sender<bool>, CancelSignal) {
170    let (tx, rx) = watch::channel(false);
171    (tx, CancelSignal(rx))
172}
173
174impl ToolContext {
175    pub(crate) fn new(call_id: String, cancelled: CancelSignal) -> Self {
176        Self { call_id, cancelled }
177    }
178}
179
180/// A tool's answer (R8.4.2): text, or text and one image.
181#[derive(Debug, Clone, PartialEq, Eq)]
182pub struct ToolOutput {
183    text: String,
184    image: Option<(String, Vec<u8>)>,
185}
186
187impl ToolOutput {
188    /// Text.
189    pub fn text(text: impl Into<String>) -> Self {
190        Self {
191            text: text.into(),
192            image: None,
193        }
194    }
195
196    /// Text and one image: its media type (`image/png`, `image/jpeg`,
197    /// `image/gif`, `image/webp`) and its bytes.
198    pub fn with_image(
199        text: impl Into<String>,
200        media_type: impl Into<String>,
201        bytes: impl Into<Vec<u8>>,
202    ) -> Self {
203        Self {
204            text: text.into(),
205            image: Some((media_type.into(), bytes.into())),
206        }
207    }
208}
209
210/// The `client_tool_result` line for call `id` (R3.3, R8.4.6): the output,
211/// or the error cut to 64 KiB; a line over [`RESULT_LIMIT`] is replaced by
212/// the oversize error.
213pub(crate) fn result_line(id: &str, outcome: &Result<ToolOutput, String>) -> String {
214    let line = match outcome {
215        Ok(out) => {
216            let mut content = vec![json!({ "type": "text", "text": out.text })];
217            if let Some((media_type, bytes)) = &out.image {
218                content.push(json!({
219                    "type": "image",
220                    "media_type": media_type,
221                    "data": base64(bytes),
222                }));
223            }
224            encoded(id, Value::Array(content), false)
225        }
226        Err(message) => error_line(id, message),
227    };
228    if line.len() > RESULT_LIMIT {
229        return error_line(
230            id,
231            &format!(
232                "the tool's result was too large ({} bytes; the limit is 8 MiB)",
233                line.len()
234            ),
235        );
236    }
237    line
238}
239
240/// An `is_error` result with `message` cut to 64 KiB.
241pub(crate) fn error_line(id: &str, message: &str) -> String {
242    let text = wire::cut(message, REASON_LIMIT);
243    encoded(id, json!([{ "type": "text", "text": text }]), true)
244}
245
246fn encoded(id: &str, content: Value, is_error: bool) -> String {
247    let mut line = serde_json::to_string(&json!({
248        "type": "client_tool_result",
249        "id": id,
250        "content": content,
251        "is_error": is_error,
252    }))
253    .expect("a JSON value encodes");
254    line.push('\n');
255    line
256}
257
258/// The schema subset R8.4.3 checks: the arguments an object when the
259/// schema's type is `object`, its top-level `required` present, each
260/// present property of the `type` it names, and no property beyond
261/// `properties` when `additionalProperties` is `false`. Nested schemas are
262/// left to deserialization.
263pub(crate) fn check_arguments(schema: &Value, args: &Value) -> Result<(), String> {
264    let Some(args) = args.as_object() else {
265        return Err(format!(
266            "the arguments must be an object, not {}",
267            kind(args)
268        ));
269    };
270    let empty = Map::new();
271    let properties = schema
272        .get("properties")
273        .and_then(Value::as_object)
274        .unwrap_or(&empty);
275    if let Some(required) = schema.get("required").and_then(Value::as_array) {
276        for name in required.iter().filter_map(Value::as_str) {
277            if !args.contains_key(name) {
278                return Err(format!("the required field {name:?} is missing"));
279            }
280        }
281    }
282    for (name, value) in args {
283        match properties.get(name) {
284            Some(property) => {
285                if let Some(ty) = property.get("type")
286                    && !fits(ty, value)
287                {
288                    return Err(format!(
289                        "the field {name:?} must be {}, not {}",
290                        expected(ty),
291                        kind(value)
292                    ));
293                }
294            }
295            None if schema.get("additionalProperties") == Some(&Value::Bool(false)) => {
296                return Err(format!("the field {name:?} is not in the schema"));
297            }
298            None => {}
299        }
300    }
301    Ok(())
302}
303
304/// Whether `value` is of the schema type `ty` (a word or a list of words);
305/// a word this check does not know passes.
306fn fits(ty: &Value, value: &Value) -> bool {
307    match ty {
308        Value::String(word) => fits_word(word, value),
309        Value::Array(words) => words
310            .iter()
311            .any(|w| w.as_str().is_none_or(|w| fits_word(w, value))),
312        _ => true,
313    }
314}
315
316fn fits_word(word: &str, value: &Value) -> bool {
317    match word {
318        "object" => value.is_object(),
319        "array" => value.is_array(),
320        "string" => value.is_string(),
321        "boolean" => value.is_boolean(),
322        "null" => value.is_null(),
323        "number" => value.is_number(),
324        "integer" => {
325            value.is_i64()
326                || value.is_u64()
327                || value
328                    .as_f64()
329                    .is_some_and(|f| f.is_finite() && f.fract() == 0.0)
330        }
331        _ => true,
332    }
333}
334
335fn expected(ty: &Value) -> String {
336    let article = |w: &str| match w {
337        "object" | "array" | "integer" => format!("an {w}"),
338        "null" => "null".to_string(),
339        _ => format!("a {w}"),
340    };
341    match ty {
342        Value::String(w) => article(w),
343        Value::Array(words) => words
344            .iter()
345            .filter_map(Value::as_str)
346            .map(article)
347            .collect::<Vec<_>>()
348            .join(" or "),
349        other => other.to_string(),
350    }
351}
352
353fn kind(value: &Value) -> &'static str {
354    match value {
355        Value::Null => "null",
356        Value::Bool(_) => "a boolean",
357        Value::Number(n) if n.is_f64() => "a number",
358        Value::Number(_) => "an integer",
359        Value::String(_) => "a string",
360        Value::Array(_) => "an array",
361        Value::Object(_) => "an object",
362    }
363}
364
365/// Standard base64 with padding.
366fn base64(bytes: &[u8]) -> String {
367    const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
368    let mut out = String::with_capacity(bytes.len().div_ceil(3) * 4);
369    for chunk in bytes.chunks(3) {
370        let b = [
371            chunk[0],
372            chunk.get(1).copied().unwrap_or(0),
373            chunk.get(2).copied().unwrap_or(0),
374        ];
375        let n = (u32::from(b[0]) << 16) | (u32::from(b[1]) << 8) | u32::from(b[2]);
376        for i in 0..4 {
377            if i <= chunk.len() {
378                out.push(char::from(ALPHABET[(n >> (18 - 6 * i)) as usize & 63]));
379            } else {
380                out.push('=');
381            }
382        }
383    }
384    out
385}
386
387/// What ended a task that did not return, as the end of a sentence:
388/// `panicked: <message>`, or `was stopped`.
389pub(crate) fn failure_text(error: tokio::task::JoinError) -> String {
390    if !error.is_panic() {
391        return "was stopped".into();
392    }
393    let payload = error.into_panic();
394    let message = if let Some(s) = payload.downcast_ref::<&str>() {
395        (*s).to_string()
396    } else if let Some(s) = payload.downcast_ref::<String>() {
397        s.clone()
398    } else {
399        "a panic with no message".into()
400    };
401    format!("panicked: {message}")
402}
403
404#[cfg(test)]
405mod tests {
406    use super::*;
407
408    fn schema() -> Value {
409        json!({
410            "type": "object",
411            "properties": {"merchant": {"type": "string"}, "n": {"type": "integer"},
412                           "x": {"type": ["number", "null"]}},
413            "required": ["merchant"],
414            "additionalProperties": false,
415        })
416    }
417
418    #[test]
419    fn the_schema_check_names_what_does_not_fit() {
420        assert!(check_arguments(&schema(), &json!({"merchant": "a", "n": 2})).is_ok());
421        assert!(check_arguments(&schema(), &json!({"merchant": "a", "n": 2.0})).is_ok());
422        assert!(check_arguments(&schema(), &json!({"merchant": "a", "x": null})).is_ok());
423        let e = check_arguments(&schema(), &json!({"n": 1})).unwrap_err();
424        assert_eq!(e, "the required field \"merchant\" is missing");
425        let e = check_arguments(&schema(), &json!({"merchant": "a", "n": "1"})).unwrap_err();
426        assert_eq!(e, "the field \"n\" must be an integer, not a string");
427        let e = check_arguments(&schema(), &json!({"merchant": "a", "n": 1.5})).unwrap_err();
428        assert_eq!(e, "the field \"n\" must be an integer, not a number");
429        let e = check_arguments(&schema(), &json!({"merchant": "a", "x": "1"})).unwrap_err();
430        assert_eq!(e, "the field \"x\" must be a number or null, not a string");
431        let e = check_arguments(&schema(), &json!({"merchant": "a", "tip": 1})).unwrap_err();
432        assert_eq!(e, "the field \"tip\" is not in the schema");
433        let e = check_arguments(&schema(), &json!([1])).unwrap_err();
434        assert_eq!(e, "the arguments must be an object, not an array");
435        // Without `additionalProperties: false`, extra fields pass.
436        assert!(check_arguments(&json!({"type": "object"}), &json!({"any": 1})).is_ok());
437    }
438
439    #[test]
440    fn deep_nesting_is_left_to_deserialization_and_never_recursed() {
441        let mut deep = json!(1);
442        for _ in 0..100 {
443            deep = json!({ "a": deep });
444        }
445        let schema = json!({"type": "object", "properties": {"a": {"type": "object"}}});
446        assert!(check_arguments(&schema, &json!({ "a": deep })).is_ok());
447    }
448
449    #[test]
450    fn a_timeout_is_sent_in_whole_seconds_rounded_up() {
451        assert_eq!(whole_seconds(Duration::from_millis(1500)), 2);
452        assert_eq!(whole_seconds(Duration::from_secs(3)), 3);
453        assert_eq!(whole_seconds(Duration::from_millis(10)), 1);
454        assert_eq!(whole_seconds(Duration::ZERO), 1);
455    }
456
457    #[test]
458    fn base64_matches_the_standard_alphabet_and_padding() {
459        assert_eq!(base64(b""), "");
460        assert_eq!(base64(b"f"), "Zg==");
461        assert_eq!(base64(b"fo"), "Zm8=");
462        assert_eq!(base64(b"foo"), "Zm9v");
463        assert_eq!(base64(b"foob"), "Zm9vYg==");
464        assert_eq!(base64(&[0xff, 0xfe, 0xfd]), "//79");
465    }
466
467    fn line_of_len(len: usize) -> String {
468        let empty = result_line("c", &Ok(ToolOutput::text(""))).len();
469        result_line("c", &Ok(ToolOutput::text("a".repeat(len - empty))))
470    }
471
472    /// R8.4.6, Review Focus 5: a line at the limit is sent as it is; one
473    /// byte over is the oversize error naming its size.
474    #[test]
475    fn a_result_is_sent_up_to_the_limit_and_replaced_one_byte_over() {
476        let at = line_of_len(RESULT_LIMIT);
477        assert_eq!(at.len(), RESULT_LIMIT);
478        assert!(at.contains("\"is_error\":false"));
479        let over = line_of_len(RESULT_LIMIT + 1);
480        let v: Value = serde_json::from_str(&over).unwrap();
481        assert_eq!(v["is_error"], true);
482        assert_eq!(
483            v["content"][0]["text"],
484            format!(
485                "the tool's result was too large ({} bytes; the limit is 8 MiB)",
486                RESULT_LIMIT + 1
487            )
488        );
489    }
490
491    /// Review Focus 5: an image at harness's `MAX_IMAGE_BYTES` (3,932,160
492    /// bytes) fits, its base64 whole.
493    #[test]
494    fn an_image_at_the_size_limit_fits_whole() {
495        let bytes = vec![0xabu8; 3_932_160];
496        let line = result_line("c", &Ok(ToolOutput::with_image("", "image/png", bytes)));
497        let v: Value = serde_json::from_str(&line).unwrap();
498        assert_eq!(v["is_error"], false);
499        assert_eq!(v["content"][1]["data"].as_str().unwrap().len(), 5_242_880);
500    }
501
502    /// R8.4.6: an error message is cut to 64 KiB on a character boundary;
503    /// hostile text stays inside one framed line.
504    #[test]
505    fn an_error_is_cut_to_64_kib_and_hostile_text_stays_framed() {
506        let long = "é".repeat(REASON_LIMIT);
507        let v: Value = serde_json::from_str(&error_line("c", &long)).unwrap();
508        let text = v["content"][0]["text"].as_str().unwrap();
509        assert!(text.len() <= REASON_LIMIT && text.len() >= REASON_LIMIT - 1);
510        let hostile = "a\nb\r\u{1b}]0;x\u{7}\u{2028}\u{2029}";
511        let line = result_line("c", &Ok(ToolOutput::text(hostile)));
512        assert_eq!(line.matches('\n').count(), 1);
513        assert!(line.ends_with('\n'));
514        let v: Value = serde_json::from_str(&line).unwrap();
515        assert_eq!(v["content"][0]["text"], hostile);
516    }
517}
megacode_agent_sdk/tools.rs

megacode_agent_sdk/
tools.rs
