llm_tool/

rust_tool.rs

//! Strongly-typed Rust tool trait and type-erasure machinery.

use alloc::{borrow::Cow, boxed::Box, format, string::ToString};
use core::{future::Future, pin::Pin};

use super::types::{
    OutputSanitization, ToolContext, ToolDefinition, ToolEffect, ToolError, ToolOutput,
    ToolProvenance,
};

/// Convenience type for tools that take no parameters.
#[derive(
    Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize, schemars::JsonSchema,
)]
pub struct EmptyParams {}

/// A custom tool implemented entirely in Rust with strongly-typed parameters.
///
/// Define your parameters as a struct deriving [`serde::Deserialize`] and
/// `JsonSchema`, then implement this trait to provide the tool's logic.
/// The JSON Schema sent to the model is derived automatically from the
/// params struct — doc comments on fields become parameter descriptions.
///
/// Tools are async: for I/O-bound work (HTTP, filesystem, subprocess) the
/// runtime stays unblocked. Sync tools just don't `.await` anything — the
/// compiler optimizes the state machine to an immediate return.
///
/// # Example
///
/// ```rust
/// use llm_tool::{JsonSchema, RustTool, ToolContext, ToolError, ToolOutput};
/// use serde::Deserialize;
///
/// #[derive(Deserialize, JsonSchema)]
/// struct FlashParams {
///     /// Target device identifier.
///     device_id: String,
///     /// Path to the firmware image.
///     image_path: String,
/// }
///
/// struct FlashDevice;
///
/// impl RustTool for FlashDevice {
///     type Params = FlashParams;
///     const NAME: &'static str = "flash_device";
///     const DESCRIPTION: &'static str = "Flashes firmware to a connected device.";
///
///     async fn call(
///         &self,
///         params: Self::Params,
///         _ctx: &ToolContext,
///     ) -> Result<ToolOutput, ToolError> {
///         Ok(format!("Flashed {} to {}", params.image_path, params.device_id).into())
///     }
/// }
/// ```
pub trait RustTool: Send + Sync {
    /// The strongly-typed parameters struct.
    ///
    /// Derive [`serde::Deserialize`] and `JsonSchema` on your params struct.
    /// `JsonSchema` auto-generates the parameter schema sent to the model;
    /// `Deserialize` parses the model's JSON arguments into your struct.
    type Params: serde::de::DeserializeOwned + schemars::JsonSchema + Send;

    /// Unique tool name (e.g. `"flash_device"`).
    const NAME: &'static str;

    /// Human-readable description shown to the model.
    const DESCRIPTION: &'static str;

    /// Execution effect of the tool.
    const EFFECT: ToolEffect = if Self::IDEMPOTENT {
        ToolEffect::ReadOnly
    } else {
        ToolEffect::Mutating
    };

    /// Whether the tool is idempotent / read-only safe for concurrent wave execution.
    const IDEMPOTENT: bool = false;

    /// Data-provenance and trust classification of the tool.
    ///
    /// Defaults to [`ToolProvenance::Trusted`]. Tools with [`ToolProvenance::ExternalMcp`]
    /// or [`ToolProvenance::UntrustedNetwork`] automatically default [`Self::OUTPUT_SANITIZATION`]
    /// to [`OutputSanitization::Quarantine`] unless overridden.
    const PROVENANCE: ToolProvenance = ToolProvenance::Trusted;

    /// Shorthand: neutralize control tokens in this tool's output.
    ///
    /// Defaults to `true` (secure by default). Set to `false` to opt out, or prefer
    /// setting [`Self::OUTPUT_SANITIZATION`] directly.
    const SANITIZE_OUTPUT: bool = true;

    /// Shorthand: wrap this tool's output in a `<tool_output_quarantine>` envelope.
    ///
    /// Takes precedence over [`Self::SANITIZE_OUTPUT`]. Prefer setting
    /// [`Self::OUTPUT_SANITIZATION`] directly.
    const QUARANTINE_OUTPUT: bool = false;

    /// Automatic prompt-injection sanitization policy applied to this tool's output.
    ///
    /// Derived from the shorthands and provenance: [`OutputSanitization::Quarantine`] if
    /// [`Self::QUARANTINE_OUTPUT`] or [`Self::PROVENANCE`] is an untrusted external source,
    /// else [`OutputSanitization::SanitizeTokens`] if [`Self::SANITIZE_OUTPUT`] (the default),
    /// else [`OutputSanitization::None`].
    /// Override it to pick a tagged envelope, or [`OutputSanitization::None`] only for
    /// output you fully control.
    const OUTPUT_SANITIZATION: OutputSanitization =
        if Self::QUARANTINE_OUTPUT || Self::PROVENANCE.is_untrusted_source() {
            OutputSanitization::Quarantine
        } else if Self::SANITIZE_OUTPUT {
            OutputSanitization::SanitizeTokens
        } else {
            OutputSanitization::None
        };

    /// Return the execution effect of the tool.
    fn effect(&self) -> ToolEffect {
        Self::EFFECT
    }

    /// Return the data-provenance tier of the tool.
    fn provenance(&self) -> ToolProvenance {
        Self::PROVENANCE
    }

    /// Return whether the tool is idempotent (read-only).
    fn is_idempotent(&self) -> bool {
        matches!(self.effect(), ToolEffect::ReadOnly)
    }

    /// Return the automatic prompt-injection sanitization policy for this tool's output.
    fn output_sanitization(&self) -> OutputSanitization {
        Self::OUTPUT_SANITIZATION
    }

    /// Return the tool description used in [`ToolDefinition`].
    ///
    /// The default returns [`Self::DESCRIPTION`] (the static string from a
    /// doc comment or template body). When using
    /// `#[llm_tool(description_file = "...", context = ...)]`, the generated
    /// implementation overrides this to render the template with runtime
    /// variables on each call (parsed once via `LazyLock`), falling back to
    /// [`Self::DESCRIPTION`] if a render ever fails.
    fn description(&self) -> Cow<'static, str> {
        Cow::Borrowed(Self::DESCRIPTION)
    }

    /// Execute the tool with typed parameters and an execution context.
    ///
    /// Async to support I/O-bound tools (HTTP, filesystem, subprocess).
    /// Sync tools just compute and return — the async wrapper is zero-cost.
    ///
    /// The `ctx` parameter provides access to conversation metadata and a
    /// shared key-value state store. Tools that don't need context can simply
    /// ignore it with `_ctx`.
    ///
    /// # Errors
    ///
    /// Returns `Err(ToolError)` if the tool execution fails.
    fn call(
        &self,
        params: Self::Params,
        ctx: &ToolContext,
    ) -> impl Future<Output = Result<ToolOutput, ToolError>> + Send;
}

/// Shared [`alloc::sync::Arc`]-backed wrapper for any [`RustTool`] implementation,
/// enabling cheap cloning across registries and subagent scopes.
#[derive(Debug)]
pub struct SharedRustTool<T>(pub alloc::sync::Arc<T>);

impl<T> Clone for SharedRustTool<T> {
    fn clone(&self) -> Self {
        Self(alloc::sync::Arc::clone(&self.0))
    }
}

impl<T> SharedRustTool<T> {
    /// Wrap `tool` in a new [`SharedRustTool`].
    #[must_use]
    pub fn new(tool: T) -> Self {
        Self(alloc::sync::Arc::new(tool))
    }

    /// Borrow the underlying tool.
    #[must_use]
    pub fn as_inner(&self) -> &T {
        &self.0
    }
}

impl<T> From<alloc::sync::Arc<T>> for SharedRustTool<T> {
    fn from(arc: alloc::sync::Arc<T>) -> Self {
        Self(arc)
    }
}

impl<T: RustTool> RustTool for alloc::sync::Arc<T> {
    type Params = T::Params;
    const NAME: &'static str = T::NAME;
    const DESCRIPTION: &'static str = T::DESCRIPTION;
    const EFFECT: ToolEffect = T::EFFECT;
    const IDEMPOTENT: bool = T::IDEMPOTENT;
    const PROVENANCE: ToolProvenance = T::PROVENANCE;
    const SANITIZE_OUTPUT: bool = T::SANITIZE_OUTPUT;
    const QUARANTINE_OUTPUT: bool = T::QUARANTINE_OUTPUT;
    const OUTPUT_SANITIZATION: OutputSanitization = T::OUTPUT_SANITIZATION;

    fn effect(&self) -> ToolEffect {
        (**self).effect()
    }

    fn provenance(&self) -> ToolProvenance {
        (**self).provenance()
    }

    fn is_idempotent(&self) -> bool {
        (**self).is_idempotent()
    }

    fn output_sanitization(&self) -> OutputSanitization {
        (**self).output_sanitization()
    }

    fn description(&self) -> Cow<'static, str> {
        (**self).description()
    }

    fn call(
        &self,
        params: Self::Params,
        ctx: &ToolContext,
    ) -> impl Future<Output = Result<ToolOutput, ToolError>> + Send {
        (**self).call(params, ctx)
    }
}

impl<T: RustTool> RustTool for SharedRustTool<T> {
    type Params = T::Params;
    const NAME: &'static str = T::NAME;
    const DESCRIPTION: &'static str = T::DESCRIPTION;
    const EFFECT: ToolEffect = T::EFFECT;
    const IDEMPOTENT: bool = T::IDEMPOTENT;
    const PROVENANCE: ToolProvenance = T::PROVENANCE;
    const SANITIZE_OUTPUT: bool = T::SANITIZE_OUTPUT;
    const QUARANTINE_OUTPUT: bool = T::QUARANTINE_OUTPUT;
    const OUTPUT_SANITIZATION: OutputSanitization = T::OUTPUT_SANITIZATION;

    fn effect(&self) -> ToolEffect {
        self.0.effect()
    }

    fn provenance(&self) -> ToolProvenance {
        self.0.provenance()
    }

    fn is_idempotent(&self) -> bool {
        self.0.is_idempotent()
    }

    fn output_sanitization(&self) -> OutputSanitization {
        self.0.output_sanitization()
    }

    fn description(&self) -> Cow<'static, str> {
        self.0.description()
    }

    fn call(
        &self,
        params: Self::Params,
        ctx: &ToolContext,
    ) -> impl Future<Output = Result<ToolOutput, ToolError>> + Send {
        self.0.call(params, ctx)
    }
}

/// Build a [`ToolDefinition`] from any [`RustTool`] implementor.
///
/// The generated schema is sanitized for broad compatibility with LLM
/// SDKs that expect `"type"` to always be a single string
/// (not the array form `["string", "null"]` that schemars emits for
/// `Option<T>` fields).
///
/// # Errors
///
/// Returns `Err` if the JSON schema serialization fails.
pub fn definition_of<T: RustTool>(tool: &T) -> Result<ToolDefinition, ToolError> {
    let schema = schemars::schema_for!(T::Params);
    let mut parameter_schema = serde_json::to_value(schema).map_err(|e| {
        ToolError::new(format!(
            "Failed to serialize schema for tool '{}': {e}",
            T::NAME
        ))
    })?;
    sanitize_schema_types(&mut parameter_schema);
    Ok(ToolDefinition::with_effect(
        T::NAME,
        tool.description().into_owned(),
        parameter_schema,
        tool.effect(),
    )
    .with_provenance(tool.provenance())
    .with_output_sanitization(tool.output_sanitization()))
}

/// Recursively sanitize JSON Schema `"type"` fields for Go genai compatibility.
///
/// `schemars` emits `"type": ["string", "null"]` for `Option<String>` fields
/// (the nullable type-array form). The Go genai SDK's `Schema.Type`
/// is a single `genai.Type` enum, so it can't unmarshal an array.
///
/// This function walks the schema tree and replaces any array `type` with the
/// first non-`"null"` element. For example:
/// - `["string", "null"]` → `"string"`
/// - `["integer", "null"]` → `"integer"`
const SCHEMA_TYPE_KEY: &str = "type";
const SCHEMA_NULL_TYPE: &str = "null";

fn sanitize_schema_types(value: &mut serde_json::Value) {
    match value {
        serde_json::Value::Object(map) => {
            // If "type" is an array (e.g. ["string", "null"]), pick the first
            // non-"null" element and replace with it as a scalar type.
            let replacement = match map.get(SCHEMA_TYPE_KEY) {
                Some(serde_json::Value::Array(arr)) => {
                    let non_null = arr
                        .iter()
                        .find(
                            |v| !matches!(v, serde_json::Value::String(s) if s == SCHEMA_NULL_TYPE),
                        )
                        .cloned();
                    non_null.or_else(|| arr.first().cloned())
                }
                _ => None,
            };
            if let Some(val) = replacement {
                map.insert(SCHEMA_TYPE_KEY.to_string(), val);
            }
            for val in map.values_mut() {
                sanitize_schema_types(val);
            }
        }
        serde_json::Value::Array(arr) => {
            for item in arr {
                sanitize_schema_types(item);
            }
        }
        _ => {}
    }
}

/// Type-erased future returned when dispatching a boxed tool.
///
/// Produced by the crate-internal type-erasure layer that lets heterogeneous
/// [`RustTool`] implementations share one registry.
pub type BoxToolFuture<'a> =
    Pin<Box<dyn Future<Output = Result<ToolOutput, ToolError>> + Send + 'a>>;

/// Type-erased wrapper enabling heterogeneous tool storage.
///
/// Boxes the future from [`RustTool::call`] so we can store different tool
/// types in the same `HashMap<String, Box<dyn ErasedTool>>`.
pub trait ErasedTool: Send + Sync {
    /// Deserialize `args` and call the handler, returning a boxed future.
    fn call_erased<'a>(
        &'a self,
        args: serde_json::Value,
        ctx: &'a ToolContext,
    ) -> BoxToolFuture<'a>;

    /// Deserialize `args_json` directly from a JSON string and call the handler,
    /// returning a boxed future without intermediate DOM allocation.
    fn call_erased_str<'a>(&'a self, args_json: &'a str, ctx: &'a ToolContext)
    -> BoxToolFuture<'a>;
}

impl<T: RustTool> ErasedTool for T {
    fn call_erased<'a>(
        &'a self,
        args: serde_json::Value,
        ctx: &'a ToolContext,
    ) -> BoxToolFuture<'a> {
        Box::pin(async move {
            let params: T::Params = serde_json::from_value(args).map_err(|e| {
                ToolError::new(format!("Failed to deserialize tool parameters: {e}"))
            })?;
            let output = self.call(params, ctx).await?;
            Ok(self.output_sanitization().apply(output))
        })
    }

    fn call_erased_str<'a>(
        &'a self,
        args_json: &'a str,
        ctx: &'a ToolContext,
    ) -> BoxToolFuture<'a> {
        Box::pin(async move {
            let params: T::Params = serde_json::from_str(args_json).map_err(|e| {
                ToolError::new(format!("Failed to deserialize tool parameters: {e}"))
            })?;
            let output = self.call(params, ctx).await?;
            Ok(self.output_sanitization().apply(output))
        })
    }
}