1use 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}