Skip to main content

tea_tools/
invocation.rs

1use std::sync::Arc;
2
3use serde_json::Value;
4use tea_protocol::{ProtocolMetadata, ToolCallId};
5use thiserror::Error;
6
7use crate::{SchedulerClass, ToolName, ToolResource, ToolSource, ToolSpec};
8
9/// Untrusted complete tool invocation before registry validation.
10#[derive(Debug, Clone, PartialEq)]
11pub struct ToolInvocation {
12    tool_call_id: ToolCallId,
13    name: ToolName,
14    arguments: Value,
15    metadata: ProtocolMetadata,
16}
17
18impl ToolInvocation {
19    /// Creates a bounded invocation with object arguments.
20    ///
21    /// # Errors
22    ///
23    /// Returns an error for non-object, oversized, or deeply nested arguments.
24    pub fn new(
25        tool_call_id: ToolCallId,
26        name: ToolName,
27        arguments: Value,
28        metadata: ProtocolMetadata,
29    ) -> Result<Self, ToolInvocationError> {
30        if !arguments.is_object() {
31            return Err(ToolInvocationError::ArgumentsMustBeObject);
32        }
33        if serde_json::to_vec(&arguments)
34            .map_err(|_| ToolInvocationError::ArgumentsOutOfBounds)?
35            .len()
36            > 256 * 1024
37            || json_depth(&arguments) > 32
38        {
39            return Err(ToolInvocationError::ArgumentsOutOfBounds);
40        }
41        Ok(Self {
42            tool_call_id,
43            name,
44            arguments,
45            metadata,
46        })
47    }
48
49    /// Returns canonical tool-call ID.
50    #[must_use]
51    pub const fn tool_call_id(&self) -> &ToolCallId {
52        &self.tool_call_id
53    }
54    /// Returns requested tool name.
55    #[must_use]
56    pub const fn name(&self) -> &ToolName {
57        &self.name
58    }
59    /// Returns untrusted object arguments.
60    #[must_use]
61    pub const fn arguments(&self) -> &Value {
62        &self.arguments
63    }
64    /// Returns bounded invocation metadata.
65    #[must_use]
66    pub const fn metadata(&self) -> &ProtocolMetadata {
67        &self.metadata
68    }
69}
70
71/// Invocation proven valid against a registered tool schema and resource resolver.
72#[derive(Debug, Clone)]
73pub struct ValidatedToolInvocation {
74    invocation: ToolInvocation,
75    spec: Arc<ToolSpec>,
76    source: ToolSource,
77    resources: Vec<ToolResource>,
78}
79
80impl ValidatedToolInvocation {
81    #[cfg(feature = "execution")]
82    pub(crate) fn new(
83        invocation: ToolInvocation,
84        spec: Arc<ToolSpec>,
85        resources: Vec<ToolResource>,
86    ) -> Self {
87        let source = spec.source().clone();
88        Self {
89            invocation,
90            spec,
91            source,
92            resources,
93        }
94    }
95
96    /// Returns canonical tool-call ID.
97    #[must_use]
98    pub const fn tool_call_id(&self) -> &ToolCallId {
99        self.invocation.tool_call_id()
100    }
101    /// Returns registered tool name.
102    #[must_use]
103    pub const fn name(&self) -> &ToolName {
104        self.invocation.name()
105    }
106    /// Returns schema-validated object arguments.
107    #[must_use]
108    pub const fn arguments(&self) -> &Value {
109        self.invocation.arguments()
110    }
111    /// Returns bounded invocation metadata.
112    #[must_use]
113    pub const fn metadata(&self) -> &ProtocolMetadata {
114        self.invocation.metadata()
115    }
116    /// Returns registered tool specification.
117    #[must_use]
118    pub fn spec(&self) -> &ToolSpec {
119        &self.spec
120    }
121    /// Returns provenance frozen when registry validation succeeded.
122    #[must_use]
123    pub const fn source(&self) -> &ToolSource {
124        &self.source
125    }
126    /// Returns sorted resolved resources.
127    #[must_use]
128    pub fn resources(&self) -> &[ToolResource] {
129        &self.resources
130    }
131    /// Returns metadata-derived scheduler class.
132    #[must_use]
133    pub fn scheduler_class(&self) -> SchedulerClass {
134        self.spec.scheduler_class()
135    }
136}
137
138/// Error constructing an untrusted invocation.
139#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
140pub enum ToolInvocationError {
141    /// Arguments must be a JSON object.
142    #[error("tool arguments must be a JSON object")]
143    ArgumentsMustBeObject,
144    /// Arguments exceed byte or nesting bounds.
145    #[error("tool arguments exceed supported bounds")]
146    ArgumentsOutOfBounds,
147}
148
149fn json_depth(value: &Value) -> usize {
150    match value {
151        Value::Array(values) => 1 + values.iter().map(json_depth).max().unwrap_or(0),
152        Value::Object(values) => 1 + values.values().map(json_depth).max().unwrap_or(0),
153        _ => 1,
154    }
155}