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