llm_tool/types/
definition.rs

1//! Tool execution effect classification and [`ToolDefinition`] metadata.
2
3use alloc::{string::String, sync::Arc};
4
5use serde::{Deserialize, Serialize};
6
7use super::OutputSanitization;
8use crate::compat::OnceLock;
9
10/// Effect classification for a tool.
11#[derive(
12    Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default, schemars::JsonSchema,
13)]
14#[serde(rename_all = "snake_case")]
15pub enum ToolEffect {
16    /// Read-only, side-effect free invocation. Safe for concurrent wave execution.
17    ReadOnly,
18    /// Mutates system state (default). Safe for normal execution.
19    #[default]
20    Mutating,
21    /// Destructive operation (e.g. deleting files, wiping data, dropping resources).
22    /// Requires explicit confirmation before parallel batch dispatch.
23    Destructive,
24}
25
26impl ToolEffect {
27    /// Returns `true` if the tool effect is read-only / idempotent.
28    #[must_use]
29    pub const fn is_read_only(self) -> bool {
30        matches!(self, Self::ReadOnly)
31    }
32
33    /// Returns `true` if the tool effect is destructive.
34    #[must_use]
35    pub const fn is_destructive(self) -> bool {
36        matches!(self, Self::Destructive)
37    }
38}
39
40/// Describes a custom tool that can be registered with an agent.
41#[non_exhaustive]
42#[derive(Debug, Serialize)]
43pub struct ToolDefinition {
44    /// Unique tool name (e.g. `"flash_device"`).
45    pub name: String,
46    /// Human-readable description shown to the model.
47    pub description: String,
48    /// JSON Schema describing the tool's parameters.
49    pub parameter_schema: serde_json::Value,
50    /// Execution effect classification of the tool.
51    pub effect: ToolEffect,
52    /// Automatic prompt-injection sanitization policy applied to the tool's output.
53    #[serde(default, skip_serializing_if = "OutputSanitization::is_none")]
54    pub output_sanitization: OutputSanitization,
55    /// Pre-serialized compact JSON Schema string cached lazily on first access.
56    #[serde(skip)]
57    parameter_schema_json: OnceLock<Arc<str>>,
58}
59
60impl Clone for ToolDefinition {
61    fn clone(&self) -> Self {
62        Self {
63            name: self.name.clone(),
64            description: self.description.clone(),
65            parameter_schema: self.parameter_schema.clone(),
66            effect: self.effect,
67            output_sanitization: self.output_sanitization,
68            parameter_schema_json: OnceLock::new(),
69        }
70    }
71}
72
73impl PartialEq for ToolDefinition {
74    fn eq(&self, other: &Self) -> bool {
75        self.name == other.name
76            && self.description == other.description
77            && self.parameter_schema == other.parameter_schema
78            && self.effect == other.effect
79            && self.output_sanitization == other.output_sanitization
80    }
81}
82
83impl Eq for ToolDefinition {}
84
85impl<'de> Deserialize<'de> for ToolDefinition {
86    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
87    where
88        D: serde::Deserializer<'de>,
89    {
90        #[derive(Deserialize)]
91        struct RawDef {
92            name: String,
93            description: String,
94            parameter_schema: serde_json::Value,
95            #[serde(default)]
96            effect: Option<ToolEffect>,
97            #[serde(default)]
98            idempotent: Option<bool>,
99            #[serde(default)]
100            output_sanitization: OutputSanitization,
101        }
102
103        let raw = RawDef::deserialize(deserializer)?;
104        let effect = match (raw.effect, raw.idempotent) {
105            (Some(eff), _) => eff,
106            (None, Some(true)) => ToolEffect::ReadOnly,
107            (None, Some(false) | None) => ToolEffect::Mutating,
108        };
109
110        Ok(Self {
111            name: raw.name,
112            description: raw.description,
113            parameter_schema: raw.parameter_schema,
114            effect,
115            output_sanitization: raw.output_sanitization,
116            parameter_schema_json: OnceLock::new(),
117        })
118    }
119}
120
121impl ToolDefinition {
122    /// Construct a [`ToolDefinition`] with default [`ToolEffect::Mutating`]
123    /// and the secure default [`OutputSanitization::SanitizeTokens`].
124    #[must_use]
125    pub fn new(
126        name: impl Into<String>,
127        description: impl Into<String>,
128        parameter_schema: serde_json::Value,
129    ) -> Self {
130        Self::with_effect(name, description, parameter_schema, ToolEffect::Mutating)
131    }
132
133    /// Construct a [`ToolDefinition`] directly from a parameter struct `T` implementing [`schemars::JsonSchema`].
134    ///
135    /// # Errors
136    ///
137    /// Returns the [`serde_json::Error`] if the generated schema cannot be converted to a JSON value.
138    pub fn from_schema<T: schemars::JsonSchema>(
139        name: impl Into<String>,
140        description: impl Into<String>,
141    ) -> Result<Self, serde_json::Error> {
142        let parameter_schema = serde_json::to_value(schemars::schema_for!(T))?;
143        Ok(Self::new(name, description, parameter_schema))
144    }
145
146    /// Construct a [`ToolDefinition`] with an explicit [`ToolEffect`].
147    #[must_use]
148    pub fn with_effect(
149        name: impl Into<String>,
150        description: impl Into<String>,
151        parameter_schema: serde_json::Value,
152        effect: ToolEffect,
153    ) -> Self {
154        Self {
155            name: name.into(),
156            description: description.into(),
157            parameter_schema,
158            effect,
159            output_sanitization: OutputSanitization::default(),
160            parameter_schema_json: OnceLock::new(),
161        }
162    }
163
164    /// Attach an [`OutputSanitization`] policy to this [`ToolDefinition`]. Chainable.
165    #[must_use]
166    pub const fn with_output_sanitization(mut self, sanitization: OutputSanitization) -> Self {
167        self.output_sanitization = sanitization;
168        self
169    }
170
171    /// Update `parameter_schema` and invalidate any cached JSON serialization.
172    pub fn set_parameter_schema(&mut self, parameter_schema: serde_json::Value) {
173        self.parameter_schema = parameter_schema;
174        self.parameter_schema_json = OnceLock::new();
175    }
176
177    /// Borrow `parameter_schema` mutably and invalidate any cached JSON serialization.
178    pub fn parameter_schema_mut(&mut self) -> &mut serde_json::Value {
179        self.parameter_schema_json = OnceLock::new();
180        &mut self.parameter_schema
181    }
182
183    /// Return the pre-serialized JSON schema string (`&str`), lazily derived
184    /// from `parameter_schema`.
185    #[must_use]
186    pub fn parameter_schema_json(&self) -> &str {
187        self.parameter_schema_json_arc()
188    }
189
190    /// Return the pre-serialized JSON schema string (`&Arc<str>`), lazily derived
191    /// from `parameter_schema`.
192    #[must_use]
193    pub fn parameter_schema_json_arc(&self) -> &Arc<str> {
194        self.parameter_schema_json.get_or_init(|| {
195            match serde_json::to_string(&self.parameter_schema) {
196                Ok(s) => Arc::from(s),
197                Err(err) => {
198                    tracing::warn!(error = %err, tool = %self.name, "Failed to serialize parameter_schema to JSON");
199                    Arc::from("{}")
200                }
201            }
202        })
203    }
204
205    /// Whether repeated invocations with identical arguments are side-effect free / idempotent.
206    ///
207    /// Preserved for backward compatibility.
208    #[must_use]
209    pub const fn is_idempotent(&self) -> bool {
210        matches!(self.effect, ToolEffect::ReadOnly)
211    }
212}
213
214#[cfg(test)]
215mod tests {
216    use super::*;
217
218    #[derive(schemars::JsonSchema, serde::Deserialize)]
219    struct LookupParams {
220        path: String,
221        max_lines: u32,
222    }
223
224    #[test]
225    fn from_schema_derives_typed_parameter_schema() {
226        let def = ToolDefinition::from_schema::<LookupParams>("lookup", "Look up a file")
227            .expect("schema serializes");
228        assert_eq!(def.name, "lookup");
229        assert_eq!(def.description, "Look up a file");
230        let props = def.parameter_schema["properties"]
231            .as_object()
232            .expect("object schema has properties");
233        assert!(props.contains_key("path"));
234        assert!(props.contains_key("max_lines"));
235        assert!(def.parameter_schema_json().contains("\"max_lines\""));
236
237        let parsed: LookupParams =
238            serde_json::from_str(r#"{"path":"src/lib.rs","max_lines":40}"#).expect("valid args");
239        assert_eq!(parsed.path, "src/lib.rs");
240        assert_eq!(parsed.max_lines, 40);
241    }
242}
llm_tool/types/definition.rs

llm_tool/types/
definition.rs
