Skip to main content

rskit_tool/
io.rs

1//! Typed tool IO wrappers.
2
3use std::collections::HashMap;
4use std::ops::Index;
5
6use rskit_errors::{AppError, AppResult, ErrorCode};
7use rskit_schema::Json;
8use serde::{Deserialize, Deserializer, Serialize, de};
9
10/// Machine-validatable JSON Schema for tool input or output.
11#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
12#[serde(transparent)]
13pub struct ToolSchema(Json);
14
15impl ToolSchema {
16    /// Create a schema from a JSON object value.
17    pub fn new(value: Json) -> AppResult<Self> {
18        if value.is_object() {
19            Ok(Self(value))
20        } else {
21            Err(AppError::new(
22                ErrorCode::InvalidInput,
23                "tool schema must be a JSON object",
24            ))
25        }
26    }
27
28    /// Create the permissive object schema used for tools without typed params.
29    #[must_use]
30    pub fn any_object() -> Self {
31        Self(serde_json::json!({"type": "object"}))
32    }
33
34    /// Borrow the underlying JSON Schema document.
35    #[must_use]
36    pub const fn as_json(&self) -> &Json {
37        &self.0
38    }
39
40    /// Borrow this schema as a JSON object.
41    #[must_use]
42    pub fn as_object(&self) -> Option<&serde_json::Map<String, Json>> {
43        self.0.as_object()
44    }
45
46    /// Return whether this schema is represented by a JSON object.
47    #[must_use]
48    pub fn is_object(&self) -> bool {
49        self.0.is_object()
50    }
51
52    /// Consume this wrapper and return the JSON Schema document.
53    #[must_use]
54    pub fn into_json(self) -> Json {
55        self.0
56    }
57}
58
59impl Default for ToolSchema {
60    fn default() -> Self {
61        Self::any_object()
62    }
63}
64
65impl TryFrom<Json> for ToolSchema {
66    type Error = AppError;
67
68    fn try_from(value: Json) -> Result<Self, Self::Error> {
69        Self::new(value)
70    }
71}
72
73impl<'de> Deserialize<'de> for ToolSchema {
74    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
75    where
76        D: Deserializer<'de>,
77    {
78        Self::new(Json::deserialize(deserializer)?).map_err(de::Error::custom)
79    }
80}
81
82/// Validated tool invocation input.
83#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
84#[serde(transparent)]
85pub struct ToolInput(Json);
86
87impl ToolInput {
88    /// Create tool input from a JSON object value.
89    pub fn new(value: Json) -> AppResult<Self> {
90        if value.is_object() {
91            Ok(Self(value))
92        } else {
93            Err(AppError::new(
94                ErrorCode::InvalidInput,
95                "tool input must be a JSON object",
96            ))
97        }
98    }
99
100    /// Create empty object input.
101    #[must_use]
102    pub fn empty() -> Self {
103        Self(Json::Object(serde_json::Map::new()))
104    }
105
106    /// Create tool input from a JSON object map.
107    #[must_use]
108    pub const fn from_object(object: serde_json::Map<String, Json>) -> Self {
109        Self(Json::Object(object))
110    }
111
112    /// Borrow the JSON object payload.
113    #[must_use]
114    pub const fn as_json(&self) -> &Json {
115        &self.0
116    }
117
118    /// Consume this wrapper and return the JSON object payload.
119    #[must_use]
120    pub fn into_json(self) -> Json {
121        self.0
122    }
123}
124
125impl Default for ToolInput {
126    fn default() -> Self {
127        Self::empty()
128    }
129}
130
131impl TryFrom<Json> for ToolInput {
132    type Error = AppError;
133
134    fn try_from(value: Json) -> Result<Self, Self::Error> {
135        Self::new(value)
136    }
137}
138
139impl<'de> Deserialize<'de> for ToolInput {
140    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
141    where
142        D: Deserializer<'de>,
143    {
144        Self::new(Json::deserialize(deserializer)?).map_err(de::Error::custom)
145    }
146}
147
148/// Structured tool output or metadata value.
149#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
150#[serde(transparent)]
151pub struct ToolOutput(Json);
152
153impl ToolOutput {
154    /// Create structured output from a serializable value.
155    pub fn from_serializable<T: Serialize>(value: &T) -> AppResult<Self> {
156        serde_json::to_value(value).map(Self).map_err(|err| {
157            AppError::new(
158                ErrorCode::Internal,
159                format!("failed to serialize tool output: {err}"),
160            )
161            .with_cause(err)
162        })
163    }
164
165    /// Borrow the structured JSON output.
166    #[must_use]
167    pub const fn as_json(&self) -> &Json {
168        &self.0
169    }
170
171    /// Consume this wrapper and return the structured JSON output.
172    #[must_use]
173    pub fn into_json(self) -> Json {
174        self.0
175    }
176}
177
178impl From<Json> for ToolOutput {
179    fn from(value: Json) -> Self {
180        Self(value)
181    }
182}
183
184impl PartialEq<Json> for ToolOutput {
185    fn eq(&self, other: &Json) -> bool {
186        &self.0 == other
187    }
188}
189
190impl Index<&str> for ToolOutput {
191    type Output = Json;
192
193    fn index(&self, index: &str) -> &Self::Output {
194        &self.0[index]
195    }
196}
197
198/// Metadata attached to a tool context or result.
199pub type ToolMetadata = HashMap<String, ToolOutput>;
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204    use serde::ser::{Error as _, SerializeStruct};
205
206    #[test]
207    fn schema_deserialization_rejects_non_object() {
208        let err = serde_json::from_str::<ToolSchema>(r#""not-a-schema""#).unwrap_err();
209        assert!(
210            err.to_string()
211                .contains("tool schema must be a JSON object")
212        );
213    }
214
215    #[test]
216    fn input_deserialization_rejects_non_object() {
217        let err = serde_json::from_str::<ToolInput>("[]").unwrap_err();
218        assert!(err.to_string().contains("tool input must be a JSON object"));
219    }
220
221    #[test]
222    fn output_serialization_error_has_context() {
223        struct FailingOutput;
224
225        impl Serialize for FailingOutput {
226            fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
227            where
228                S: serde::Serializer,
229            {
230                let mut state = serializer.serialize_struct("FailingOutput", 1)?;
231                state.serialize_field("bad", &FailingField)?;
232                state.end()
233            }
234        }
235
236        struct FailingField;
237
238        impl Serialize for FailingField {
239            fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
240            where
241                S: serde::Serializer,
242            {
243                Err(S::Error::custom("boom"))
244            }
245        }
246
247        let err = ToolOutput::from_serializable(&FailingOutput).unwrap_err();
248        assert!(
249            err.message()
250                .contains("failed to serialize tool output: boom")
251        );
252    }
253
254    #[test]
255    fn schema_and_input_helpers_preserve_json_object_contract() {
256        let schema = ToolSchema::any_object();
257        assert!(schema.is_object());
258        assert_eq!(schema.as_json()["type"], "object");
259        assert_eq!(
260            ToolSchema::default().into_json(),
261            serde_json::json!({"type": "object"})
262        );
263        assert_eq!(
264            ToolSchema::try_from(serde_json::json!("string"))
265                .unwrap_err()
266                .code(),
267            ErrorCode::InvalidInput
268        );
269
270        let mut object = serde_json::Map::new();
271        object.insert("query".to_owned(), serde_json::json!("rust"));
272        let input = ToolInput::from_object(object);
273        assert_eq!(input.as_json()["query"], "rust");
274        assert_eq!(ToolInput::empty().into_json(), serde_json::json!({}));
275        assert_eq!(
276            ToolInput::try_from(serde_json::json!(["not", "object"]))
277                .unwrap_err()
278                .code(),
279            ErrorCode::InvalidInput
280        );
281    }
282
283    #[test]
284    fn output_helpers_expose_and_consume_json() {
285        let output = ToolOutput::from(serde_json::json!({"ok": true}));
286        assert_eq!(output.as_json()["ok"], true);
287        assert_eq!(output["ok"], true);
288        assert_eq!(output.into_json(), serde_json::json!({"ok": true}));
289    }
290}