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