pub(crate) mod edit_file;
pub(crate) mod glob_search;
pub(crate) mod grep;
pub(crate) mod list_dir;
pub(crate) mod read_file;
use crate::chat::edits::PendingEditStore;
use crate::workspace_fs::WorkspaceFs;
use async_trait::async_trait;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::sync::mpsc;
pub(crate) use crate::fswalk::{default_walker, path_to_forward_slash};
#[derive(Clone)]
pub(crate) struct ToolContext {
workspace_fs: Arc<WorkspaceFs>,
pub pending_edits: Option<Arc<PendingEditStore>>,
pub event_sink: Option<mpsc::Sender<crate::chat::agent::AgentEvent>>,
}
impl ToolContext {
#[cfg(test)]
pub(crate) fn new(root: impl AsRef<Path>) -> Result<Self, ToolError> {
let workspace_root = dunce::canonicalize(root.as_ref())
.map_err(|e| ToolError::Io(format!("workspace root: {e}")))?;
let workspace_fs = Arc::new(WorkspaceFs::new(workspace_root.clone(), None));
Ok(Self {
workspace_fs,
pending_edits: None,
event_sink: None,
})
}
#[cfg(test)]
pub(crate) fn for_single_file(root: impl AsRef<Path>, file: &str) -> Result<Self, ToolError> {
let workspace_root = dunce::canonicalize(root.as_ref())
.map_err(|e| ToolError::Io(format!("workspace root: {e}")))?;
Self::for_workspace(Arc::new(WorkspaceFs::new(workspace_root, Some(file))))
}
pub(crate) fn for_workspace(workspace_fs: Arc<WorkspaceFs>) -> Result<Self, ToolError> {
if !workspace_fs.ambient_root().is_dir() {
return Err(ToolError::Io("workspace root is unavailable".to_string()));
}
Ok(Self {
workspace_fs,
pending_edits: None,
event_sink: None,
})
}
pub(crate) fn content_files(&self, limit: usize) -> Vec<(String, PathBuf)> {
self.workspace_fs
.content_files(limit)
.into_iter()
.map(|(rel, abs)| (rel.as_route(), abs))
.collect()
}
pub(crate) fn directory_root(&self) -> Option<&Path> {
self.workspace_fs.directory_root()
}
pub(crate) fn route_for_path(&self, path: &Path) -> Option<String> {
self.workspace_fs.route_for_path(path)
}
pub(crate) fn with_chat_state(
mut self,
pending_edits: Arc<PendingEditStore>,
event_sink: mpsc::Sender<crate::chat::agent::AgentEvent>,
) -> Self {
self.pending_edits = Some(pending_edits);
self.event_sink = Some(event_sink);
self
}
pub(crate) fn resolve(&self, rel: &str) -> Result<PathBuf, ToolError> {
self.workspace_fs
.resolve_content(rel)
.map_err(|error| match error {
crate::workspace_fs::WorkspaceFsError::InvalidPath
| crate::workspace_fs::WorkspaceFsError::Denied => ToolError::OutsideWorkspace,
crate::workspace_fs::WorkspaceFsError::NotFound => {
ToolError::NotFound(rel.to_string())
}
crate::workspace_fs::WorkspaceFsError::Io(message) => ToolError::Io(message),
})
}
}
#[derive(Debug, thiserror::Error)]
pub(crate) enum ToolError {
#[error("not found: {0}")]
NotFound(String),
#[error("invalid argument: {0}")]
InvalidArgument(String),
#[error("path escapes workspace root")]
OutsideWorkspace,
#[error("file too large")]
TooLarge,
#[error("binary or non-utf8 file")]
Binary,
#[error("io error: {0}")]
Io(String),
#[error("internal error: {0}")]
Internal(String),
}
impl ToolError {
pub(crate) fn to_tool_message(&self) -> String {
self.to_string()
}
}
#[derive(Debug, Clone)]
pub(crate) struct ToolSchema {
pub name: String,
pub description: String,
pub input_schema: serde_json::Value,
}
#[async_trait]
pub(crate) trait Tool: Send + Sync {
fn name(&self) -> &'static str;
fn description(&self) -> &'static str;
fn input_schema(&self) -> serde_json::Value;
async fn run(&self, ctx: &ToolContext, input: serde_json::Value) -> Result<String, ToolError>;
fn schema(&self) -> ToolSchema {
ToolSchema {
name: self.name().to_string(),
description: self.description().to_string(),
input_schema: self.input_schema(),
}
}
}
pub(crate) struct ToolRegistry {
tools: HashMap<&'static str, Arc<dyn Tool>>,
}
impl ToolRegistry {
pub(crate) fn new() -> Self {
Self {
tools: HashMap::new(),
}
}
pub(crate) fn with_default_tools() -> Self {
let mut r = Self::new();
r.register(Arc::new(read_file::ReadFileTool));
r.register(Arc::new(list_dir::ListDirTool));
r.register(Arc::new(glob_search::GlobTool));
r.register(Arc::new(grep::GrepTool));
r
}
pub(crate) fn for_workspace(enable_edit: bool) -> Self {
let mut r = Self::with_default_tools();
if enable_edit {
r.register(Arc::new(edit_file::EditFileTool));
}
r
}
pub(crate) fn register(&mut self, tool: Arc<dyn Tool>) {
self.tools.insert(tool.name(), tool);
}
pub(crate) fn get(&self, name: &str) -> Option<&Arc<dyn Tool>> {
self.tools.get(name)
}
pub(crate) fn schemas(&self) -> Vec<ToolSchema> {
self.tools.values().map(|t| t.schema()).collect()
}
pub(crate) async fn dispatch(
&self,
ctx: &ToolContext,
name: &str,
input: serde_json::Value,
) -> Result<String, ToolError> {
let tool = self
.get(name)
.ok_or_else(|| ToolError::NotFound(format!("tool '{name}' is unknown")))?;
tool.run(ctx, input).await
}
}
pub(crate) const MAX_TOOL_OUTPUT_BYTES: usize = 64 * 1024;
pub(crate) const MAX_FILE_BYTES: u64 = 1024 * 1024;
pub(crate) fn looks_binary(bytes: &[u8]) -> bool {
let scan = bytes.len().min(8192);
bytes[..scan].contains(&0)
}
#[cfg(test)]
mod resolve_tests {
use super::*;
use std::fs;
use tempfile::TempDir;
fn ctx() -> (TempDir, ToolContext) {
let td = TempDir::new().unwrap();
fs::write(td.path().join("ok.md"), b"hi").unwrap();
let ctx = ToolContext::new(td.path()).unwrap();
(td, ctx)
}
#[test]
fn accepts_normal_relative_path() {
let (_td, ctx) = ctx();
let p = ctx.resolve("ok.md").unwrap();
assert!(p.ends_with("ok.md"));
assert_eq!(ctx.route_for_path(&p).as_deref(), Some("ok.md"));
}
#[test]
fn accepts_curdir_component() {
let (_td, ctx) = ctx();
ctx.resolve("./ok.md").unwrap();
}
#[test]
fn rejects_parent_traversal_even_when_target_missing() {
let (_td, ctx) = ctx();
let err = ctx.resolve("../etc/passwd").unwrap_err();
assert!(matches!(err, ToolError::OutsideWorkspace), "got {err:?}");
}
#[test]
fn rejects_parent_traversal_with_existing_outside_target() {
let (_td, ctx) = ctx();
let err = ctx.resolve("../").unwrap_err();
assert!(matches!(err, ToolError::OutsideWorkspace), "got {err:?}");
}
#[test]
fn rejects_absolute_path() {
let (_td, ctx) = ctx();
#[cfg(unix)]
let abs = "/etc/passwd";
#[cfg(windows)]
let abs = "C:\\Windows\\System32\\drivers\\etc\\hosts";
let err = ctx.resolve(abs).unwrap_err();
assert!(matches!(err, ToolError::OutsideWorkspace), "got {err:?}");
}
#[test]
fn missing_file_returns_not_found() {
let (_td, ctx) = ctx();
let err = ctx.resolve("does-not-exist.md").unwrap_err();
assert!(matches!(err, ToolError::NotFound(_)), "got {err:?}");
}
#[cfg(unix)]
#[test]
fn rejects_symlink_escaping_workspace() {
use std::os::unix::fs::symlink;
let td = TempDir::new().unwrap();
let outside = TempDir::new().unwrap();
fs::write(outside.path().join("secret.txt"), b"nope").unwrap();
symlink(outside.path(), td.path().join("escape")).unwrap();
let ctx = ToolContext::new(td.path()).unwrap();
let err = ctx.resolve("escape/secret.txt").unwrap_err();
assert!(matches!(err, ToolError::OutsideWorkspace), "got {err:?}");
}
#[tokio::test]
async fn single_file_scope_is_enforced_across_all_chat_tools() {
let td = TempDir::new().unwrap();
fs::write(td.path().join("opened.md"), "visible marker\n").unwrap();
fs::write(td.path().join("sibling.md"), "secret marker\n").unwrap();
let ctx = ToolContext::for_single_file(td.path(), "opened.md").unwrap();
let tools = ToolRegistry::for_workspace(true);
let read = tools
.dispatch(
&ctx,
"read_file",
serde_json::json!({ "path": "opened.md" }),
)
.await
.unwrap();
assert!(read.contains("visible marker"));
assert!(matches!(
tools
.dispatch(
&ctx,
"read_file",
serde_json::json!({ "path": "sibling.md" })
)
.await,
Err(ToolError::OutsideWorkspace)
));
let listing = tools
.dispatch(&ctx, "list_dir", serde_json::json!({}))
.await
.unwrap();
assert!(listing.contains("opened.md"), "got: {listing}");
assert!(!listing.contains("sibling.md"), "got: {listing}");
let glob = tools
.dispatch(&ctx, "glob", serde_json::json!({ "pattern": "**/*.md" }))
.await
.unwrap();
assert_eq!(glob, "opened.md");
let grep = tools
.dispatch(&ctx, "grep", serde_json::json!({ "pattern": "marker" }))
.await
.unwrap();
assert!(grep.contains("opened.md:1:visible marker"), "got: {grep}");
assert!(!grep.contains("sibling.md"), "got: {grep}");
assert!(matches!(
tools
.dispatch(
&ctx,
"grep",
serde_json::json!({ "pattern": "secret", "path": "sibling.md" })
)
.await,
Err(ToolError::OutsideWorkspace)
));
assert!(matches!(
tools
.dispatch(
&ctx,
"edit_file",
serde_json::json!({
"path": "sibling.md",
"old_string": "secret",
"new_string": "stolen"
})
)
.await,
Err(ToolError::OutsideWorkspace)
));
assert_eq!(
fs::read_to_string(td.path().join("sibling.md")).unwrap(),
"secret marker\n"
);
}
}