1use std::collections::HashMap;
7use std::path::{Path, PathBuf};
8use std::sync::Arc;
9use std::time::Duration;
10
11#[cfg(unix)]
12#[allow(unused_imports)]
13use std::os::unix::process::CommandExt;
14
15use glob::glob as glob_match;
16use regex::{Regex, RegexBuilder};
17use serde_json::Value;
18use tokio::io::AsyncReadExt;
19use tokio::process::Command;
20use tracing::{debug, warn};
21
22use crate::error::{AgentError, Result};
23use crate::types::tools::{
24 BashInput, FileEditInput, FileReadInput, FileWriteInput, GlobInput, GrepInput, GrepOutputMode,
25};
26
27#[derive(Debug, Clone)]
29pub struct ToolResult {
30 pub content: String,
32 pub is_error: bool,
34 pub raw_content: Option<serde_json::Value>,
37}
38
39impl ToolResult {
40 fn ok(content: String) -> Self {
42 Self {
43 content,
44 is_error: false,
45 raw_content: None,
46 }
47 }
48
49 fn err(content: String) -> Self {
51 Self {
52 content,
53 is_error: true,
54 raw_content: None,
55 }
56 }
57}
58
59async fn drain_pipe<R: tokio::io::AsyncRead + Unpin>(handle: Option<R>) -> Vec<u8> {
63 let mut buf = Vec::new();
64 let Some(mut reader) = handle else {
65 return buf;
66 };
67 let mut chunk = [0u8; 65536];
68 loop {
69 match tokio::time::timeout(Duration::from_millis(10), reader.read(&mut chunk)).await {
70 Ok(Ok(0)) => break, Ok(Ok(n)) => buf.extend_from_slice(&chunk[..n]),
72 Ok(Err(_)) => break, Err(_) => break, }
75 }
76 buf
77}
78
79pub struct ToolExecutor {
81 cwd: PathBuf,
83 boundary: Option<PathBoundary>,
88 env_blocklist: Vec<String>,
91 env_inject: HashMap<String, String>,
94 #[cfg(unix)]
97 pre_exec: Option<Arc<dyn Fn() -> std::io::Result<()> + Send + Sync>>,
98}
99
100struct PathBoundary {
109 allowed: Vec<PathBuf>,
111}
112
113impl PathBoundary {
114 fn new(cwd: &Path, additional: &[PathBuf]) -> Self {
116 let mut allowed = Vec::with_capacity(1 + additional.len());
117
118 let push_canon = |dirs: &mut Vec<PathBuf>, p: &Path| {
121 dirs.push(p.canonicalize().unwrap_or_else(|_| p.to_path_buf()));
122 };
123
124 push_canon(&mut allowed, cwd);
125 for dir in additional {
126 push_canon(&mut allowed, dir);
127 }
128
129 Self { allowed }
130 }
131
132 fn check(&self, path: &Path) -> std::result::Result<(), ToolResult> {
140 let normalized = Self::normalize(path)?;
141
142 for allowed in &self.allowed {
143 if normalized.starts_with(allowed) {
144 return Ok(());
145 }
146 }
147
148 Err(ToolResult::err(format!(
149 "Access denied: {} is outside the allowed directories",
150 path.display()
151 )))
152 }
153
154 fn normalize(path: &Path) -> std::result::Result<PathBuf, ToolResult> {
161 if let Ok(canon) = path.canonicalize() {
163 return Ok(canon);
164 }
165
166 let mut remaining = Vec::new();
168 let mut ancestor = path.to_path_buf();
169
170 loop {
171 if ancestor.exists() {
172 let base = ancestor.canonicalize().map_err(|_| {
173 ToolResult::err(format!("Access denied: cannot resolve {}", path.display()))
174 })?;
175
176 let mut result = base;
178 for component in remaining.iter().rev() {
179 result = result.join(component);
180 }
181 return Ok(result);
182 }
183
184 match ancestor.file_name() {
185 Some(name) => {
186 let name = name.to_os_string();
187 remaining.push(name);
188 if !ancestor.pop() {
189 break;
190 }
191 }
192 None => break,
193 }
194 }
195
196 Err(ToolResult::err(format!(
198 "Access denied: cannot resolve {}",
199 path.display()
200 )))
201 }
202}
203
204impl ToolExecutor {
205 pub fn new(cwd: PathBuf) -> Self {
207 Self {
208 cwd,
209 boundary: None,
210 env_blocklist: Vec::new(),
211 env_inject: HashMap::new(),
212 #[cfg(unix)]
213 pre_exec: None,
214 }
215 }
216
217 pub fn with_allowed_dirs(cwd: PathBuf, additional: Vec<PathBuf>) -> Self {
220 let boundary = PathBoundary::new(&cwd, &additional);
221 Self {
222 cwd,
223 boundary: Some(boundary),
224 env_blocklist: Vec::new(),
225 env_inject: HashMap::new(),
226 #[cfg(unix)]
227 pre_exec: None,
228 }
229 }
230
231 pub fn with_env_blocklist(mut self, blocklist: Vec<String>) -> Self {
233 self.env_blocklist = blocklist;
234 self
235 }
236
237 pub fn with_env_inject(mut self, env: HashMap<String, String>) -> Self {
240 self.env_inject = env;
241 self
242 }
243
244 #[cfg(unix)]
247 pub fn with_pre_exec(mut self, f: Box<dyn Fn() -> std::io::Result<()> + Send + Sync>) -> Self {
248 self.pre_exec = Some(Arc::from(f));
249 self
250 }
251
252 pub async fn execute(&self, tool_name: &str, input: Value) -> Result<ToolResult> {
254 debug!(tool = tool_name, "executing built-in tool");
255
256 match tool_name {
257 "Read" => {
258 let params: FileReadInput = serde_json::from_value(input)?;
259 self.execute_read(¶ms).await
260 }
261 "Write" => {
262 let params: FileWriteInput = serde_json::from_value(input)?;
263 self.execute_write(¶ms).await
264 }
265 "Edit" => {
266 let params: FileEditInput = serde_json::from_value(input)?;
267 self.execute_edit(¶ms).await
268 }
269 "Bash" => {
270 let params: BashInput = serde_json::from_value(input)?;
271 self.execute_bash(¶ms).await
272 }
273 "Glob" => {
274 let params: GlobInput = serde_json::from_value(input)?;
275 self.execute_glob(¶ms).await
276 }
277 "Grep" => {
278 let params: GrepInput = serde_json::from_value(input)?;
279 self.execute_grep(¶ms).await
280 }
281 _ => Err(AgentError::ToolExecution(format!(
282 "unsupported built-in tool: {tool_name}. Supported built-in tools: Read, Write, Edit, Bash, Glob, Grep"
283 ))),
284 }
285 }
286
287 fn resolve_and_check(&self, path: &str) -> std::result::Result<PathBuf, ToolResult> {
292 let p = Path::new(path);
293 let resolved = if p.is_absolute() {
294 p.to_path_buf()
295 } else {
296 self.cwd.join(p)
297 };
298
299 if let Some(ref boundary) = self.boundary {
300 boundary.check(&resolved)?;
301 }
302
303 Ok(resolved)
304 }
305
306 async fn execute_read(&self, input: &FileReadInput) -> Result<ToolResult> {
313 let path = match self.resolve_and_check(&input.file_path) {
314 Ok(p) => p,
315 Err(denied) => return Ok(denied),
316 };
317
318 let ext = path
320 .extension()
321 .unwrap_or_default()
322 .to_string_lossy()
323 .to_lowercase();
324
325 let media_type = match ext.as_str() {
326 "png" => Some("image/png"),
327 "jpg" | "jpeg" => Some("image/jpeg"),
328 "gif" => Some("image/gif"),
329 "webp" => Some("image/webp"),
330 _ => None,
331 };
332
333 if let Some(media_type) = media_type {
334 return self.read_image(&path, media_type).await;
335 }
336
337 let content = match tokio::fs::read_to_string(&path).await {
338 Ok(c) => c,
339 Err(e) => {
340 return Ok(ToolResult::err(format!(
341 "Failed to read {}: {e}",
342 path.display()
343 )));
344 }
345 };
346
347 let lines: Vec<&str> = content.lines().collect();
348 let total = lines.len();
349
350 let offset = input.offset.unwrap_or(0) as usize;
351 let limit = input.limit.unwrap_or(total as u64) as usize;
352
353 if offset >= total {
354 return Ok(ToolResult::ok(String::new()));
355 }
356
357 let end = (offset + limit).min(total);
358 let selected = &lines[offset..end];
359
360 let width = format!("{}", end).len();
362 let mut output = String::new();
363 for (i, line) in selected.iter().enumerate() {
364 let line_no = offset + i + 1; output.push_str(&format!("{line_no:>width$}\t{line}\n", width = width));
366 }
367
368 Ok(ToolResult::ok(output))
369 }
370
371 async fn read_image(&self, path: &Path, media_type: &str) -> Result<ToolResult> {
373 let bytes = match tokio::fs::read(path).await {
374 Ok(b) => b,
375 Err(e) => {
376 return Ok(ToolResult::err(format!(
377 "Failed to read {}: {e}",
378 path.display()
379 )));
380 }
381 };
382
383 if bytes.len() > 20 * 1024 * 1024 {
385 return Ok(ToolResult::err(format!(
386 "Image too large ({:.1} MB, max 20 MB): {}",
387 bytes.len() as f64 / (1024.0 * 1024.0),
388 path.display()
389 )));
390 }
391
392 use base64::Engine;
393 let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes);
394
395 Ok(ToolResult {
396 content: format!("Image: {}", path.display()),
397 is_error: false,
398 raw_content: Some(serde_json::json!([
399 {
400 "type": "image",
401 "source": {
402 "type": "base64",
403 "media_type": media_type,
404 "data": b64,
405 }
406 }
407 ])),
408 })
409 }
410
411 async fn execute_write(&self, input: &FileWriteInput) -> Result<ToolResult> {
415 let path = match self.resolve_and_check(&input.file_path) {
416 Ok(p) => p,
417 Err(denied) => return Ok(denied),
418 };
419
420 if let Some(parent) = path.parent() {
422 if let Err(e) = tokio::fs::create_dir_all(parent).await {
423 return Ok(ToolResult::err(format!(
424 "Failed to create directories for {}: {e}",
425 path.display()
426 )));
427 }
428 }
429
430 match tokio::fs::write(&path, &input.content).await {
431 Ok(()) => Ok(ToolResult::ok(format!(
432 "Successfully wrote to {}",
433 path.display()
434 ))),
435 Err(e) => Ok(ToolResult::err(format!(
436 "Failed to write {}: {e}",
437 path.display()
438 ))),
439 }
440 }
441
442 async fn execute_edit(&self, input: &FileEditInput) -> Result<ToolResult> {
450 let path = match self.resolve_and_check(&input.file_path) {
451 Ok(p) => p,
452 Err(denied) => return Ok(denied),
453 };
454
455 let content = match tokio::fs::read_to_string(&path).await {
456 Ok(c) => c,
457 Err(e) => {
458 return Ok(ToolResult::err(format!(
459 "Failed to read {}: {e}",
460 path.display()
461 )));
462 }
463 };
464
465 let replace_all = input.replace_all.unwrap_or(false);
466 let count = content.matches(&input.old_string).count();
467
468 if count == 0 {
469 return Ok(ToolResult::err(format!(
470 "old_string not found in {}. Make sure it matches exactly, including whitespace and indentation.",
471 path.display()
472 )));
473 }
474
475 if count > 1 && !replace_all {
476 return Ok(ToolResult::err(format!(
477 "old_string found {count} times in {}. Provide more surrounding context to make it unique, or set replace_all to true.",
478 path.display()
479 )));
480 }
481
482 let new_content = if replace_all {
483 content.replace(&input.old_string, &input.new_string)
484 } else {
485 content.replacen(&input.old_string, &input.new_string, 1)
487 };
488
489 match tokio::fs::write(&path, &new_content).await {
490 Ok(()) => {
491 let replacements = if replace_all {
492 format!("{count} replacement(s)")
493 } else {
494 "1 replacement".to_string()
495 };
496 Ok(ToolResult::ok(format!(
497 "Successfully edited {} ({replacements})",
498 path.display()
499 )))
500 }
501 Err(e) => Ok(ToolResult::err(format!(
502 "Failed to write {}: {e}",
503 path.display()
504 ))),
505 }
506 }
507
508 async fn execute_bash(&self, input: &BashInput) -> Result<ToolResult> {
513 if input.run_in_background == Some(true) {
515 let mut cmd = Command::new("/bin/bash");
516 cmd.arg("-c")
517 .arg(&input.command)
518 .current_dir(&self.cwd)
519 .env("HOME", &self.cwd)
520 .stdout(std::process::Stdio::null())
521 .stderr(std::process::Stdio::null());
522 for (key, value) in &self.env_inject {
523 cmd.env(key, value);
524 }
525 for key in &self.env_blocklist {
526 cmd.env_remove(key);
527 }
528 #[cfg(unix)]
529 if let Some(ref hook) = self.pre_exec {
530 let hook = Arc::clone(hook);
531 unsafe {
534 cmd.pre_exec(move || (hook)());
535 }
536 }
537 let child = cmd.spawn();
538
539 return match child {
540 Ok(child) => {
541 let pid = child.id().unwrap_or(0);
542 Ok(ToolResult {
543 content: format!("Process started in background (pid: {pid})"),
544 is_error: false,
545 raw_content: None,
546 })
547 }
548 Err(e) => Ok(ToolResult::err(format!("Failed to spawn process: {e}"))),
549 };
550 }
551
552 let timeout_ms = input.timeout.unwrap_or(120_000);
553 let timeout_dur = Duration::from_millis(timeout_ms);
554
555 let mut cmd = Command::new("/bin/bash");
556 cmd.arg("-c")
557 .arg(&input.command)
558 .current_dir(&self.cwd)
559 .env("HOME", &self.cwd)
560 .stdout(std::process::Stdio::piped())
561 .stderr(std::process::Stdio::piped());
562 for (key, value) in &self.env_inject {
563 cmd.env(key, value);
564 }
565 for key in &self.env_blocklist {
566 cmd.env_remove(key);
567 }
568 #[cfg(unix)]
569 if let Some(ref hook) = self.pre_exec {
570 let hook = Arc::clone(hook);
571 unsafe {
572 cmd.pre_exec(move || (hook)());
573 }
574 }
575 let child = cmd.spawn();
576
577 let mut child = match child {
578 Ok(c) => c,
579 Err(e) => {
580 return Ok(ToolResult::err(format!("Failed to spawn process: {e}")));
581 }
582 };
583
584 let stdout_handle = child.stdout.take();
586 let stderr_handle = child.stderr.take();
587
588 let wait_result = tokio::time::timeout(timeout_dur, child.wait()).await;
590
591 match wait_result {
592 Ok(Ok(status)) => {
593 let (stdout_bytes, stderr_bytes) =
597 tokio::join!(drain_pipe(stdout_handle), drain_pipe(stderr_handle),);
598
599 let stdout = String::from_utf8_lossy(&stdout_bytes);
600 let stderr = String::from_utf8_lossy(&stderr_bytes);
601
602 let mut combined = String::new();
603 if !stdout.is_empty() {
604 combined.push_str(&stdout);
605 }
606 if !stderr.is_empty() {
607 if !combined.is_empty() {
608 combined.push('\n');
609 }
610 combined.push_str(&stderr);
611 }
612
613 let is_error = !status.success();
614 if is_error && combined.is_empty() {
615 combined = format!("Process exited with code {}", status.code().unwrap_or(-1));
616 }
617
618 Ok(ToolResult {
619 content: combined,
620 is_error,
621 raw_content: None,
622 })
623 }
624 Ok(Err(e)) => Ok(ToolResult::err(format!("Process IO error: {e}"))),
625 Err(_) => {
626 let _ = child.kill().await;
628 Ok(ToolResult::err(format!(
629 "Command timed out after {timeout_ms}ms"
630 )))
631 }
632 }
633 }
634
635 async fn execute_glob(&self, input: &GlobInput) -> Result<ToolResult> {
639 let base = match &input.path {
640 Some(p) => match self.resolve_and_check(p) {
641 Ok(resolved) => resolved,
642 Err(denied) => return Ok(denied),
643 },
644 None => self.cwd.clone(),
645 };
646
647 let full_pattern = base.join(&input.pattern);
648 let pattern_str = full_pattern.to_string_lossy().to_string();
649
650 let result =
652 tokio::task::spawn_blocking(move || -> std::result::Result<Vec<String>, String> {
653 let entries =
654 glob_match(&pattern_str).map_err(|e| format!("Invalid glob pattern: {e}"))?;
655
656 let mut paths: Vec<String> = Vec::new();
657 for entry in entries {
658 match entry {
659 Ok(p) => paths.push(p.to_string_lossy().to_string()),
660 Err(e) => {
661 warn!("glob entry error: {e}");
662 }
663 }
664 }
665 paths.sort();
666 Ok(paths)
667 })
668 .await
669 .map_err(|e| AgentError::ToolExecution(format!("glob task panicked: {e}")))?;
670
671 match result {
672 Ok(paths) => {
673 if paths.is_empty() {
674 Ok(ToolResult::ok("No files matched the pattern.".to_string()))
675 } else {
676 Ok(ToolResult::ok(paths.join("\n")))
677 }
678 }
679 Err(e) => Ok(ToolResult::err(e)),
680 }
681 }
682
683 async fn execute_grep(&self, input: &GrepInput) -> Result<ToolResult> {
688 if let Some(ref p) = input.path {
689 if let Err(denied) = self.resolve_and_check(p) {
690 return Ok(denied);
691 }
692 }
693
694 let input = input.clone();
695 let cwd = self.cwd.clone();
696
697 let result = tokio::task::spawn_blocking(move || grep_sync(&input, &cwd))
699 .await
700 .map_err(|e| AgentError::ToolExecution(format!("grep task panicked: {e}")))?;
701
702 result
703 }
704}
705
706fn extensions_for_type(file_type: &str) -> Option<Vec<&'static str>> {
710 let map: HashMap<&str, Vec<&str>> = HashMap::from([
711 ("rust", vec!["rs"]),
712 ("rs", vec!["rs"]),
713 ("py", vec!["py", "pyi"]),
714 ("python", vec!["py", "pyi"]),
715 ("js", vec!["js", "mjs", "cjs"]),
716 ("ts", vec!["ts", "tsx", "mts", "cts"]),
717 ("go", vec!["go"]),
718 ("java", vec!["java"]),
719 ("c", vec!["c", "h"]),
720 ("cpp", vec!["cpp", "cxx", "cc", "hpp", "hxx", "hh", "h"]),
721 ("rb", vec!["rb"]),
722 ("ruby", vec!["rb"]),
723 ("html", vec!["html", "htm"]),
724 ("css", vec!["css"]),
725 ("json", vec!["json"]),
726 ("yaml", vec!["yaml", "yml"]),
727 ("toml", vec!["toml"]),
728 ("md", vec!["md", "markdown"]),
729 ("sh", vec!["sh", "bash", "zsh"]),
730 ("sql", vec!["sql"]),
731 ("xml", vec!["xml"]),
732 ("swift", vec!["swift"]),
733 ("kt", vec!["kt", "kts"]),
734 ("scala", vec!["scala"]),
735 ]);
736 map.get(file_type).cloned()
737}
738
739fn matches_file_filter(
741 path: &Path,
742 glob_filter: &Option<glob::Pattern>,
743 type_exts: &Option<Vec<&str>>,
744) -> bool {
745 if let Some(pat) = glob_filter {
746 let name = path.file_name().unwrap_or_default().to_string_lossy();
747 if !pat.matches(&name) {
748 return false;
749 }
750 }
751 if let Some(exts) = type_exts {
752 let ext = path
753 .extension()
754 .unwrap_or_default()
755 .to_string_lossy()
756 .to_lowercase();
757 if !exts.contains(&ext.as_str()) {
758 return false;
759 }
760 }
761 true
762}
763
764fn walk_files(dir: &Path) -> Vec<PathBuf> {
766 let mut files = Vec::new();
767 walk_files_recursive(dir, &mut files);
768 files.sort();
769 files
770}
771
772fn walk_files_recursive(dir: &Path, out: &mut Vec<PathBuf>) {
773 let entries = match std::fs::read_dir(dir) {
774 Ok(e) => e,
775 Err(_) => return,
776 };
777 for entry in entries.flatten() {
778 let path = entry.path();
779 let name = entry.file_name();
780 let name_str = name.to_string_lossy();
781
782 if name_str.starts_with('.') || name_str == "node_modules" || name_str == "target" {
784 continue;
785 }
786
787 if path.is_dir() {
788 walk_files_recursive(&path, out);
789 } else if path.is_file() {
790 out.push(path);
791 }
792 }
793}
794
795fn grep_sync(input: &GrepInput, cwd: &Path) -> Result<ToolResult> {
797 let output_mode = input
798 .output_mode
799 .clone()
800 .unwrap_or(GrepOutputMode::FilesWithMatches);
801 let case_insensitive = input.case_insensitive.unwrap_or(false);
802 let show_line_numbers = input.line_numbers.unwrap_or(true);
803 let multiline = input.multiline.unwrap_or(false);
804
805 let context_lines = input.context.or(input.context_alias);
807 let before_context = input.before_context.or(context_lines).unwrap_or(0) as usize;
808 let after_context = input.after_context.or(context_lines).unwrap_or(0) as usize;
809
810 let head_limit = input.head_limit.unwrap_or(0) as usize;
811 let offset = input.offset.unwrap_or(0) as usize;
812
813 let re = RegexBuilder::new(&input.pattern)
815 .case_insensitive(case_insensitive)
816 .multi_line(multiline)
817 .dot_matches_new_line(multiline)
818 .build()?;
819
820 let search_path = match &input.path {
822 Some(p) => {
823 let resolved = if Path::new(p).is_absolute() {
824 PathBuf::from(p)
825 } else {
826 cwd.join(p)
827 };
828 resolved
829 }
830 None => cwd.to_path_buf(),
831 };
832
833 let glob_filter = input
835 .glob
836 .as_ref()
837 .map(|g| glob::Pattern::new(g).unwrap_or_else(|_| glob::Pattern::new("*").unwrap()));
838 let type_exts = input
839 .file_type
840 .as_ref()
841 .and_then(|t| extensions_for_type(t).map(|v| v.into_iter().collect::<Vec<_>>()));
842
843 let files = if search_path.is_file() {
845 vec![search_path.clone()]
846 } else {
847 walk_files(&search_path)
848 };
849
850 let files: Vec<PathBuf> = files
852 .into_iter()
853 .filter(|f| matches_file_filter(f, &glob_filter, &type_exts))
854 .collect();
855
856 match output_mode {
857 GrepOutputMode::FilesWithMatches => {
858 grep_files_with_matches(&re, &files, offset, head_limit)
859 }
860 GrepOutputMode::Count => grep_count(&re, &files, offset, head_limit),
861 GrepOutputMode::Content => grep_content(
862 &re,
863 &files,
864 before_context,
865 after_context,
866 show_line_numbers,
867 offset,
868 head_limit,
869 ),
870 }
871}
872
873fn grep_files_with_matches(
874 re: &Regex,
875 files: &[PathBuf],
876 offset: usize,
877 head_limit: usize,
878) -> Result<ToolResult> {
879 let mut matched: Vec<String> = Vec::new();
880 for file in files {
881 if let Ok(content) = std::fs::read_to_string(file) {
882 if re.is_match(&content) {
883 matched.push(file.to_string_lossy().to_string());
884 }
885 }
886 }
887
888 let result = apply_offset_limit(matched, offset, head_limit);
889 if result.is_empty() {
890 Ok(ToolResult::ok("No matches found.".to_string()))
891 } else {
892 Ok(ToolResult::ok(result.join("\n")))
893 }
894}
895
896fn grep_count(
897 re: &Regex,
898 files: &[PathBuf],
899 offset: usize,
900 head_limit: usize,
901) -> Result<ToolResult> {
902 let mut entries: Vec<String> = Vec::new();
903 for file in files {
904 if let Ok(content) = std::fs::read_to_string(file) {
905 let count = re.find_iter(&content).count();
906 if count > 0 {
907 entries.push(format!("{}:{count}", file.to_string_lossy()));
908 }
909 }
910 }
911
912 let result = apply_offset_limit(entries, offset, head_limit);
913 if result.is_empty() {
914 Ok(ToolResult::ok("No matches found.".to_string()))
915 } else {
916 Ok(ToolResult::ok(result.join("\n")))
917 }
918}
919
920fn grep_content(
921 re: &Regex,
922 files: &[PathBuf],
923 before_context: usize,
924 after_context: usize,
925 show_line_numbers: bool,
926 offset: usize,
927 head_limit: usize,
928) -> Result<ToolResult> {
929 let mut output_lines: Vec<String> = Vec::new();
930
931 for file in files {
932 let content = match std::fs::read_to_string(file) {
933 Ok(c) => c,
934 Err(_) => continue,
935 };
936
937 let lines: Vec<&str> = content.lines().collect();
938 let file_display = file.to_string_lossy();
939
940 let mut matching_line_indices: Vec<usize> = Vec::new();
942 for (i, line) in lines.iter().enumerate() {
943 if re.is_match(line) {
944 matching_line_indices.push(i);
945 }
946 }
947
948 if matching_line_indices.is_empty() {
949 continue;
950 }
951
952 let mut display_set = Vec::new();
954 for &idx in &matching_line_indices {
955 let start = idx.saturating_sub(before_context);
956 let end = (idx + after_context + 1).min(lines.len());
957 for i in start..end {
958 display_set.push(i);
959 }
960 }
961 display_set.sort();
962 display_set.dedup();
963
964 let mut prev: Option<usize> = None;
966 for &line_idx in &display_set {
967 if let Some(p) = prev {
968 if line_idx > p + 1 {
969 output_lines.push("--".to_string());
970 }
971 }
972
973 let line_content = lines[line_idx];
974 if show_line_numbers {
975 let sep = if matching_line_indices.contains(&line_idx) {
976 ':'
977 } else {
978 '-'
979 };
980 output_lines.push(format!(
981 "{file_display}{sep}{}{sep}{line_content}",
982 line_idx + 1
983 ));
984 } else {
985 output_lines.push(format!("{file_display}:{line_content}"));
986 }
987
988 prev = Some(line_idx);
989 }
990 }
991
992 let result = apply_offset_limit(output_lines, offset, head_limit);
993 if result.is_empty() {
994 Ok(ToolResult::ok("No matches found.".to_string()))
995 } else {
996 Ok(ToolResult::ok(result.join("\n")))
997 }
998}
999
1000fn apply_offset_limit(items: Vec<String>, offset: usize, head_limit: usize) -> Vec<String> {
1002 let after_offset: Vec<String> = items.into_iter().skip(offset).collect();
1003 if head_limit > 0 {
1004 after_offset.into_iter().take(head_limit).collect()
1005 } else {
1006 after_offset
1007 }
1008}
1009
1010#[cfg(test)]
1011mod tests {
1012 use super::*;
1013 use serde_json::json;
1014 use tempfile::TempDir;
1015
1016 fn setup() -> (TempDir, ToolExecutor) {
1017 let tmp = TempDir::new().unwrap();
1018 let executor = ToolExecutor::new(tmp.path().to_path_buf());
1019 (tmp, executor)
1020 }
1021
1022 #[tokio::test]
1025 async fn read_text_file() {
1026 let (tmp, executor) = setup();
1027 let file = tmp.path().join("hello.txt");
1028 std::fs::write(&file, "line one\nline two\nline three\n").unwrap();
1029
1030 let result = executor
1031 .execute("Read", json!({ "file_path": file.to_str().unwrap() }))
1032 .await
1033 .unwrap();
1034
1035 assert!(!result.is_error);
1036 assert!(result.raw_content.is_none());
1037 assert!(result.content.contains("line one"));
1038 assert!(result.content.contains("line three"));
1039 }
1040
1041 #[tokio::test]
1042 async fn read_text_file_with_offset_and_limit() {
1043 let (tmp, executor) = setup();
1044 let file = tmp.path().join("lines.txt");
1045 std::fs::write(&file, "a\nb\nc\nd\ne\n").unwrap();
1046
1047 let result = executor
1048 .execute(
1049 "Read",
1050 json!({ "file_path": file.to_str().unwrap(), "offset": 1, "limit": 2 }),
1051 )
1052 .await
1053 .unwrap();
1054
1055 assert!(!result.is_error);
1056 assert!(result.content.contains("b"));
1058 assert!(result.content.contains("c"));
1059 assert!(!result.content.contains("a"));
1060 assert!(!result.content.contains("d"));
1061 }
1062
1063 #[tokio::test]
1064 async fn read_missing_file_returns_error() {
1065 let (tmp, executor) = setup();
1066 let file = tmp.path().join("nope.txt");
1067
1068 let result = executor
1069 .execute("Read", json!({ "file_path": file.to_str().unwrap() }))
1070 .await
1071 .unwrap();
1072
1073 assert!(result.is_error);
1074 assert!(result.content.contains("Failed to read"));
1075 }
1076
1077 #[tokio::test]
1080 async fn read_png_returns_image_content_block() {
1081 let (tmp, executor) = setup();
1082 let file = tmp.path().join("test.png");
1083 let png_bytes = b"\x89PNG\r\n\x1a\nfake-png-payload";
1084 std::fs::write(&file, png_bytes).unwrap();
1085
1086 let result = executor
1087 .execute("Read", json!({ "file_path": file.to_str().unwrap() }))
1088 .await
1089 .unwrap();
1090
1091 assert!(!result.is_error);
1092 assert!(result.raw_content.is_some(), "image should set raw_content");
1093
1094 let blocks = result.raw_content.unwrap();
1095 let block = blocks.as_array().unwrap().first().unwrap();
1096 assert_eq!(block["type"], "image");
1097 assert_eq!(block["source"]["type"], "base64");
1098 assert_eq!(block["source"]["media_type"], "image/png");
1099 let data = block["source"]["data"].as_str().unwrap();
1101 assert!(!data.is_empty());
1102 use base64::Engine;
1103 let decoded = base64::engine::general_purpose::STANDARD
1104 .decode(data)
1105 .unwrap();
1106 assert_eq!(decoded, png_bytes);
1107 }
1108
1109 #[tokio::test]
1110 async fn read_jpeg_returns_image_content_block() {
1111 let (tmp, executor) = setup();
1112 let file = tmp.path().join("photo.jpg");
1115 let fake_jpeg = b"\xFF\xD8\xFF\xE0fake-jpeg-data";
1116 std::fs::write(&file, fake_jpeg).unwrap();
1117
1118 let result = executor
1119 .execute("Read", json!({ "file_path": file.to_str().unwrap() }))
1120 .await
1121 .unwrap();
1122
1123 assert!(!result.is_error);
1124 let blocks = result.raw_content.unwrap();
1125 let block = blocks.as_array().unwrap().first().unwrap();
1126 assert_eq!(block["source"]["media_type"], "image/jpeg");
1127 }
1128
1129 #[tokio::test]
1130 async fn read_jpeg_extension_detected() {
1131 let (tmp, executor) = setup();
1132 let file = tmp.path().join("photo.jpeg");
1133 std::fs::write(&file, b"data").unwrap();
1134
1135 let result = executor
1136 .execute("Read", json!({ "file_path": file.to_str().unwrap() }))
1137 .await
1138 .unwrap();
1139
1140 assert!(!result.is_error);
1141 let blocks = result.raw_content.unwrap();
1142 assert_eq!(blocks[0]["source"]["media_type"], "image/jpeg");
1143 }
1144
1145 #[tokio::test]
1146 async fn read_gif_returns_image_content_block() {
1147 let (tmp, executor) = setup();
1148 let file = tmp.path().join("anim.gif");
1149 std::fs::write(&file, b"GIF89adata").unwrap();
1150
1151 let result = executor
1152 .execute("Read", json!({ "file_path": file.to_str().unwrap() }))
1153 .await
1154 .unwrap();
1155
1156 assert!(!result.is_error);
1157 let blocks = result.raw_content.unwrap();
1158 assert_eq!(blocks[0]["source"]["media_type"], "image/gif");
1159 }
1160
1161 #[tokio::test]
1162 async fn read_webp_returns_image_content_block() {
1163 let (tmp, executor) = setup();
1164 let file = tmp.path().join("img.webp");
1165 std::fs::write(&file, b"RIFF\x00\x00\x00\x00WEBP").unwrap();
1166
1167 let result = executor
1168 .execute("Read", json!({ "file_path": file.to_str().unwrap() }))
1169 .await
1170 .unwrap();
1171
1172 assert!(!result.is_error);
1173 let blocks = result.raw_content.unwrap();
1174 assert_eq!(blocks[0]["source"]["media_type"], "image/webp");
1175 }
1176
1177 #[tokio::test]
1178 async fn read_missing_image_returns_error() {
1179 let (tmp, executor) = setup();
1180 let file = tmp.path().join("nope.png");
1181
1182 let result = executor
1183 .execute("Read", json!({ "file_path": file.to_str().unwrap() }))
1184 .await
1185 .unwrap();
1186
1187 assert!(result.is_error);
1188 assert!(result.content.contains("Failed to read"));
1189 assert!(result.raw_content.is_none());
1190 }
1191
1192 #[tokio::test]
1193 async fn read_non_image_extension_returns_text() {
1194 let (tmp, executor) = setup();
1195 let file = tmp.path().join("data.csv");
1196 std::fs::write(&file, "a,b,c\n1,2,3\n").unwrap();
1197
1198 let result = executor
1199 .execute("Read", json!({ "file_path": file.to_str().unwrap() }))
1200 .await
1201 .unwrap();
1202
1203 assert!(!result.is_error);
1204 assert!(
1205 result.raw_content.is_none(),
1206 "csv should not be treated as image"
1207 );
1208 assert!(result.content.contains("a,b,c"));
1209 }
1210
1211 #[test]
1214 fn tool_result_ok_has_no_raw_content() {
1215 let r = ToolResult::ok("hello".into());
1216 assert!(!r.is_error);
1217 assert!(r.raw_content.is_none());
1218 }
1219
1220 #[test]
1221 fn tool_result_err_has_no_raw_content() {
1222 let r = ToolResult::err("boom".into());
1223 assert!(r.is_error);
1224 assert!(r.raw_content.is_none());
1225 }
1226
1227 fn setup_sandboxed() -> (TempDir, TempDir, ToolExecutor) {
1230 let project = TempDir::new().unwrap();
1231 let data = TempDir::new().unwrap();
1232 let executor = ToolExecutor::with_allowed_dirs(
1233 project.path().to_path_buf(),
1234 vec![data.path().to_path_buf()],
1235 );
1236 (project, data, executor)
1237 }
1238
1239 #[tokio::test]
1240 async fn sandbox_allows_read_inside_cwd() {
1241 let (project, _data, executor) = setup_sandboxed();
1242 let file = project.path().join("hello.txt");
1243 std::fs::write(&file, "ok").unwrap();
1244
1245 let result = executor
1246 .execute("Read", json!({ "file_path": file.to_str().unwrap() }))
1247 .await
1248 .unwrap();
1249 assert!(!result.is_error);
1250 }
1251
1252 #[tokio::test]
1253 async fn sandbox_allows_read_inside_additional_dir() {
1254 let (_project, data, executor) = setup_sandboxed();
1255 let file = data.path().join("MEMORY.md");
1256 std::fs::write(&file, "# Memory").unwrap();
1257
1258 let result = executor
1259 .execute("Read", json!({ "file_path": file.to_str().unwrap() }))
1260 .await
1261 .unwrap();
1262 assert!(!result.is_error);
1263 assert!(result.content.contains("Memory"));
1264 }
1265
1266 #[tokio::test]
1267 async fn sandbox_denies_read_outside_boundaries() {
1268 let (_project, _data, executor) = setup_sandboxed();
1269 let outside = TempDir::new().unwrap();
1270 let file = outside.path().join("secret.txt");
1271 std::fs::write(&file, "secret data").unwrap();
1272
1273 let result = executor
1274 .execute("Read", json!({ "file_path": file.to_str().unwrap() }))
1275 .await
1276 .unwrap();
1277 assert!(result.is_error);
1278 assert!(result.content.contains("Access denied"));
1279 }
1280
1281 #[tokio::test]
1282 async fn sandbox_denies_write_outside_boundaries() {
1283 let (_project, _data, executor) = setup_sandboxed();
1284 let outside = TempDir::new().unwrap();
1285 let file = outside.path().join("hack.txt");
1286
1287 let result = executor
1288 .execute(
1289 "Write",
1290 json!({ "file_path": file.to_str().unwrap(), "content": "pwned" }),
1291 )
1292 .await
1293 .unwrap();
1294 assert!(result.is_error);
1295 assert!(result.content.contains("Access denied"));
1296 assert!(!file.exists());
1297 }
1298
1299 #[tokio::test]
1300 async fn sandbox_denies_edit_outside_boundaries() {
1301 let (_project, _data, executor) = setup_sandboxed();
1302 let outside = TempDir::new().unwrap();
1303 let file = outside.path().join("target.txt");
1304 std::fs::write(&file, "original").unwrap();
1305
1306 let result = executor
1307 .execute(
1308 "Edit",
1309 json!({
1310 "file_path": file.to_str().unwrap(),
1311 "old_string": "original",
1312 "new_string": "modified"
1313 }),
1314 )
1315 .await
1316 .unwrap();
1317 assert!(result.is_error);
1318 assert!(result.content.contains("Access denied"));
1319 }
1320
1321 #[tokio::test]
1322 async fn no_sandbox_when_allowed_dirs_empty() {
1323 let outside = TempDir::new().unwrap();
1325 let file = outside.path().join("free.txt");
1326 std::fs::write(&file, "open access").unwrap();
1327
1328 let executor = ToolExecutor::new(TempDir::new().unwrap().path().to_path_buf());
1329 let result = executor
1330 .execute("Read", json!({ "file_path": file.to_str().unwrap() }))
1331 .await
1332 .unwrap();
1333 assert!(!result.is_error);
1334 }
1335
1336 #[tokio::test]
1337 async fn sandbox_denies_dotdot_traversal() {
1338 let (project, _data, executor) = setup_sandboxed();
1339 let outside = TempDir::new().unwrap();
1341 let secret = outside.path().join("secret.txt");
1342 std::fs::write(&secret, "sensitive").unwrap();
1343
1344 let traversal = project
1346 .path()
1347 .join("..")
1348 .join("..")
1349 .join(outside.path().strip_prefix("/").unwrap())
1350 .join("secret.txt");
1351
1352 let result = executor
1353 .execute("Read", json!({ "file_path": traversal.to_str().unwrap() }))
1354 .await
1355 .unwrap();
1356 assert!(result.is_error);
1357 assert!(result.content.contains("Access denied"));
1358 }
1359
1360 #[tokio::test]
1361 async fn sandbox_allows_write_new_file_inside_cwd() {
1362 let (project, _data, executor) = setup_sandboxed();
1363 let file = project.path().join("subdir").join("new.txt");
1364
1365 let result = executor
1366 .execute(
1367 "Write",
1368 json!({ "file_path": file.to_str().unwrap(), "content": "hello" }),
1369 )
1370 .await
1371 .unwrap();
1372 assert!(!result.is_error);
1373 assert!(file.exists());
1374 }
1375
1376 #[tokio::test]
1377 async fn sandbox_denies_symlink_escape() {
1378 let (project, _data, executor) = setup_sandboxed();
1379 let outside = TempDir::new().unwrap();
1380 let secret = outside.path().join("secret.txt");
1381 std::fs::write(&secret, "sensitive").unwrap();
1382
1383 let link = project.path().join("escape");
1385 std::os::unix::fs::symlink(outside.path(), &link).unwrap();
1386
1387 let via_link = link.join("secret.txt");
1388 let result = executor
1389 .execute("Read", json!({ "file_path": via_link.to_str().unwrap() }))
1390 .await
1391 .unwrap();
1392 assert!(result.is_error);
1393 assert!(result.content.contains("Access denied"));
1394 }
1395
1396 #[tokio::test]
1397 async fn bash_sets_home_to_cwd() {
1398 let (tmp, executor) = setup();
1399
1400 let result = executor
1401 .execute("Bash", json!({ "command": "echo $HOME" }))
1402 .await
1403 .unwrap();
1404 assert!(!result.is_error, "Bash should succeed");
1405
1406 let expected = tmp.path().to_string_lossy().to_string();
1407 assert!(
1408 result.content.trim().contains(&expected),
1409 "HOME should be set to cwd ({}), got: {}",
1410 expected,
1411 result.content.trim()
1412 );
1413 }
1414
1415 #[tokio::test]
1416 async fn bash_tilde_resolves_to_cwd() {
1417 let (tmp, executor) = setup();
1418
1419 std::fs::write(tmp.path().join("marker.txt"), "found").unwrap();
1421
1422 let result = executor
1423 .execute("Bash", json!({ "command": "cat ~/marker.txt" }))
1424 .await
1425 .unwrap();
1426 assert!(
1427 !result.is_error,
1428 "Should read file via ~: {}",
1429 result.content
1430 );
1431 assert!(result.content.contains("found"), "~ should resolve to cwd");
1432 }
1433
1434 #[tokio::test]
1435 async fn bash_env_blocklist_strips_vars() {
1436 let tmp = TempDir::new().unwrap();
1437 unsafe {
1439 std::env::set_var("STARPOD_TEST_SECRET", "leaked");
1440 }
1441
1442 let executor = ToolExecutor::new(tmp.path().to_path_buf())
1443 .with_env_blocklist(vec!["STARPOD_TEST_SECRET".to_string()]);
1444
1445 let result = executor
1446 .execute(
1447 "Bash",
1448 json!({ "command": "echo \"val=${STARPOD_TEST_SECRET}\"" }),
1449 )
1450 .await
1451 .unwrap();
1452 assert!(!result.is_error);
1453 assert_eq!(
1454 result.content.trim(),
1455 "val=",
1456 "Blocked env var should not be visible to child process"
1457 );
1458
1459 std::env::remove_var("STARPOD_TEST_SECRET");
1461 }
1462
1463 #[tokio::test]
1464 async fn bash_env_blocklist_does_not_affect_other_vars() {
1465 let tmp = TempDir::new().unwrap();
1466 unsafe {
1467 std::env::set_var("STARPOD_TEST_ALLOWED", "visible");
1468 std::env::set_var("STARPOD_TEST_BLOCKED", "hidden");
1469 }
1470
1471 let executor = ToolExecutor::new(tmp.path().to_path_buf())
1472 .with_env_blocklist(vec!["STARPOD_TEST_BLOCKED".to_string()]);
1473
1474 let result = executor
1475 .execute("Bash", json!({ "command": "echo $STARPOD_TEST_ALLOWED" }))
1476 .await
1477 .unwrap();
1478 assert!(
1479 result.content.contains("visible"),
1480 "Non-blocked vars should still be inherited"
1481 );
1482
1483 std::env::remove_var("STARPOD_TEST_ALLOWED");
1485 std::env::remove_var("STARPOD_TEST_BLOCKED");
1486 }
1487}