use async_trait::async_trait;
use serde_json::{json, Value};
use std::path::PathBuf;
use super::{resolve_within_any, AccessTier, SharedSandboxRoots, Tool};
use crate::error::{Error, Result};
use crate::llm::ToolSpec;
#[derive(Debug, Clone)]
pub struct CountLines {
pub root: PathBuf,
pub extra_roots: Vec<(PathBuf, AccessTier)>,
pub session_roots: Option<SharedSandboxRoots>,
}
impl CountLines {
pub fn new(root: impl Into<PathBuf>) -> Self {
Self {
root: root.into(),
extra_roots: Vec::new(),
session_roots: None,
}
}
pub fn with_extra_roots(
mut self,
extra: impl IntoIterator<Item = (PathBuf, AccessTier)>,
) -> Self {
self.extra_roots.extend(extra);
self
}
pub fn with_session_roots(mut self, slot: SharedSandboxRoots) -> Self {
self.session_roots = Some(slot);
self
}
pub fn with_session_roots_opt(mut self, slot: Option<SharedSandboxRoots>) -> Self {
if let Some(s) = slot {
self.session_roots = Some(s);
}
self
}
fn all_roots(&self) -> Vec<(PathBuf, AccessTier)> {
let mut v: Vec<(PathBuf, AccessTier)> = Vec::with_capacity(self.extra_roots.len() + 1);
v.push((self.root.clone(), AccessTier::ReadWrite));
v.extend(self.extra_roots.iter().cloned());
if let Some(slot) = &self.session_roots {
if let Ok(roots) = slot.read() {
v.extend(roots.iter().cloned());
}
}
v
}
}
#[async_trait]
impl Tool for CountLines {
fn spec(&self) -> ToolSpec {
ToolSpec {
name: "count_lines".into(),
description: "Return the number of lines in a text file inside the workspace.".into(),
parameters: json!({
"type": "object",
"properties": {
"path": {
"type": "string",
"description": "Path relative to the workspace root"
}
},
"required": ["path"]
}),
}
}
fn side_effect_class(&self) -> crate::tools::ToolSideEffect {
crate::tools::ToolSideEffect::ReadOnly
}
async fn execute(&self, args: Value) -> Result<String> {
let path = args["path"].as_str().ok_or_else(|| Error::BadToolArgs {
name: "count_lines".into(),
message: "missing `path`".into(),
})?;
let abs = resolve_within_any(&self.all_roots(), path, false)?;
let content = tokio::fs::read_to_string(&abs)
.await
.map_err(|e| Error::Tool {
name: "count_lines".into(),
call_id: None,
message: format!("{}: {e}", abs.display()),
})?;
let count = content.lines().count();
Ok(count.to_string())
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[tokio::test]
async fn count_lines_happy_path() {
let tmp = TempDir::new().unwrap();
let contents = "line1\nline2\nline3\n";
std::fs::write(tmp.path().join("test.txt"), contents).unwrap();
let tool = CountLines::new(tmp.path());
let result = tool.execute(json!({"path": "test.txt"})).await.unwrap();
assert_eq!(result, "3");
}
#[tokio::test]
async fn count_lines_rejects_escape() {
let tmp = TempDir::new().unwrap();
let tool = CountLines::new(tmp.path());
let err = tool
.execute(json!({"path": "../etc/passwd"}))
.await
.unwrap_err();
assert!(matches!(err, Error::BadToolArgs { .. }));
}
#[tokio::test]
async fn count_lines_empty_file_returns_zero() {
let tmp = TempDir::new().unwrap();
std::fs::write(tmp.path().join("empty.txt"), "").unwrap();
let tool = CountLines::new(tmp.path());
let result = tool.execute(json!({"path": "empty.txt"})).await.unwrap();
assert_eq!(result, "0");
}
#[tokio::test]
async fn count_lines_no_trailing_newline() {
let tmp = TempDir::new().unwrap();
std::fs::write(tmp.path().join("no_nl.txt"), "a\nb").unwrap();
let tool = CountLines::new(tmp.path());
let result = tool.execute(json!({"path": "no_nl.txt"})).await.unwrap();
assert_eq!(result, "2");
}
#[tokio::test]
async fn count_lines_missing_path_argument_errors() {
let tmp = TempDir::new().unwrap();
let tool = CountLines::new(tmp.path());
let res = tool.execute(json!({})).await;
assert!(
matches!(res, Err(Error::BadToolArgs { .. })),
"missing 'path' must return BadToolArgs"
);
}
#[tokio::test]
async fn count_lines_nonexistent_file_errors() {
let tmp = TempDir::new().unwrap();
let tool = CountLines::new(tmp.path());
let res = tool.execute(json!({"path": "nonexistent.txt"})).await;
assert!(
matches!(res, Err(Error::Tool { .. })),
"nonexistent file must return Tool error"
);
}
}