1use async_trait::async_trait;
7use serde_json::Value;
8use std::collections::BTreeMap;
9use std::sync::Arc;
10use tracing::Instrument;
11
12use crate::error::{Error, Result};
13use crate::llm::ToolSpec;
14
15pub mod apply_patch;
16pub mod estimate_tokens;
17pub mod fs;
18pub mod load_skill;
19pub mod memory;
20pub mod run_background;
21pub mod run_skill_script;
22pub mod search;
23pub mod shell;
24pub mod sub_agent;
25pub mod transport;
26#[cfg(feature = "web_fetch")]
27pub mod web_fetch;
28
29pub use apply_patch::ApplyPatch;
30pub use estimate_tokens::EstimateTokens;
31pub use fs::{ListDir, ReadFile, WriteFile};
32pub use load_skill::LoadSkill;
33pub use memory::{Forget, Recall, Remember};
34pub use run_background::{BackgroundJobManager, CheckBackground, RunBackground};
35pub use run_skill_script::RunSkillScript;
36pub use search::SearchFiles;
37pub use shell::RunShell;
38pub use sub_agent::SubAgent;
39pub use transport::{DirEntry, ExecResult, LocalTransport, ReadResult, ToolTransport};
40#[cfg(feature = "web_fetch")]
41pub use web_fetch::WebFetch;
42
43impl Default for ToolRegistry {
44 fn default() -> Self {
45 Self::local()
46 }
47}
48
49#[async_trait]
50pub trait Tool: Send + Sync {
51 fn spec(&self) -> ToolSpec;
52 async fn execute(&self, arguments: Value) -> Result<String>;
53 fn is_readonly(&self) -> bool {
56 false
57 }
58}
59
60#[derive(Clone)]
61pub struct ToolRegistry {
62 tools: BTreeMap<String, Arc<dyn Tool>>,
63 transport: Arc<dyn ToolTransport>,
64}
65
66impl ToolRegistry {
67 pub fn new(transport: Arc<dyn ToolTransport>) -> Self {
68 Self {
69 tools: BTreeMap::new(),
70 transport,
71 }
72 }
73
74 pub fn local() -> Self {
76 Self::new(Arc::new(LocalTransport))
77 }
78
79 pub fn transport(&self) -> &Arc<dyn ToolTransport> {
81 &self.transport
82 }
83
84 pub fn with_same_transport(&self) -> Self {
86 Self {
87 tools: BTreeMap::new(),
88 transport: self.transport.clone(),
89 }
90 }
91
92 pub fn register(mut self, tool: Arc<dyn Tool>) -> Self {
93 let name = tool.spec().name;
94 self.tools.insert(name, tool);
95 self
96 }
97
98 pub fn register_mut(&mut self, tool: Arc<dyn Tool>) {
100 let name = tool.spec().name;
101 self.tools.insert(name, tool);
102 }
103
104 pub fn specs(&self) -> Vec<ToolSpec> {
105 self.tools.values().map(|t| t.spec()).collect()
106 }
107
108 pub fn names(&self) -> Vec<String> {
109 self.tools.keys().cloned().collect()
110 }
111
112 pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
113 self.tools.get(name).cloned()
114 }
115
116 pub fn is_readonly(&self, name: &str) -> bool {
118 self.tools
119 .get(name)
120 .map(|t| t.is_readonly())
121 .unwrap_or(false)
122 }
123
124 pub async fn invoke(&self, name: &str, arguments: Value) -> Result<String> {
125 let args_size = arguments.to_string().len();
126 let span = tracing::info_span!("tool.execute", name = %name, args_size);
127 async move {
128 let tool = self
129 .get(name)
130 .ok_or_else(|| Error::UnknownTool(name.into()))?;
131 tool.execute(arguments).await.map_err(|e| match e {
132 Error::Tool { .. } | Error::BadToolArgs { .. } | Error::UnknownTool(_) => e,
133 other => Error::Tool {
134 name: name.into(),
135 message: other.to_string(),
136 },
137 })
138 }
139 .instrument(span)
140 .await
141 }
142}
143
144pub(crate) fn resolve_within(root: &std::path::Path, path: &str) -> Result<std::path::PathBuf> {
150 let candidate = std::path::Path::new(path);
151 let joined = if candidate.is_absolute() {
152 candidate.to_path_buf()
153 } else {
154 root.join(candidate)
155 };
156 let abs_root = absolutise(root);
157 let abs_joined = absolutise(&joined);
158 if !abs_joined.starts_with(&abs_root) {
159 return Err(Error::BadToolArgs {
160 name: "<fs>".into(),
161 message: format!(
162 "path `{}` escapes workspace root `{}`",
163 path,
164 abs_root.display()
165 ),
166 });
167 }
168 Ok(abs_joined)
169}
170
171fn absolutise(p: &std::path::Path) -> std::path::PathBuf {
174 let abs = if p.is_absolute() {
175 p.to_path_buf()
176 } else {
177 std::env::current_dir()
178 .unwrap_or_else(|_| std::path::PathBuf::from("."))
179 .join(p)
180 };
181 normalise(&abs)
182}
183
184fn normalise(p: &std::path::Path) -> std::path::PathBuf {
185 let mut out = std::path::PathBuf::new();
186 for c in p.components() {
187 use std::path::Component::*;
188 match c {
189 ParentDir => {
190 out.pop();
191 }
192 CurDir => {}
193 other => out.push(other.as_os_str()),
194 }
195 }
196 out
197}
198
199#[cfg(test)]
200mod tests {
201 use super::*;
202 use async_trait::async_trait;
203
204 struct Echo;
205
206 #[async_trait]
207 impl Tool for Echo {
208 fn spec(&self) -> ToolSpec {
209 ToolSpec {
210 name: "echo".into(),
211 description: "echo".into(),
212 parameters: serde_json::json!({"type":"object","properties":{"msg":{"type":"string"}}}),
213 }
214 }
215 async fn execute(&self, args: Value) -> Result<String> {
216 Ok(args["msg"].as_str().unwrap_or("").into())
217 }
218 }
219
220 #[tokio::test]
221 async fn registry_dispatches_and_errors_on_unknown() {
222 let reg = ToolRegistry::local().register(Arc::new(Echo));
223 let out = reg
224 .invoke("echo", serde_json::json!({"msg":"hi"}))
225 .await
226 .unwrap();
227 assert_eq!(out, "hi");
228 let err = reg.invoke("nope", serde_json::json!({})).await.unwrap_err();
229 assert!(matches!(err, Error::UnknownTool(_)));
230 }
231
232 #[test]
233 fn resolve_within_rejects_escape() {
234 let root = std::path::Path::new("/work");
235 assert!(resolve_within(root, "../etc/passwd").is_err());
236 assert!(resolve_within(root, "/elsewhere").is_err());
237 assert!(resolve_within(root, "src/lib.rs").is_ok());
238 }
239
240 #[test]
241 fn resolve_within_handles_relative_root() {
242 let cwd = std::env::current_dir().unwrap();
244 let resolved = resolve_within(std::path::Path::new("."), "src/lib.rs").unwrap();
245 assert!(resolved.starts_with(&cwd));
246 assert!(resolved.ends_with("src/lib.rs"));
247 }
248}