Skip to main content

lash_remote_protocol/
tools.rs

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