megacode-agent-sdk 0.31.0

//! Client tools (spec slice-ar R8.4): a definition with a typed handler, the
//! context a call runs with, its output, and the `client_tool_result` line
//! the SDK sends for it.

use std::fmt::{self, Display};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;

use serde::de::DeserializeOwned;
use serde_json::{Map, Value, json};
use tokio::sync::watch;

use crate::wire::{self, LINE_LIMIT, REASON_LIMIT};

/// A tool's default timeout (R3.1.4).
const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5 * 60);

/// The longest result line the SDK sends: harness's line limit less room
/// for the envelope (R8.4.6).
pub(crate) const RESULT_LIMIT: usize = LINE_LIMIT - 64 * 1024;

type BoxFuture<T> = Pin<Box<dyn Future<Output = T> + Send + 'static>>;

/// A handler with its argument type erased: the arguments are deserialized
/// inside the call, so a failure there is the call's failure.
type Handler =
    Arc<dyn Fn(Value, ToolContext) -> BoxFuture<Result<ToolOutput, String>> + Send + Sync>;

/// A tool the program runs (R8.4.1): a name, a description, an input
/// schema sent to the model as given, whether it asks before every call,
/// how long harness waits for it, and an async handler.
#[derive(Clone)]
pub struct Tool {
    pub(crate) name: String,
    pub(crate) description: String,
    pub(crate) schema: Value,
    pub(crate) ask: bool,
    pub(crate) timeout: Duration,
    pub(crate) handler: Handler,
}

impl fmt::Debug for Tool {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Tool")
            .field("name", &self.name)
            .field("description", &self.description)
            .field("schema", &self.schema)
            .field("ask", &self.ask)
            .field("timeout", &self.timeout)
            .finish_non_exhaustive()
    }
}

impl Tool {
    /// A tool whose calls run `handler` with the model's arguments as a `T`.
    ///
    /// Before the handler runs, the arguments are checked against the
    /// schema's top-level `required`, each property's `type` and
    /// `additionalProperties: false`, then deserialized; a call that fails
    /// either, an `Err`, or a panic answers `is_error` with the reason.
    pub fn new<T, E, F, Fut>(
        name: impl Into<String>,
        description: impl Into<String>,
        schema: Value,
        handler: F,
    ) -> Self
    where
        T: DeserializeOwned + Send + 'static,
        E: Display + 'static,
        F: Fn(T, ToolContext) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<ToolOutput, E>> + Send + 'static,
    {
        let handler = Arc::new(handler);
        let checked = schema.clone();
        let erased: Handler = Arc::new(move |args: Value, ctx: ToolContext| {
            if let Err(why) = check_arguments(&checked, &args) {
                return Box::pin(std::future::ready(Err(format!("invalid arguments: {why}"))))
                    as BoxFuture<_>;
            }
            match serde_json::from_value::<T>(args) {
                Err(e) => Box::pin(std::future::ready(Err(format!(
                    "the arguments do not fit the tool's input: {e}"
                )))),
                Ok(args) => {
                    let running = handler(args, ctx);
                    Box::pin(async move { running.await.map_err(|e| e.to_string()) })
                }
            }
        });
        Self {
            name: name.into(),
            description: description.into(),
            schema,
            ask: false,
            timeout: DEFAULT_TIMEOUT,
            handler: erased,
        }
    }

    /// Ask before every call (R3.5.1); default `false`.
    pub fn ask(mut self, ask: bool) -> Self {
        self.ask = ask;
        self
    }

    /// How long harness waits for a call before cancelling it; default 5
    /// minutes, at most 24 hours. Sent in whole seconds, rounded up.
    pub fn timeout(mut self, timeout: Duration) -> Self {
        self.timeout = timeout;
        self
    }

    /// The tool's name.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// The definition as the options file's `client_tools` takes it (R3.1).
    pub(crate) fn definition(&self) -> Value {
        json!({
            "name": self.name,
            "description": self.description,
            "input_schema": self.schema,
            "ask": self.ask,
            "timeout": format!("{}s", whole_seconds(self.timeout)),
        })
    }
}

/// `timeout` in whole seconds, rounded up, at least one.
pub(crate) fn whole_seconds(timeout: Duration) -> u64 {
    let up = u64::from(timeout.subsec_nanos() > 0);
    timeout.as_secs().saturating_add(up).max(1)
}

/// What a call runs with (R8.4.4).
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct ToolContext {
    /// The call's id: the `tool_use` event's `id`.
    pub call_id: String,
    /// Fires when harness cancels the call, the session's process exits or
    /// its connection fails, or the session is closed. A result returned
    /// after it fired is dropped.
    pub cancelled: CancelSignal,
}

/// A call's cancellation signal.
#[derive(Clone, Debug)]
pub struct CancelSignal(watch::Receiver<bool>);

impl CancelSignal {
    /// Whether the signal fired.
    pub fn is_fired(&self) -> bool {
        *self.0.borrow()
    }

    /// Waits until the signal fires.
    pub async fn fired(&self) {
        let mut rx = self.0.clone();
        // A closed sender is a session gone: that is a cancellation too.
        let _ = rx.wait_for(|fired| *fired).await;
    }
}

/// A new signal and its trigger.
pub(crate) fn signal() -> (watch::Sender<bool>, CancelSignal) {
    let (tx, rx) = watch::channel(false);
    (tx, CancelSignal(rx))
}

impl ToolContext {
    pub(crate) fn new(call_id: String, cancelled: CancelSignal) -> Self {
        Self { call_id, cancelled }
    }
}

/// A tool's answer (R8.4.2): text, or text and one image.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ToolOutput {
    text: String,
    image: Option<(String, Vec<u8>)>,
}

impl ToolOutput {
    /// Text.
    pub fn text(text: impl Into<String>) -> Self {
        Self {
            text: text.into(),
            image: None,
        }
    }

    /// Text and one image: its media type (`image/png`, `image/jpeg`,
    /// `image/gif`, `image/webp`) and its bytes.
    pub fn with_image(
        text: impl Into<String>,
        media_type: impl Into<String>,
        bytes: impl Into<Vec<u8>>,
    ) -> Self {
        Self {
            text: text.into(),
            image: Some((media_type.into(), bytes.into())),
        }
    }
}

/// The `client_tool_result` line for call `id` (R3.3, R8.4.6): the output,
/// or the error cut to 64 KiB; a line over [`RESULT_LIMIT`] is replaced by
/// the oversize error.
pub(crate) fn result_line(id: &str, outcome: &Result<ToolOutput, String>) -> String {
    let line = match outcome {
        Ok(out) => {
            let mut content = vec![json!({ "type": "text", "text": out.text })];
            if let Some((media_type, bytes)) = &out.image {
                content.push(json!({
                    "type": "image",
                    "media_type": media_type,
                    "data": base64(bytes),
                }));
            }
            encoded(id, Value::Array(content), false)
        }
        Err(message) => error_line(id, message),
    };
    if line.len() > RESULT_LIMIT {
        return error_line(
            id,
            &format!(
                "the tool's result was too large ({} bytes; the limit is 8 MiB)",
                line.len()
            ),
        );
    }
    line
}

/// An `is_error` result with `message` cut to 64 KiB.
pub(crate) fn error_line(id: &str, message: &str) -> String {
    let text = wire::cut(message, REASON_LIMIT);
    encoded(id, json!([{ "type": "text", "text": text }]), true)
}

fn encoded(id: &str, content: Value, is_error: bool) -> String {
    let mut line = serde_json::to_string(&json!({
        "type": "client_tool_result",
        "id": id,
        "content": content,
        "is_error": is_error,
    }))
    .expect("a JSON value encodes");
    line.push('\n');
    line
}

/// The schema subset R8.4.3 checks: the arguments an object when the
/// schema's type is `object`, its top-level `required` present, each
/// present property of the `type` it names, and no property beyond
/// `properties` when `additionalProperties` is `false`. Nested schemas are
/// left to deserialization.
pub(crate) fn check_arguments(schema: &Value, args: &Value) -> Result<(), String> {
    let Some(args) = args.as_object() else {
        return Err(format!(
            "the arguments must be an object, not {}",
            kind(args)
        ));
    };
    let empty = Map::new();
    let properties = schema
        .get("properties")
        .and_then(Value::as_object)
        .unwrap_or(&empty);
    if let Some(required) = schema.get("required").and_then(Value::as_array) {
        for name in required.iter().filter_map(Value::as_str) {
            if !args.contains_key(name) {
                return Err(format!("the required field {name:?} is missing"));
            }
        }
    }
    for (name, value) in args {
        match properties.get(name) {
            Some(property) => {
                if let Some(ty) = property.get("type")
                    && !fits(ty, value)
                {
                    return Err(format!(
                        "the field {name:?} must be {}, not {}",
                        expected(ty),
                        kind(value)
                    ));
                }
            }
            None if schema.get("additionalProperties") == Some(&Value::Bool(false)) => {
                return Err(format!("the field {name:?} is not in the schema"));
            }
            None => {}
        }
    }
    Ok(())
}

/// Whether `value` is of the schema type `ty` (a word or a list of words);
/// a word this check does not know passes.
fn fits(ty: &Value, value: &Value) -> bool {
    match ty {
        Value::String(word) => fits_word(word, value),
        Value::Array(words) => words
            .iter()
            .any(|w| w.as_str().is_none_or(|w| fits_word(w, value))),
        _ => true,
    }
}

fn fits_word(word: &str, value: &Value) -> bool {
    match word {
        "object" => value.is_object(),
        "array" => value.is_array(),
        "string" => value.is_string(),
        "boolean" => value.is_boolean(),
        "null" => value.is_null(),
        "number" => value.is_number(),
        "integer" => {
            value.is_i64()
                || value.is_u64()
                || value
                    .as_f64()
                    .is_some_and(|f| f.is_finite() && f.fract() == 0.0)
        }
        _ => true,
    }
}

fn expected(ty: &Value) -> String {
    let article = |w: &str| match w {
        "object" | "array" | "integer" => format!("an {w}"),
        "null" => "null".to_string(),
        _ => format!("a {w}"),
    };
    match ty {
        Value::String(w) => article(w),
        Value::Array(words) => words
            .iter()
            .filter_map(Value::as_str)
            .map(article)
            .collect::<Vec<_>>()
            .join(" or "),
        other => other.to_string(),
    }
}

fn kind(value: &Value) -> &'static str {
    match value {
        Value::Null => "null",
        Value::Bool(_) => "a boolean",
        Value::Number(n) if n.is_f64() => "a number",
        Value::Number(_) => "an integer",
        Value::String(_) => "a string",
        Value::Array(_) => "an array",
        Value::Object(_) => "an object",
    }
}

/// Standard base64 with padding.
fn base64(bytes: &[u8]) -> String {
    const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    let mut out = String::with_capacity(bytes.len().div_ceil(3) * 4);
    for chunk in bytes.chunks(3) {
        let b = [
            chunk[0],
            chunk.get(1).copied().unwrap_or(0),
            chunk.get(2).copied().unwrap_or(0),
        ];
        let n = (u32::from(b[0]) << 16) | (u32::from(b[1]) << 8) | u32::from(b[2]);
        for i in 0..4 {
            if i <= chunk.len() {
                out.push(char::from(ALPHABET[(n >> (18 - 6 * i)) as usize & 63]));
            } else {
                out.push('=');
            }
        }
    }
    out
}

/// What ended a task that did not return, as the end of a sentence:
/// `panicked: <message>`, or `was stopped`.
pub(crate) fn failure_text(error: tokio::task::JoinError) -> String {
    if !error.is_panic() {
        return "was stopped".into();
    }
    let payload = error.into_panic();
    let message = if let Some(s) = payload.downcast_ref::<&str>() {
        (*s).to_string()
    } else if let Some(s) = payload.downcast_ref::<String>() {
        s.clone()
    } else {
        "a panic with no message".into()
    };
    format!("panicked: {message}")
}

#[cfg(test)]
mod tests {
    use super::*;

    fn schema() -> Value {
        json!({
            "type": "object",
            "properties": {"merchant": {"type": "string"}, "n": {"type": "integer"},
                           "x": {"type": ["number", "null"]}},
            "required": ["merchant"],
            "additionalProperties": false,
        })
    }

    #[test]
    fn the_schema_check_names_what_does_not_fit() {
        assert!(check_arguments(&schema(), &json!({"merchant": "a", "n": 2})).is_ok());
        assert!(check_arguments(&schema(), &json!({"merchant": "a", "n": 2.0})).is_ok());
        assert!(check_arguments(&schema(), &json!({"merchant": "a", "x": null})).is_ok());
        let e = check_arguments(&schema(), &json!({"n": 1})).unwrap_err();
        assert_eq!(e, "the required field \"merchant\" is missing");
        let e = check_arguments(&schema(), &json!({"merchant": "a", "n": "1"})).unwrap_err();
        assert_eq!(e, "the field \"n\" must be an integer, not a string");
        let e = check_arguments(&schema(), &json!({"merchant": "a", "n": 1.5})).unwrap_err();
        assert_eq!(e, "the field \"n\" must be an integer, not a number");
        let e = check_arguments(&schema(), &json!({"merchant": "a", "x": "1"})).unwrap_err();
        assert_eq!(e, "the field \"x\" must be a number or null, not a string");
        let e = check_arguments(&schema(), &json!({"merchant": "a", "tip": 1})).unwrap_err();
        assert_eq!(e, "the field \"tip\" is not in the schema");
        let e = check_arguments(&schema(), &json!([1])).unwrap_err();
        assert_eq!(e, "the arguments must be an object, not an array");
        // Without `additionalProperties: false`, extra fields pass.
        assert!(check_arguments(&json!({"type": "object"}), &json!({"any": 1})).is_ok());
    }

    #[test]
    fn deep_nesting_is_left_to_deserialization_and_never_recursed() {
        let mut deep = json!(1);
        for _ in 0..100 {
            deep = json!({ "a": deep });
        }
        let schema = json!({"type": "object", "properties": {"a": {"type": "object"}}});
        assert!(check_arguments(&schema, &json!({ "a": deep })).is_ok());
    }

    #[test]
    fn a_timeout_is_sent_in_whole_seconds_rounded_up() {
        assert_eq!(whole_seconds(Duration::from_millis(1500)), 2);
        assert_eq!(whole_seconds(Duration::from_secs(3)), 3);
        assert_eq!(whole_seconds(Duration::from_millis(10)), 1);
        assert_eq!(whole_seconds(Duration::ZERO), 1);
    }

    #[test]
    fn base64_matches_the_standard_alphabet_and_padding() {
        assert_eq!(base64(b""), "");
        assert_eq!(base64(b"f"), "Zg==");
        assert_eq!(base64(b"fo"), "Zm8=");
        assert_eq!(base64(b"foo"), "Zm9v");
        assert_eq!(base64(b"foob"), "Zm9vYg==");
        assert_eq!(base64(&[0xff, 0xfe, 0xfd]), "//79");
    }

    fn line_of_len(len: usize) -> String {
        let empty = result_line("c", &Ok(ToolOutput::text(""))).len();
        result_line("c", &Ok(ToolOutput::text("a".repeat(len - empty))))
    }

    /// R8.4.6, Review Focus 5: a line at the limit is sent as it is; one
    /// byte over is the oversize error naming its size.
    #[test]
    fn a_result_is_sent_up_to_the_limit_and_replaced_one_byte_over() {
        let at = line_of_len(RESULT_LIMIT);
        assert_eq!(at.len(), RESULT_LIMIT);
        assert!(at.contains("\"is_error\":false"));
        let over = line_of_len(RESULT_LIMIT + 1);
        let v: Value = serde_json::from_str(&over).unwrap();
        assert_eq!(v["is_error"], true);
        assert_eq!(
            v["content"][0]["text"],
            format!(
                "the tool's result was too large ({} bytes; the limit is 8 MiB)",
                RESULT_LIMIT + 1
            )
        );
    }

    /// Review Focus 5: an image at harness's `MAX_IMAGE_BYTES` (3,932,160
    /// bytes) fits, its base64 whole.
    #[test]
    fn an_image_at_the_size_limit_fits_whole() {
        let bytes = vec![0xabu8; 3_932_160];
        let line = result_line("c", &Ok(ToolOutput::with_image("", "image/png", bytes)));
        let v: Value = serde_json::from_str(&line).unwrap();
        assert_eq!(v["is_error"], false);
        assert_eq!(v["content"][1]["data"].as_str().unwrap().len(), 5_242_880);
    }

    /// R8.4.6: an error message is cut to 64 KiB on a character boundary;
    /// hostile text stays inside one framed line.
    #[test]
    fn an_error_is_cut_to_64_kib_and_hostile_text_stays_framed() {
        let long = "é".repeat(REASON_LIMIT);
        let v: Value = serde_json::from_str(&error_line("c", &long)).unwrap();
        let text = v["content"][0]["text"].as_str().unwrap();
        assert!(text.len() <= REASON_LIMIT && text.len() >= REASON_LIMIT - 1);
        let hostile = "a\nb\r\u{1b}]0;x\u{7}\u{2028}\u{2029}";
        let line = result_line("c", &Ok(ToolOutput::text(hostile)));
        assert_eq!(line.matches('\n').count(), 1);
        assert!(line.ends_with('\n'));
        let v: Value = serde_json::from_str(&line).unwrap();
        assert_eq!(v["content"][0]["text"], hostile);
    }
}