Skip to main content

lash_remote_protocol/
tools.rs

1//! Tool grants: schemas, call-path bindings, activation, scheduling, and
2//! retry policies.
3
4use std::collections::{BTreeMap, HashSet};
5
6use schemars::JsonSchema;
7use serde::{Deserialize, Serialize};
8
9use crate::ensure_protocol_version;
10use crate::llm::{RemoteSchemaContract, default_remote_input_schema};
11use crate::registry_errors::RemoteProtocolError;
12
13#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
14pub struct RemoteToolGrant {
15    pub protocol_version: u32,
16    pub id: String,
17    pub name: String,
18    #[serde(default, skip_serializing_if = "String::is_empty")]
19    pub description: String,
20    #[serde(default = "default_remote_input_schema")]
21    pub input_schema: RemoteSchemaContract,
22    #[serde(default)]
23    pub output_schema: RemoteSchemaContract,
24    #[serde(default, skip_serializing_if = "RemoteToolOutputContract::is_static")]
25    pub output_contract: RemoteToolOutputContract,
26    #[serde(default, skip_serializing_if = "Vec::is_empty")]
27    pub examples: Vec<String>,
28    #[serde(default, skip_serializing_if = "Option::is_none")]
29    pub activation: Option<RemoteToolActivation>,
30    #[serde(default, skip_serializing_if = "Option::is_none")]
31    pub argument_projection: Option<RemoteToolArgumentProjectionPolicy>,
32    #[serde(default, skip_serializing_if = "Option::is_none")]
33    pub scheduling: Option<RemoteToolScheduling>,
34    #[serde(default, skip_serializing_if = "Option::is_none")]
35    pub retry_policy: Option<RemoteToolRetryPolicy>,
36    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
37    pub bindings: BTreeMap<String, serde_json::Value>,
38}
39
40impl RemoteToolGrant {
41    pub fn binding_call_path(&self, binding_key: &str) -> Result<String, RemoteProtocolError> {
42        let binding = self.required_call_path_binding(binding_key)?;
43        Ok(format!(
44            "{}.{}",
45            binding.module_path.join("."),
46            binding.operation
47        ))
48    }
49
50    pub fn validate(&self) -> Result<(), RemoteProtocolError> {
51        ensure_protocol_version(self.protocol_version)?;
52        if self.id.trim().is_empty() {
53            return Err(RemoteProtocolError::InvalidToolGrant {
54                tool_name: self.name.clone(),
55                message: "tool grant id cannot be empty".to_string(),
56            });
57        }
58        if self.name.trim().is_empty() {
59            return Err(RemoteProtocolError::InvalidToolGrant {
60                tool_name: self.name.clone(),
61                message: "tool grant name cannot be empty".to_string(),
62            });
63        }
64        for key in self.bindings.keys() {
65            if key.trim().is_empty() {
66                return Err(RemoteProtocolError::InvalidToolGrant {
67                    tool_name: self.name.clone(),
68                    message: "tool grant binding keys cannot be empty".to_string(),
69                });
70            }
71        }
72        Ok(())
73    }
74
75    pub fn validate_all(grants: &[Self]) -> Result<(), RemoteProtocolError> {
76        let mut seen_ids = HashSet::new();
77        let mut seen_names = HashSet::new();
78        let mut seen_call_paths = HashSet::new();
79        for grant in grants {
80            grant.validate()?;
81            if !seen_ids.insert(grant.id.clone()) {
82                return Err(RemoteProtocolError::InvalidToolGrant {
83                    tool_name: grant.name.clone(),
84                    message: format!("duplicate tool grant id `{}`", grant.id),
85                });
86            }
87            if !seen_names.insert(grant.name.clone()) {
88                return Err(RemoteProtocolError::InvalidToolGrant {
89                    tool_name: grant.name.clone(),
90                    message: format!("duplicate tool grant name `{}`", grant.name),
91                });
92            }
93            for call_path in grant.call_path_bindings()? {
94                if !seen_call_paths.insert(call_path.clone()) {
95                    return Err(RemoteProtocolError::DuplicateRemoteCallPath { call_path });
96                }
97            }
98        }
99        Ok(())
100    }
101
102    pub fn call_path_bindings(&self) -> Result<Vec<String>, RemoteProtocolError> {
103        let mut paths = Vec::new();
104        for (key, value) in &self.bindings {
105            if let Some(binding) = RemoteCallPathBinding::from_value(value) {
106                validate_call_path_binding(&self.name, key, &binding)?;
107                paths.push(format!(
108                    "{}.{}",
109                    binding.module_path.join("."),
110                    binding.operation
111                ));
112            }
113        }
114        Ok(paths)
115    }
116
117    fn required_call_path_binding(
118        &self,
119        binding_key: &str,
120    ) -> Result<RemoteCallPathBinding, RemoteProtocolError> {
121        let Some(value) = self.bindings.get(binding_key) else {
122            return Err(RemoteProtocolError::MissingToolBinding {
123                tool_name: self.name.clone(),
124                binding: binding_key.to_string(),
125            });
126        };
127        let Some(binding) = RemoteCallPathBinding::from_value(value) else {
128            return Err(RemoteProtocolError::InvalidToolGrant {
129                tool_name: self.name.clone(),
130                message: format!("tool binding `{binding_key}` does not expose a call path"),
131            });
132        };
133        validate_call_path_binding(&self.name, binding_key, &binding)?;
134        Ok(binding)
135    }
136}
137
138#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
139pub struct RemoteCallPathBinding {
140    pub module_path: Vec<String>,
141    pub operation: String,
142}
143
144impl RemoteCallPathBinding {
145    fn from_value(value: &serde_json::Value) -> Option<Self> {
146        let module_path = value
147            .get("module_path")?
148            .as_array()?
149            .iter()
150            .map(|part| part.as_str().map(ToOwned::to_owned))
151            .collect::<Option<Vec<_>>>()?;
152        let operation = value.get("operation")?.as_str()?.to_string();
153        Some(Self {
154            module_path,
155            operation,
156        })
157    }
158}
159
160fn validate_call_path_binding(
161    tool_name: &str,
162    binding_key: &str,
163    binding: &RemoteCallPathBinding,
164) -> Result<(), RemoteProtocolError> {
165    if binding.module_path.is_empty() {
166        return Err(RemoteProtocolError::InvalidToolGrant {
167            tool_name: tool_name.to_string(),
168            message: format!("tool binding `{binding_key}` requires an explicit module path"),
169        });
170    }
171    if binding
172        .module_path
173        .iter()
174        .any(|part| part.trim().is_empty())
175    {
176        return Err(RemoteProtocolError::InvalidToolGrant {
177            tool_name: tool_name.to_string(),
178            message: format!(
179                "tool binding `{binding_key}` module path cannot contain empty segments"
180            ),
181        });
182    }
183    if binding.operation.trim().is_empty() {
184        return Err(RemoteProtocolError::InvalidToolGrant {
185            tool_name: tool_name.to_string(),
186            message: format!("tool binding `{binding_key}` requires an explicit operation"),
187        });
188    }
189    Ok(())
190}
191
192#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
193#[serde(rename_all = "snake_case")]
194pub enum RemoteToolActivation {
195    #[default]
196    Always,
197    Internal,
198}
199
200#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
201#[serde(rename_all = "snake_case")]
202pub enum RemoteToolScheduling {
203    #[default]
204    Parallel,
205    Serial,
206}
207
208#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
209#[serde(tag = "kind", rename_all = "snake_case")]
210pub enum RemoteToolOutputContract {
211    #[default]
212    Static,
213    FromInputSchema {
214        input_field: String,
215        #[serde(default, skip_serializing_if = "Option::is_none")]
216        default_schema: Option<serde_json::Value>,
217    },
218}
219
220impl RemoteToolOutputContract {
221    pub(crate) fn is_static(&self) -> bool {
222        matches!(self, Self::Static)
223    }
224}
225
226#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
227#[serde(tag = "kind", rename_all = "snake_case")]
228pub enum RemoteToolArgumentProjectionPolicy {
229    #[default]
230    MaterializeProjectedValues,
231    PreserveProjectedRefsInField {
232        field: String,
233    },
234}
235
236#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
237#[serde(tag = "type", rename_all = "snake_case")]
238pub enum RemoteToolRetryPolicy {
239    #[default]
240    Never,
241    Safe {
242        max_attempts: u32,
243        base_delay_ms: u64,
244        max_delay_ms: u64,
245    },
246    Idempotent {
247        max_attempts: u32,
248        base_delay_ms: u64,
249        max_delay_ms: u64,
250    },
251}