lib.rs 8.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256
  1. use runtime::{
  2. edit_file, execute_bash, glob_search, grep_search, read_file, write_file, BashCommandInput,
  3. GrepSearchInput,
  4. };
  5. use serde::Deserialize;
  6. use serde_json::{json, Value};
  7. #[derive(Debug, Clone, PartialEq, Eq)]
  8. pub struct ToolManifestEntry {
  9. pub name: String,
  10. pub source: ToolSource,
  11. }
  12. #[derive(Debug, Clone, Copy, PartialEq, Eq)]
  13. pub enum ToolSource {
  14. Base,
  15. Conditional,
  16. }
  17. #[derive(Debug, Clone, Default, PartialEq, Eq)]
  18. pub struct ToolRegistry {
  19. entries: Vec<ToolManifestEntry>,
  20. }
  21. impl ToolRegistry {
  22. #[must_use]
  23. pub fn new(entries: Vec<ToolManifestEntry>) -> Self {
  24. Self { entries }
  25. }
  26. #[must_use]
  27. pub fn entries(&self) -> &[ToolManifestEntry] {
  28. &self.entries
  29. }
  30. }
  31. #[derive(Debug, Clone, PartialEq, Eq)]
  32. pub struct ToolSpec {
  33. pub name: &'static str,
  34. pub description: &'static str,
  35. pub input_schema: Value,
  36. }
  37. #[must_use]
  38. pub fn mvp_tool_specs() -> Vec<ToolSpec> {
  39. vec![
  40. ToolSpec {
  41. name: "bash",
  42. description: "Execute a shell command in the current workspace.",
  43. input_schema: json!({
  44. "type": "object",
  45. "properties": {
  46. "command": { "type": "string" },
  47. "timeout": { "type": "integer", "minimum": 1 },
  48. "description": { "type": "string" },
  49. "run_in_background": { "type": "boolean" },
  50. "dangerouslyDisableSandbox": { "type": "boolean" }
  51. },
  52. "required": ["command"],
  53. "additionalProperties": false
  54. }),
  55. },
  56. ToolSpec {
  57. name: "read_file",
  58. description: "Read a text file from the workspace.",
  59. input_schema: json!({
  60. "type": "object",
  61. "properties": {
  62. "path": { "type": "string" },
  63. "offset": { "type": "integer", "minimum": 0 },
  64. "limit": { "type": "integer", "minimum": 1 }
  65. },
  66. "required": ["path"],
  67. "additionalProperties": false
  68. }),
  69. },
  70. ToolSpec {
  71. name: "write_file",
  72. description: "Write a text file in the workspace.",
  73. input_schema: json!({
  74. "type": "object",
  75. "properties": {
  76. "path": { "type": "string" },
  77. "content": { "type": "string" }
  78. },
  79. "required": ["path", "content"],
  80. "additionalProperties": false
  81. }),
  82. },
  83. ToolSpec {
  84. name: "edit_file",
  85. description: "Replace text in a workspace file.",
  86. input_schema: json!({
  87. "type": "object",
  88. "properties": {
  89. "path": { "type": "string" },
  90. "old_string": { "type": "string" },
  91. "new_string": { "type": "string" },
  92. "replace_all": { "type": "boolean" }
  93. },
  94. "required": ["path", "old_string", "new_string"],
  95. "additionalProperties": false
  96. }),
  97. },
  98. ToolSpec {
  99. name: "glob_search",
  100. description: "Find files by glob pattern.",
  101. input_schema: json!({
  102. "type": "object",
  103. "properties": {
  104. "pattern": { "type": "string" },
  105. "path": { "type": "string" }
  106. },
  107. "required": ["pattern"],
  108. "additionalProperties": false
  109. }),
  110. },
  111. ToolSpec {
  112. name: "grep_search",
  113. description: "Search file contents with a regex pattern.",
  114. input_schema: json!({
  115. "type": "object",
  116. "properties": {
  117. "pattern": { "type": "string" },
  118. "path": { "type": "string" },
  119. "glob": { "type": "string" },
  120. "output_mode": { "type": "string" },
  121. "-B": { "type": "integer", "minimum": 0 },
  122. "-A": { "type": "integer", "minimum": 0 },
  123. "-C": { "type": "integer", "minimum": 0 },
  124. "context": { "type": "integer", "minimum": 0 },
  125. "-n": { "type": "boolean" },
  126. "-i": { "type": "boolean" },
  127. "type": { "type": "string" },
  128. "head_limit": { "type": "integer", "minimum": 1 },
  129. "offset": { "type": "integer", "minimum": 0 },
  130. "multiline": { "type": "boolean" }
  131. },
  132. "required": ["pattern"],
  133. "additionalProperties": false
  134. }),
  135. },
  136. ]
  137. }
  138. pub fn execute_tool(name: &str, input: &Value) -> Result<String, String> {
  139. match name {
  140. "bash" => from_value::<BashCommandInput>(input).and_then(run_bash),
  141. "read_file" => from_value::<ReadFileInput>(input).and_then(|input| run_read_file(&input)),
  142. "write_file" => {
  143. from_value::<WriteFileInput>(input).and_then(|input| run_write_file(&input))
  144. }
  145. "edit_file" => from_value::<EditFileInput>(input).and_then(|input| run_edit_file(&input)),
  146. "glob_search" => {
  147. from_value::<GlobSearchInputValue>(input).and_then(|input| run_glob_search(&input))
  148. }
  149. "grep_search" => {
  150. from_value::<GrepSearchInput>(input).and_then(|input| run_grep_search(&input))
  151. }
  152. _ => Err(format!("unsupported tool: {name}")),
  153. }
  154. }
  155. fn from_value<T: for<'de> Deserialize<'de>>(input: &Value) -> Result<T, String> {
  156. serde_json::from_value(input.clone()).map_err(|error| error.to_string())
  157. }
  158. fn run_bash(input: BashCommandInput) -> Result<String, String> {
  159. serde_json::to_string_pretty(&execute_bash(input).map_err(|error| error.to_string())?)
  160. .map_err(|error| error.to_string())
  161. }
  162. fn run_read_file(input: &ReadFileInput) -> Result<String, String> {
  163. to_pretty_json(
  164. read_file(&input.path, input.offset, input.limit).map_err(|error| error.to_string())?,
  165. )
  166. }
  167. fn run_write_file(input: &WriteFileInput) -> Result<String, String> {
  168. to_pretty_json(write_file(&input.path, &input.content).map_err(|error| error.to_string())?)
  169. }
  170. fn run_edit_file(input: &EditFileInput) -> Result<String, String> {
  171. to_pretty_json(
  172. edit_file(
  173. &input.path,
  174. &input.old_string,
  175. &input.new_string,
  176. input.replace_all.unwrap_or(false),
  177. )
  178. .map_err(|error| error.to_string())?,
  179. )
  180. }
  181. fn run_glob_search(input: &GlobSearchInputValue) -> Result<String, String> {
  182. to_pretty_json(
  183. glob_search(&input.pattern, input.path.as_deref()).map_err(|error| error.to_string())?,
  184. )
  185. }
  186. fn run_grep_search(input: &GrepSearchInput) -> Result<String, String> {
  187. to_pretty_json(grep_search(input).map_err(|error| error.to_string())?)
  188. }
  189. fn to_pretty_json<T: serde::Serialize>(value: T) -> Result<String, String> {
  190. serde_json::to_string_pretty(&value).map_err(|error| error.to_string())
  191. }
  192. #[derive(Debug, Deserialize)]
  193. struct ReadFileInput {
  194. path: String,
  195. offset: Option<usize>,
  196. limit: Option<usize>,
  197. }
  198. #[derive(Debug, Deserialize)]
  199. struct WriteFileInput {
  200. path: String,
  201. content: String,
  202. }
  203. #[derive(Debug, Deserialize)]
  204. struct EditFileInput {
  205. path: String,
  206. old_string: String,
  207. new_string: String,
  208. replace_all: Option<bool>,
  209. }
  210. #[derive(Debug, Deserialize)]
  211. struct GlobSearchInputValue {
  212. pattern: String,
  213. path: Option<String>,
  214. }
  215. #[cfg(test)]
  216. mod tests {
  217. use super::{execute_tool, mvp_tool_specs};
  218. use serde_json::json;
  219. #[test]
  220. fn exposes_mvp_tools() {
  221. let names = mvp_tool_specs()
  222. .into_iter()
  223. .map(|spec| spec.name)
  224. .collect::<Vec<_>>();
  225. assert!(names.contains(&"bash"));
  226. assert!(names.contains(&"read_file"));
  227. }
  228. #[test]
  229. fn rejects_unknown_tool_names() {
  230. let error = execute_tool("nope", &json!({})).expect_err("tool should be rejected");
  231. assert!(error.contains("unsupported tool"));
  232. }
  233. }