1mod ctx;
14mod error;
15mod registry;
16mod tool;
17mod truncate;
18
19pub use ctx::ToolCtx;
20pub use error::ToolError;
21pub use locode_protocol::{GrammarSyntax, ToolInputFormat, ToolSpec};
22pub use registry::{Dispatched, DynTool, Registry, ToolRunResult};
23pub use tool::{Tool, ToolKind, ToolOutput};
24pub use truncate::{MODEL_OUTPUT_BUDGET, truncate_for_model};
25
26#[cfg(test)]
27mod tests {
28 #![allow(clippy::unnecessary_literal_bound)]
31
32 use super::*;
33 use async_trait::async_trait;
34 use locode_protocol::{ContentBlock, ResultChunk};
35 use schemars::JsonSchema;
36 use serde::{Deserialize, Serialize};
37 use serde_json::{Value, json};
38 use std::path::PathBuf;
39 use tokio_util::sync::CancellationToken;
40
41 #[derive(Debug, Deserialize, JsonSchema)]
44 struct EchoArgs {
45 message: String,
46 }
47
48 #[derive(Debug, Serialize)]
49 struct EchoOut {
50 echoed: String,
51 }
52
53 impl ToolOutput for EchoOut {
54 fn to_prompt_text(&self) -> String {
55 self.echoed.clone()
56 }
57 }
58
59 struct Echo;
60
61 #[async_trait]
62 impl Tool for Echo {
63 type Args = EchoArgs;
64 type Output = EchoOut;
65
66 fn kind(&self) -> ToolKind {
67 ToolKind::Shell
68 }
69 fn description(&self) -> &str {
70 "echoes its input"
71 }
72 async fn run(&self, _ctx: &ToolCtx, args: EchoArgs) -> Result<EchoOut, ToolError> {
73 Ok(EchoOut {
74 echoed: args.message,
75 })
76 }
77 }
78
79 struct Boom;
81
82 #[async_trait]
83 impl Tool for Boom {
84 type Args = Value;
85 type Output = EchoOut;
86
87 fn kind(&self) -> ToolKind {
88 ToolKind::Shell
89 }
90 fn description(&self) -> &str {
91 "always fatal"
92 }
93 async fn run(&self, _ctx: &ToolCtx, _args: Value) -> Result<EchoOut, ToolError> {
94 Err(ToolError::Fatal("unrecoverable".into()))
95 }
96 }
97
98 fn ctx(call_id: &str) -> ToolCtx {
99 ToolCtx::new(
100 PathBuf::from("/repo"),
101 call_id.to_owned(),
102 PathBuf::from("/repo"),
103 CancellationToken::new(),
104 )
105 }
106
107 fn registry_with_echo() -> Registry {
108 let mut reg = Registry::new();
109 reg.register("echo", Echo);
110 reg
111 }
112
113 #[test]
114 fn schema_is_derived_from_args() {
115 let schema = Echo.parameters_schema();
116 assert_eq!(
117 schema["properties"]["message"]["type"],
118 json!("string"),
119 "derived schema should describe the `message: String` arg"
120 );
121 assert_eq!(schema["required"], json!(["message"]));
122 }
123
124 #[tokio::test]
125 async fn echo_round_trips_output_and_prompt_text() {
126 let reg = registry_with_echo();
127 let out = reg
128 .dispatch("echo", json!({ "message": "hi" }), &ctx("c1"))
129 .await;
130
131 let ContentBlock::ToolResult {
133 tool_use_id,
134 content,
135 is_error,
136 } = &out.tool_result
137 else {
138 panic!("expected a tool_result block");
139 };
140 assert_eq!(tool_use_id, "c1");
141 assert!(!is_error);
142 assert_eq!(content, &vec![ResultChunk::Text { text: "hi".into() }]);
143
144 assert!(out.record.ok);
146 assert_eq!(out.record.name, "echo");
147 assert_eq!(out.record.kind, "shell");
148 assert_eq!(out.record.output, json!({ "echoed": "hi" }));
149 assert!(out.fatal.is_none());
150 }
151
152 #[tokio::test]
153 async fn bad_args_are_soft() {
154 let reg = registry_with_echo();
155 let out = reg.dispatch("echo", json!({}), &ctx("c1")).await;
157
158 let ContentBlock::ToolResult { is_error, .. } = &out.tool_result else {
159 panic!("expected a tool_result block");
160 };
161 assert!(is_error, "bad args should be a soft is_error result");
162 assert!(!out.record.ok);
163 assert!(out.fatal.is_none());
164 }
165
166 #[tokio::test]
167 async fn unknown_tool_is_soft() {
168 let reg = registry_with_echo();
169 let out = reg.dispatch("nope", json!({}), &ctx("c1")).await;
170
171 let ContentBlock::ToolResult { is_error, .. } = &out.tool_result else {
172 panic!("expected a tool_result block");
173 };
174 assert!(is_error);
175 assert!(!out.record.ok);
176 assert_eq!(out.record.kind, "other");
177 assert!(out.fatal.is_none());
178 }
179
180 #[tokio::test]
181 async fn fatal_sets_flag_and_still_pairs() {
182 let mut reg = Registry::new();
183 reg.register("boom", Boom);
184 let out = reg.dispatch("boom", json!({}), &ctx("c9")).await;
185
186 let ContentBlock::ToolResult {
188 tool_use_id,
189 is_error,
190 ..
191 } = &out.tool_result
192 else {
193 panic!("expected a tool_result block");
194 };
195 assert_eq!(tool_use_id, "c9");
196 assert!(is_error);
197 assert_eq!(out.fatal.as_deref(), Some("unrecoverable"));
198 }
199
200 #[test]
201 #[should_panic(expected = "duplicate tool registration")]
202 fn duplicate_registration_panics() {
203 let mut reg = Registry::new();
204 reg.register("echo", Echo);
205 reg.register("echo", Echo);
206 }
207
208 struct FakeMcpTool;
211
212 #[async_trait]
213 impl DynTool for FakeMcpTool {
214 fn kind(&self) -> ToolKind {
215 ToolKind::Other
216 }
217 fn description(&self) -> &str {
218 "a dynamically-described tool"
219 }
220 fn parameters_schema(&self) -> Value {
221 json!({ "type": "object", "properties": {} })
222 }
223 async fn call(&self, _ctx: &ToolCtx, raw_args: Value) -> Result<ToolRunResult, ToolError> {
224 Ok(ToolRunResult {
225 output: raw_args.clone(),
226 prompt_text: raw_args.to_string(),
227 })
228 }
229 }
230
231 #[tokio::test]
232 async fn register_dyn_supports_mcp_like_tools() {
233 let mut reg = Registry::new();
234 reg.register_dyn("mcp__thing", Box::new(FakeMcpTool));
235 assert!(reg.contains("mcp__thing"));
236
237 let out = reg
238 .dispatch("mcp__thing", json!({ "x": 1 }), &ctx("m1"))
239 .await;
240 assert!(out.record.ok);
241 assert_eq!(out.record.kind, "other");
242 assert_eq!(out.record.output, json!({ "x": 1 }));
243 }
244
245 #[test]
246 fn tool_kind_key_matches_serde() {
247 for kind in [
248 ToolKind::Shell,
249 ToolKind::Read,
250 ToolKind::Write,
251 ToolKind::Edit,
252 ToolKind::Glob,
253 ToolKind::Grep,
254 ToolKind::Other,
255 ] {
256 assert_eq!(serde_json::to_value(kind).unwrap(), json!(kind.as_str()));
257 }
258 }
259}