Skip to main content

lash_remote_protocol/protocol/
tools.rs

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