agentos_native_sidecar_core/
tools.rs

1use agentos_sidecar_protocol::protocol::RegisterHostCallbacksRequest;
2use serde_json::Value;
3use std::collections::{BTreeMap, BTreeSet};
4use std::error::Error;
5use std::fmt;
6
7pub const DEFAULT_TOOL_TIMEOUT_MS: u64 = 30_000;
8pub const MAX_TOOL_TIMEOUT_MS: u64 = 300_000;
9pub const MAX_REGISTERED_TOOLKITS: usize = 64;
10pub const MAX_REGISTERED_TOOLS_PER_VM: usize = 256;
11pub const MAX_TOOLS_PER_TOOLKIT: usize = 64;
12pub const MAX_TOOLKIT_NAME_LENGTH: usize = 64;
13pub const MAX_TOOL_NAME_LENGTH: usize = 64;
14pub const MAX_TOOL_DESCRIPTION_LENGTH: usize = 200;
15pub const MAX_TOOL_SCHEMA_BYTES: usize = 16 * 1024;
16pub const MAX_TOOL_SCHEMA_DEPTH: usize = 32;
17pub const MAX_TOOL_EXAMPLES_PER_TOOL: usize = 16;
18pub const MAX_TOOL_EXAMPLE_INPUT_BYTES: usize = 4 * 1024;
19
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum ToolRegistrationError {
22    InvalidState(String),
23    Conflict(String),
24}
25
26impl fmt::Display for ToolRegistrationError {
27    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
28        match self {
29            Self::InvalidState(message) | Self::Conflict(message) => f.write_str(message),
30        }
31    }
32}
33
34impl Error for ToolRegistrationError {}
35
36pub fn validate_toolkit_registration(
37    payload: &RegisterHostCallbacksRequest,
38) -> Result<(), ToolRegistrationError> {
39    validate_toolkit_name(&payload.name)?;
40    if payload.description.is_empty() {
41        return Err(ToolRegistrationError::InvalidState(format!(
42            "toolkit {} is missing a description",
43            payload.name
44        )));
45    }
46    validate_description_length(
47        &format!("Toolkit \"{}\"", payload.name),
48        &payload.description,
49    )?;
50    validate_command_aliases("command alias", &payload.command_aliases)?;
51    validate_command_aliases("registry command alias", &payload.registry_command_aliases)?;
52    for alias in &payload.command_aliases {
53        if payload.registry_command_aliases.contains(alias) {
54            return Err(ToolRegistrationError::InvalidState(format!(
55                "host callback command alias must not also be a registry command alias: {alias}"
56            )));
57        }
58    }
59    if payload.callbacks.is_empty() {
60        return Err(ToolRegistrationError::InvalidState(format!(
61            "toolkit {} must define at least one tool",
62            payload.name
63        )));
64    }
65    if payload.callbacks.len() > MAX_TOOLS_PER_TOOLKIT {
66        return Err(ToolRegistrationError::InvalidState(format!(
67            "toolkit {} defines {} tools, max is {MAX_TOOLS_PER_TOOLKIT}",
68            payload.name,
69            payload.callbacks.len()
70        )));
71    }
72    for (tool_name, tool) in &payload.callbacks {
73        validate_tool_name(tool_name)?;
74        if tool.description.is_empty() {
75            return Err(ToolRegistrationError::InvalidState(format!(
76                "tool {} in toolkit {} is missing a description",
77                tool_name, payload.name
78            )));
79        }
80        validate_description_length(
81            &format!("Tool \"{}/{}\"", payload.name, tool_name),
82            &tool.description,
83        )?;
84        let tool_input_schema: Value =
85            serde_json::from_str(&tool.input_schema).map_err(|error| {
86                ToolRegistrationError::InvalidState(format!(
87                    "Tool \"{}/{}\" input schema is invalid JSON: {error}",
88                    payload.name, tool_name
89                ))
90            })?;
91        validate_tool_schema_shape(
92            &format!("Tool \"{}/{}\" input schema", payload.name, tool_name),
93            &tool_input_schema,
94        )?;
95        if let Some(timeout_ms) = tool.timeout_ms {
96            if timeout_ms > MAX_TOOL_TIMEOUT_MS {
97                return Err(ToolRegistrationError::InvalidState(format!(
98                    "Tool \"{}/{}\" timeout is {timeout_ms}ms, max is {MAX_TOOL_TIMEOUT_MS}ms",
99                    payload.name, tool_name
100                )));
101            }
102        }
103        if tool.examples.len() > MAX_TOOL_EXAMPLES_PER_TOOL {
104            return Err(ToolRegistrationError::InvalidState(format!(
105                "Tool \"{}/{}\" defines {} examples, max is {MAX_TOOL_EXAMPLES_PER_TOOL}",
106                payload.name,
107                tool_name,
108                tool.examples.len()
109            )));
110        }
111        for (index, example) in tool.examples.iter().enumerate() {
112            validate_description_length(
113                &format!("Tool \"{}/{}\" example {index}", payload.name, tool_name),
114                &example.description,
115            )?;
116            let example_input: Value = serde_json::from_str(&example.input).map_err(|error| {
117                ToolRegistrationError::InvalidState(format!(
118                    "Tool \"{}/{}\" example {index} input is invalid JSON: {error}",
119                    payload.name, tool_name
120                ))
121            })?;
122            validate_json_byte_length(
123                &format!(
124                    "Tool \"{}/{}\" example {index} input",
125                    payload.name, tool_name
126                ),
127                &example_input,
128                MAX_TOOL_EXAMPLE_INPUT_BYTES,
129            )?;
130        }
131    }
132    Ok(())
133}
134
135pub fn ensure_toolkit_name_available(
136    toolkits: &BTreeMap<String, RegisterHostCallbacksRequest>,
137    toolkit_name: &str,
138) -> Result<(), ToolRegistrationError> {
139    if toolkits.contains_key(toolkit_name) {
140        return Err(ToolRegistrationError::Conflict(format!(
141            "toolkit already registered: {toolkit_name}"
142        )));
143    }
144    Ok(())
145}
146
147pub fn ensure_command_aliases_available(
148    toolkits: &BTreeMap<String, RegisterHostCallbacksRequest>,
149    payload: &RegisterHostCallbacksRequest,
150) -> Result<(), ToolRegistrationError> {
151    let requested_command_aliases = payload.command_aliases.iter().collect::<BTreeSet<_>>();
152    let requested_registry_aliases = payload
153        .registry_command_aliases
154        .iter()
155        .collect::<BTreeSet<_>>();
156    for toolkit in toolkits.values() {
157        for alias in &toolkit.command_aliases {
158            if requested_command_aliases.contains(alias)
159                || requested_registry_aliases.contains(alias)
160            {
161                return Err(ToolRegistrationError::Conflict(format!(
162                    "host callback command alias already registered: {alias}"
163                )));
164            }
165        }
166        for alias in &toolkit.registry_command_aliases {
167            if requested_command_aliases.contains(alias) {
168                return Err(ToolRegistrationError::Conflict(format!(
169                    "host callback command alias already registered: {alias}"
170                )));
171            }
172        }
173    }
174    Ok(())
175}
176
177pub fn ensure_toolkit_registry_capacity(
178    toolkits: &BTreeMap<String, RegisterHostCallbacksRequest>,
179    payload: &RegisterHostCallbacksRequest,
180) -> Result<(), ToolRegistrationError> {
181    if toolkits.len() >= MAX_REGISTERED_TOOLKITS {
182        return Err(ToolRegistrationError::InvalidState(format!(
183            "VM already has {} registered toolkits, max is {MAX_REGISTERED_TOOLKITS}",
184            toolkits.len()
185        )));
186    }
187
188    let registered_tools = toolkits
189        .values()
190        .map(|toolkit| toolkit.callbacks.len())
191        .sum::<usize>();
192    let total_tools = registered_tools
193        .checked_add(payload.callbacks.len())
194        .ok_or_else(|| {
195            ToolRegistrationError::InvalidState(String::from(
196                "registered host callback count overflow",
197            ))
198        })?;
199    if total_tools > MAX_REGISTERED_TOOLS_PER_VM {
200        return Err(ToolRegistrationError::InvalidState(format!(
201            "VM would have {total_tools} registered host callbacks, max is {MAX_REGISTERED_TOOLS_PER_VM}"
202        )));
203    }
204
205    Ok(())
206}
207
208pub fn registered_tool_command_names(
209    toolkits: &BTreeMap<String, RegisterHostCallbacksRequest>,
210) -> Vec<String> {
211    let mut seen = BTreeSet::new();
212    let mut commands = Vec::new();
213    for toolkit in toolkits.values() {
214        for alias in toolkit
215            .registry_command_aliases
216            .iter()
217            .chain(toolkit.command_aliases.iter())
218        {
219            if seen.insert(alias.clone()) {
220                commands.push(alias.clone());
221            }
222        }
223    }
224    commands
225}
226
227fn validate_toolkit_name(name: &str) -> Result<(), ToolRegistrationError> {
228    if name.len() > MAX_TOOLKIT_NAME_LENGTH {
229        return Err(ToolRegistrationError::InvalidState(format!(
230            "invalid toolkit name {name}; max length is {MAX_TOOLKIT_NAME_LENGTH}"
231        )));
232    }
233    if name.is_empty()
234        || !name
235            .chars()
236            .all(|ch| ch.is_ascii_lowercase() || ch.is_ascii_digit() || ch == '-')
237    {
238        return Err(ToolRegistrationError::InvalidState(format!(
239            "invalid toolkit name {name}; expected lowercase alphanumeric characters plus hyphens"
240        )));
241    }
242    Ok(())
243}
244
245fn validate_tool_name(name: &str) -> Result<(), ToolRegistrationError> {
246    if name.len() > MAX_TOOL_NAME_LENGTH {
247        return Err(ToolRegistrationError::InvalidState(format!(
248            "invalid tool name {name}; max length is {MAX_TOOL_NAME_LENGTH}"
249        )));
250    }
251    if name.is_empty()
252        || !name
253            .chars()
254            .all(|ch| ch.is_ascii_lowercase() || ch.is_ascii_digit() || ch == '-')
255    {
256        return Err(ToolRegistrationError::InvalidState(format!(
257            "invalid tool name {name}; expected lowercase alphanumeric characters plus hyphens"
258        )));
259    }
260    Ok(())
261}
262
263fn validate_command_aliases(label: &str, aliases: &[String]) -> Result<(), ToolRegistrationError> {
264    let mut seen = BTreeSet::new();
265    for alias in aliases {
266        validate_command_alias(label, alias)?;
267        if !seen.insert(alias) {
268            return Err(ToolRegistrationError::InvalidState(format!(
269                "duplicate host callback {label}: {alias}"
270            )));
271        }
272    }
273    Ok(())
274}
275
276fn validate_command_alias(label: &str, alias: &str) -> Result<(), ToolRegistrationError> {
277    if alias.is_empty()
278        || alias == "."
279        || alias == ".."
280        || alias.contains('/')
281        || alias.contains('\0')
282    {
283        return Err(ToolRegistrationError::InvalidState(format!(
284            "invalid host callback {label}: {alias:?}"
285        )));
286    }
287    Ok(())
288}
289
290fn validate_description_length(
291    label: &str,
292    description: &str,
293) -> Result<(), ToolRegistrationError> {
294    if description.len() > MAX_TOOL_DESCRIPTION_LENGTH {
295        return Err(ToolRegistrationError::InvalidState(format!(
296            "{label} description is {} characters, max is {MAX_TOOL_DESCRIPTION_LENGTH}",
297            description.len()
298        )));
299    }
300    Ok(())
301}
302
303fn validate_tool_schema_shape(label: &str, schema: &Value) -> Result<(), ToolRegistrationError> {
304    validate_json_byte_length(label, schema, MAX_TOOL_SCHEMA_BYTES)?;
305    validate_json_depth(label, schema, 0)
306}
307
308fn validate_json_byte_length(
309    label: &str,
310    value: &Value,
311    limit: usize,
312) -> Result<(), ToolRegistrationError> {
313    let length = serde_json::to_vec(value)
314        .map_err(|error| {
315            ToolRegistrationError::InvalidState(format!("{label} is invalid JSON: {error}"))
316        })?
317        .len();
318    if length > limit {
319        return Err(ToolRegistrationError::InvalidState(format!(
320            "{label} is {length} bytes, max is {limit}"
321        )));
322    }
323    Ok(())
324}
325
326fn validate_json_depth(
327    label: &str,
328    value: &Value,
329    depth: usize,
330) -> Result<(), ToolRegistrationError> {
331    if depth > MAX_TOOL_SCHEMA_DEPTH {
332        return Err(ToolRegistrationError::InvalidState(format!(
333            "{label} exceeds max JSON depth {MAX_TOOL_SCHEMA_DEPTH}"
334        )));
335    }
336
337    match value {
338        Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => Ok(()),
339        Value::Array(values) => {
340            for value in values {
341                validate_json_depth(label, value, depth + 1)?;
342            }
343            Ok(())
344        }
345        Value::Object(object) => {
346            for value in object.values() {
347                validate_json_depth(label, value, depth + 1)?;
348            }
349            Ok(())
350        }
351    }
352}
agentos_native_sidecar_core/tools.rs

agentos_native_sidecar_core/
tools.rs
