組み込みツールの実装
This commit is contained in:
parent
a05eec42d7
commit
3d0d5ffe85
212
Cargo.lock
generated
212
Cargo.lock
generated
|
|
@ -156,6 +156,17 @@ dependencies = [
|
|||
"hybrid-array",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bstr"
|
||||
version = "1.12.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "63044e1ae8e69f3b5a92c736ca6269b8d12fa7efe39bf34ddb06d102cf0e2cab"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
"regex-automata",
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bumpalo"
|
||||
version = "3.20.2"
|
||||
|
|
@ -314,6 +325,31 @@ dependencies = [
|
|||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crossbeam-deque"
|
||||
version = "0.8.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51"
|
||||
dependencies = [
|
||||
"crossbeam-epoch",
|
||||
"crossbeam-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crossbeam-epoch"
|
||||
version = "0.9.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e"
|
||||
dependencies = [
|
||||
"crossbeam-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crossbeam-utils"
|
||||
version = "0.8.21"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28"
|
||||
|
||||
[[package]]
|
||||
name = "crossterm"
|
||||
version = "0.28.1"
|
||||
|
|
@ -525,6 +561,24 @@ version = "1.15.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
|
||||
|
||||
[[package]]
|
||||
name = "encoding_rs"
|
||||
version = "0.8.35"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "75030f3c4f45dafd7586dd6780965a8c7e8e285a5ecb86713e63a79c5b2766f3"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "encoding_rs_io"
|
||||
version = "0.1.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1cc3c5651fb62ab8aa3103998dade57efdd028544bd300516baa31840c252a83"
|
||||
dependencies = [
|
||||
"encoding_rs",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "equivalent"
|
||||
version = "1.0.2"
|
||||
|
|
@ -588,6 +642,17 @@ dependencies = [
|
|||
"winapi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "filetime"
|
||||
version = "0.2.27"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f98844151eee8917efc50bd9e8318cb963ae8b297431495d3f758616ea5c57db"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"libc",
|
||||
"libredox",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "find-msvc-tools"
|
||||
version = "0.1.9"
|
||||
|
|
@ -788,6 +853,56 @@ version = "0.3.3"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280"
|
||||
|
||||
[[package]]
|
||||
name = "globset"
|
||||
version = "0.4.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "52dfc19153a48bde0cbd630453615c8151bce3a5adfac7a0aebfbf0a1e1f57e3"
|
||||
dependencies = [
|
||||
"aho-corasick",
|
||||
"bstr",
|
||||
"log",
|
||||
"regex-automata",
|
||||
"regex-syntax",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "grep-matcher"
|
||||
version = "0.1.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "36d7b71093325ab22d780b40d7df3066ae4aebb518ba719d38c697a8228a8023"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "grep-regex"
|
||||
version = "0.1.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0ce0c256c3ad82bcc07b812c15a45ec1d398122e8e15124f96695234db7112ef"
|
||||
dependencies = [
|
||||
"bstr",
|
||||
"grep-matcher",
|
||||
"log",
|
||||
"regex-automata",
|
||||
"regex-syntax",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "grep-searcher"
|
||||
version = "0.1.16"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ac63295322dc48ebb20a25348147905d816318888e64f531bfc2a2bc0577dc34"
|
||||
dependencies = [
|
||||
"bstr",
|
||||
"encoding_rs",
|
||||
"encoding_rs_io",
|
||||
"grep-matcher",
|
||||
"log",
|
||||
"memchr",
|
||||
"memmap2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "h2"
|
||||
version = "0.4.13"
|
||||
|
|
@ -1078,6 +1193,22 @@ dependencies = [
|
|||
"icu_properties",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ignore"
|
||||
version = "0.4.25"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d3d782a365a015e0f5c04902246139249abf769125006fbe7649e2ee88169b4a"
|
||||
dependencies = [
|
||||
"crossbeam-deque",
|
||||
"globset",
|
||||
"log",
|
||||
"memchr",
|
||||
"regex-automata",
|
||||
"same-file",
|
||||
"walkdir",
|
||||
"winapi-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "indexmap"
|
||||
version = "2.13.1"
|
||||
|
|
@ -1196,6 +1327,18 @@ version = "0.2.184"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "48f5d2a454e16a5ea0f4ced81bd44e4cfc7bd3a507b61887c99fd3538b28e4af"
|
||||
|
||||
[[package]]
|
||||
name = "libredox"
|
||||
version = "0.1.16"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e02f3bb43d335493c96bf3fd3a321600bf6bd07ed34bc64118e9293bdffea46c"
|
||||
dependencies = [
|
||||
"bitflags 2.11.0",
|
||||
"libc",
|
||||
"plain",
|
||||
"redox_syscall 0.7.4",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "line-clipping"
|
||||
version = "0.3.7"
|
||||
|
|
@ -1319,6 +1462,15 @@ version = "2.8.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79"
|
||||
|
||||
[[package]]
|
||||
name = "memmap2"
|
||||
version = "0.9.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "714098028fe011992e1c3962653c96b2d578c4b4bce9036e15ff220319b1e0e3"
|
||||
dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "memmem"
|
||||
version = "0.1.1"
|
||||
|
|
@ -1519,7 +1671,7 @@ checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1"
|
|||
dependencies = [
|
||||
"cfg-if",
|
||||
"libc",
|
||||
"redox_syscall",
|
||||
"redox_syscall 0.5.18",
|
||||
"smallvec",
|
||||
"windows-link",
|
||||
]
|
||||
|
|
@ -1637,6 +1789,12 @@ version = "0.3.32"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c"
|
||||
|
||||
[[package]]
|
||||
name = "plain"
|
||||
version = "0.2.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b4596b6d070b27117e987119b4dac604f3c58cfb0b191112e24771b2faeac1a6"
|
||||
|
||||
[[package]]
|
||||
name = "pod"
|
||||
version = "0.1.0"
|
||||
|
|
@ -1656,6 +1814,7 @@ dependencies = [
|
|||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"toml",
|
||||
"tools",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
|
|
@ -1849,6 +2008,15 @@ dependencies = [
|
|||
"bitflags 2.11.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "redox_syscall"
|
||||
version = "0.7.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f450ad9c3b1da563fb6948a8e0fb0fb9269711c9c73d9ea1de5058c79c8d643a"
|
||||
dependencies = [
|
||||
"bitflags 2.11.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ref-cast"
|
||||
version = "1.0.25"
|
||||
|
|
@ -2032,6 +2200,15 @@ version = "1.0.23"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f"
|
||||
|
||||
[[package]]
|
||||
name = "same-file"
|
||||
version = "1.0.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502"
|
||||
dependencies = [
|
||||
"winapi-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "scc"
|
||||
version = "2.4.0"
|
||||
|
|
@ -2676,6 +2853,29 @@ version = "1.1.1+spec-1.1.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "756daf9b1013ebe47a8776667b466417e2d4c5679d441c26230efd9ef78692db"
|
||||
|
||||
[[package]]
|
||||
name = "tools"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"filetime",
|
||||
"globset",
|
||||
"grep-matcher",
|
||||
"grep-regex",
|
||||
"grep-searcher",
|
||||
"ignore",
|
||||
"llm-worker",
|
||||
"manifest",
|
||||
"schemars",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2 0.11.0",
|
||||
"tempfile",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tower"
|
||||
version = "0.5.3"
|
||||
|
|
@ -2931,6 +3131,16 @@ dependencies = [
|
|||
"utf8parse",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "walkdir"
|
||||
version = "2.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b"
|
||||
dependencies = [
|
||||
"same-file",
|
||||
"winapi-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "want"
|
||||
version = "0.3.1"
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ members = [
|
|||
"crates/pod",
|
||||
"crates/protocol",
|
||||
"crates/provider",
|
||||
"crates/tools",
|
||||
"crates/tui",
|
||||
]
|
||||
|
||||
|
|
|
|||
2
TODO.md
2
TODO.md
|
|
@ -15,6 +15,8 @@
|
|||
- [x] Hook モジュールの llm-worker からの除去 → [tickets/remove-hook-module.md](tickets/remove-hook-module.md)
|
||||
- [x] api_key_file: ファイルパスによるAPIキー解決 → [tickets/api-key-file.md](tickets/api-key-file.md)
|
||||
- [x] コンテキスト圧縮 (Prune + Compact) → [tickets/context-compaction.md](tickets/context-compaction.md)
|
||||
- [ ] LlmClient へ Tokenizer の導入 → [tickets/token-counter.md](tickets/token-counter.md)
|
||||
- [ ] ToolOutput.referenced_files: ツール参照ファイル追跡 → [tickets/tool-output-referenced-files.md](tickets/tool-output-referenced-files.md)
|
||||
- [ ] Compact の改善(要約品質 + 挙動詳細) → [tickets/compact-improvements.md](tickets/compact-improvements.md)
|
||||
- [ ] Protocol の設計 → [tickets/protocol-design.md](tickets/protocol-design.md)
|
||||
- [x] Protocol: request-response パターン (GetHistory等) → [tickets/request-response-protocol.md](tickets/request-response-protocol.md)
|
||||
|
|
|
|||
|
|
@ -24,17 +24,19 @@ impl Scope {
|
|||
|
||||
/// Check whether `path` falls within this scope.
|
||||
///
|
||||
/// The path is canonicalized before comparison.
|
||||
/// The path is canonicalized before comparison. If the path does not
|
||||
/// exist yet (typical for new-file writes), the closest existing
|
||||
/// ancestor is canonicalized and checked, so deep new directory
|
||||
/// hierarchies inside the scope are also accepted.
|
||||
pub fn contains(&self, path: &Path) -> bool {
|
||||
match path.canonicalize() {
|
||||
Ok(canonical) => canonical.starts_with(&self.root),
|
||||
Err(_) => {
|
||||
// Path doesn't exist yet — check the parent directory instead.
|
||||
// This handles write_file to a new file inside the scope.
|
||||
match path.parent().and_then(|p| p.canonicalize().ok()) {
|
||||
Some(parent) => parent.starts_with(&self.root),
|
||||
None => false,
|
||||
let mut cur = path;
|
||||
loop {
|
||||
if let Ok(canonical) = cur.canonicalize() {
|
||||
return canonical.starts_with(&self.root);
|
||||
}
|
||||
match cur.parent() {
|
||||
Some(parent) if parent != cur => cur = parent,
|
||||
_ => return false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -98,4 +100,16 @@ mod tests {
|
|||
let traversal = dir.path().join("../../../etc/passwd");
|
||||
assert!(!scope.contains(&traversal));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn contains_deeply_nested_new_path() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = Scope::new(dir.path()).unwrap();
|
||||
|
||||
// Neither the file nor any of its ancestors (a, a/b, a/b/c) exist yet
|
||||
// under the scope; contains should still accept because the closest
|
||||
// existing ancestor (the scope root) is inside the scope.
|
||||
let deep = dir.path().join("a/b/c/new.txt");
|
||||
assert!(scope.contains(&deep));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ thiserror = "2.0"
|
|||
tokio = { version = "1.49", features = ["fs", "io-util", "macros", "net", "rt-multi-thread", "signal", "sync"] }
|
||||
toml = "1.1.2"
|
||||
tracing = "0.1.44"
|
||||
tools = { version = "0.1.0", path = "../tools" }
|
||||
|
||||
[dev-dependencies]
|
||||
async-trait = "0.1.89"
|
||||
|
|
|
|||
|
|
@ -83,6 +83,10 @@ impl PodController {
|
|||
// Keep the server alive by moving it into the controller task
|
||||
// (it will be dropped when the task ends)
|
||||
|
||||
// Grab the scope before the mutable borrow of the worker so we can
|
||||
// build a `ScopedFs` for the builtin tools. `Scope` is cheap to clone.
|
||||
let scope_for_tools = pod.scope().cloned();
|
||||
|
||||
// Register event bridge callbacks on the worker
|
||||
{
|
||||
let worker = pod.worker_mut();
|
||||
|
|
@ -155,6 +159,19 @@ impl PodController {
|
|||
message: event.message.clone(),
|
||||
});
|
||||
});
|
||||
|
||||
// Register the builtin file-manipulation tools (Read / Write /
|
||||
// Edit / Glob / Grep) when the manifest declares a scope.
|
||||
//
|
||||
// `ScopedFs` carries the pod-lifetime write boundary (derived
|
||||
// from the manifest scope). `ReadTracker` is session-scoped —
|
||||
// a fresh instance per controller spawn ensures state from a
|
||||
// previous process lifetime cannot be reused after a resume.
|
||||
if let Some(scope) = scope_for_tools {
|
||||
let fs = tools::ScopedFs::new(scope);
|
||||
let tracker = tools::ReadTracker::new();
|
||||
worker.register_tools(tools::builtin_tools(fs, tracker));
|
||||
}
|
||||
}
|
||||
|
||||
// Clone cancel sender before moving pod
|
||||
|
|
|
|||
27
crates/tools/Cargo.toml
Normal file
27
crates/tools/Cargo.toml
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
[package]
|
||||
name = "tools"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
|
||||
[dependencies]
|
||||
async-trait = "0.1.89"
|
||||
globset = "0.4.18"
|
||||
grep-matcher = "0.1.8"
|
||||
grep-regex = "0.1.14"
|
||||
grep-searcher = "0.1.16"
|
||||
ignore = "0.4.25"
|
||||
llm-worker = { version = "0.2.1", path = "../llm-worker" }
|
||||
manifest = { version = "0.1.0", path = "../manifest" }
|
||||
schemars = "1.2.1"
|
||||
serde = { version = "1.0.228", features = ["derive"] }
|
||||
serde_json = "1.0.149"
|
||||
sha2 = "0.11.0"
|
||||
tempfile = "3.27.0"
|
||||
thiserror = "2.0.18"
|
||||
tokio = { version = "1.51.1", features = ["rt"] }
|
||||
tracing = "0.1.44"
|
||||
|
||||
[dev-dependencies]
|
||||
filetime = "0.2.27"
|
||||
tokio = { version = "1.51.1", features = ["macros", "rt-multi-thread"] }
|
||||
292
crates/tools/src/edit.rs
Normal file
292
crates/tools/src/edit.rs
Normal file
|
|
@ -0,0 +1,292 @@
|
|||
//! `Edit` tool — partial string replacement with uniqueness check.
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::ToolsError;
|
||||
use crate::read_tracker::ReadTracker;
|
||||
use crate::scoped_fs::ScopedFs;
|
||||
|
||||
const DESCRIPTION: &str = "Replace a substring in an existing file. By default \
|
||||
`old_string` must be unique in the file; set `replace_all: true` to replace \
|
||||
every occurrence. The file must have been read first (via the Read tool) in \
|
||||
this session. Paths must be absolute.";
|
||||
|
||||
#[derive(Debug, Deserialize, schemars::JsonSchema)]
|
||||
pub(crate) struct EditParams {
|
||||
/// Absolute path to the file.
|
||||
pub file_path: PathBuf,
|
||||
/// String to replace. Must be unique in the file unless `replace_all` is true.
|
||||
pub old_string: String,
|
||||
/// Replacement string. Must differ from `old_string`.
|
||||
pub new_string: String,
|
||||
/// Replace all occurrences. Defaults to false.
|
||||
#[serde(default)]
|
||||
pub replace_all: bool,
|
||||
}
|
||||
|
||||
pub(crate) struct EditTool {
|
||||
fs: ScopedFs,
|
||||
tracker: ReadTracker,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for EditTool {
|
||||
async fn execute(&self, input_json: &str) -> Result<ToolOutput, ToolError> {
|
||||
let params: EditParams = serde_json::from_str(input_json)
|
||||
.map_err(|e| ToolError::InvalidArgument(format!("invalid Edit input: {e}")))?;
|
||||
|
||||
tracing::debug!(
|
||||
path = %params.file_path.display(),
|
||||
replace_all = params.replace_all,
|
||||
"Edit"
|
||||
);
|
||||
|
||||
if params.old_string.is_empty() {
|
||||
return Err(ToolError::InvalidArgument(
|
||||
"old_string must not be empty".into(),
|
||||
));
|
||||
}
|
||||
if params.old_string == params.new_string {
|
||||
return Err(ToolError::InvalidArgument(
|
||||
"old_string and new_string are identical".into(),
|
||||
));
|
||||
}
|
||||
|
||||
// Load current content and verify it matches the recorded hash.
|
||||
let current_bytes = self.fs.read_bytes(¶ms.file_path)?;
|
||||
self.tracker.verify(¶ms.file_path, ¤t_bytes)?;
|
||||
|
||||
let current_text = std::str::from_utf8(¤t_bytes).map_err(|_| {
|
||||
ToolsError::InvalidArgument(format!(
|
||||
"file is not valid UTF-8: {}",
|
||||
params.file_path.display()
|
||||
))
|
||||
})?;
|
||||
|
||||
let count = current_text.matches(¶ms.old_string).count();
|
||||
if count == 0 {
|
||||
return Err(ToolsError::StringNotFound {
|
||||
path: params.file_path.clone(),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
if !params.replace_all && count > 1 {
|
||||
return Err(ToolsError::NotUnique {
|
||||
path: params.file_path.clone(),
|
||||
count,
|
||||
}
|
||||
.into());
|
||||
}
|
||||
|
||||
let new_text = if params.replace_all {
|
||||
current_text.replace(¶ms.old_string, ¶ms.new_string)
|
||||
} else {
|
||||
current_text.replacen(¶ms.old_string, ¶ms.new_string, 1)
|
||||
};
|
||||
let occurrences = if params.replace_all { count } else { 1 };
|
||||
|
||||
self.fs.write(¶ms.file_path, new_text.as_bytes())?;
|
||||
self.tracker
|
||||
.record(¶ms.file_path, new_text.as_bytes());
|
||||
|
||||
let summary = format!(
|
||||
"Edited {} ({} replacement{})",
|
||||
params.file_path.display(),
|
||||
occurrences,
|
||||
if occurrences == 1 { "" } else { "s" }
|
||||
);
|
||||
let preview = make_preview(&new_text, ¶ms.new_string);
|
||||
|
||||
Ok(ToolOutput {
|
||||
summary,
|
||||
content: Some(preview),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a small line-numbered snippet centered on the first occurrence of
|
||||
/// `needle` in `text`. Shows ±3 surrounding lines.
|
||||
fn make_preview(text: &str, needle: &str) -> String {
|
||||
let lines: Vec<&str> = text.lines().collect();
|
||||
if lines.is_empty() {
|
||||
return String::new();
|
||||
}
|
||||
let first_needle_line = needle.lines().next().unwrap_or(needle);
|
||||
let hit = lines
|
||||
.iter()
|
||||
.position(|l| l.contains(first_needle_line))
|
||||
.unwrap_or(0);
|
||||
|
||||
let start = hit.saturating_sub(3);
|
||||
let end = (hit + 4).min(lines.len());
|
||||
|
||||
use std::fmt::Write as _;
|
||||
let mut out = String::new();
|
||||
for (i, line) in lines[start..end].iter().enumerate() {
|
||||
let lineno = start + i + 1;
|
||||
let _ = writeln!(&mut out, "{:>6}\t{}", lineno, line);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Factory for the `Edit` tool.
|
||||
pub fn edit_tool(fs: ScopedFs, tracker: ReadTracker) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(EditParams);
|
||||
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
|
||||
let meta = ToolMeta::new("Edit")
|
||||
.description(DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(EditTool {
|
||||
fs: fs.clone(),
|
||||
tracker: tracker.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::read::read_tool;
|
||||
use manifest::Scope;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn setup() -> (TempDir, ScopedFs, ReadTracker) {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let fs = ScopedFs::new(Scope::new(dir.path()).unwrap());
|
||||
(dir, fs, ReadTracker::new())
|
||||
}
|
||||
|
||||
async fn read_first(fs: &ScopedFs, tracker: &ReadTracker, file: &std::path::Path) {
|
||||
let def = read_tool(fs.clone(), tracker.clone());
|
||||
let (_, reader) = def();
|
||||
let inp = serde_json::json!({ "file_path": file.to_str().unwrap() });
|
||||
reader.execute(&inp.to_string()).await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn edit_unique_replacement() {
|
||||
let (dir, fs, tracker) = setup();
|
||||
let file = dir.path().join("a.txt");
|
||||
std::fs::write(&file, "line1\nfoo bar\nline3\n").unwrap();
|
||||
read_first(&fs, &tracker, &file).await;
|
||||
|
||||
let def = edit_tool(fs, tracker);
|
||||
let (meta, tool) = def();
|
||||
assert_eq!(meta.name, "Edit");
|
||||
|
||||
let inp = serde_json::json!({
|
||||
"file_path": file.to_str().unwrap(),
|
||||
"old_string": "foo bar",
|
||||
"new_string": "foo baz",
|
||||
});
|
||||
let out = tool.execute(&inp.to_string()).await.unwrap();
|
||||
assert!(out.summary.contains("1 replacement"));
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(&file).unwrap(),
|
||||
"line1\nfoo baz\nline3\n"
|
||||
);
|
||||
assert!(out.content.unwrap().contains("foo baz"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn edit_replace_all() {
|
||||
let (dir, fs, tracker) = setup();
|
||||
let file = dir.path().join("a.txt");
|
||||
std::fs::write(&file, "x x x\n").unwrap();
|
||||
read_first(&fs, &tracker, &file).await;
|
||||
|
||||
let def = edit_tool(fs, tracker);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({
|
||||
"file_path": file.to_str().unwrap(),
|
||||
"old_string": "x",
|
||||
"new_string": "y",
|
||||
"replace_all": true,
|
||||
});
|
||||
let out = tool.execute(&inp.to_string()).await.unwrap();
|
||||
assert!(out.summary.contains("3 replacements"));
|
||||
assert_eq!(std::fs::read_to_string(&file).unwrap(), "y y y\n");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn edit_not_unique() {
|
||||
let (dir, fs, tracker) = setup();
|
||||
let file = dir.path().join("a.txt");
|
||||
std::fs::write(&file, "a a\n").unwrap();
|
||||
read_first(&fs, &tracker, &file).await;
|
||||
|
||||
let def = edit_tool(fs, tracker);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({
|
||||
"file_path": file.to_str().unwrap(),
|
||||
"old_string": "a",
|
||||
"new_string": "b",
|
||||
});
|
||||
let err = tool.execute(&inp.to_string()).await.unwrap_err();
|
||||
assert!(matches!(err, ToolError::InvalidArgument(_)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn edit_string_not_found() {
|
||||
let (dir, fs, tracker) = setup();
|
||||
let file = dir.path().join("a.txt");
|
||||
std::fs::write(&file, "hello\n").unwrap();
|
||||
read_first(&fs, &tracker, &file).await;
|
||||
|
||||
let def = edit_tool(fs, tracker);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({
|
||||
"file_path": file.to_str().unwrap(),
|
||||
"old_string": "world",
|
||||
"new_string": "x",
|
||||
});
|
||||
let err = tool.execute(&inp.to_string()).await.unwrap_err();
|
||||
assert!(matches!(err, ToolError::InvalidArgument(_)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn edit_requires_prior_read() {
|
||||
let (dir, fs, tracker) = setup();
|
||||
let file = dir.path().join("a.txt");
|
||||
std::fs::write(&file, "foo\n").unwrap();
|
||||
|
||||
let def = edit_tool(fs, tracker);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({
|
||||
"file_path": file.to_str().unwrap(),
|
||||
"old_string": "foo",
|
||||
"new_string": "bar",
|
||||
});
|
||||
let err = tool.execute(&inp.to_string()).await.unwrap_err();
|
||||
assert!(matches!(err, ToolError::InvalidArgument(_)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn edit_detects_external_modification() {
|
||||
let (dir, fs, tracker) = setup();
|
||||
let file = dir.path().join("a.txt");
|
||||
std::fs::write(&file, "foo\n").unwrap();
|
||||
read_first(&fs, &tracker, &file).await;
|
||||
|
||||
// External tampering between read and edit
|
||||
std::fs::write(&file, "something else").unwrap();
|
||||
|
||||
let def = edit_tool(fs, tracker);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({
|
||||
"file_path": file.to_str().unwrap(),
|
||||
"old_string": "foo",
|
||||
"new_string": "bar",
|
||||
});
|
||||
let err = tool.execute(&inp.to_string()).await.unwrap_err();
|
||||
let msg = format!("{err}");
|
||||
assert!(msg.contains("modified externally"), "{msg}");
|
||||
}
|
||||
}
|
||||
85
crates/tools/src/error.rs
Normal file
85
crates/tools/src/error.rs
Normal file
|
|
@ -0,0 +1,85 @@
|
|||
//! Error type shared across the `tools` crate.
|
||||
//!
|
||||
//! `ToolsError` is the crate-level error returned by `ScopedFs` and each
|
||||
//! builtin tool's internal logic. Tool `execute()` impls convert it to
|
||||
//! [`llm_worker::tool::ToolError`] via the `From` impl defined here.
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use llm_worker::tool::ToolError;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum ToolsError {
|
||||
#[error("path must be absolute: {}", .0.display())]
|
||||
RelativePath(PathBuf),
|
||||
|
||||
#[error("path is outside allowed scope: {}", .0.display())]
|
||||
OutOfScope(PathBuf),
|
||||
|
||||
#[error("path is a directory: {}", .0.display())]
|
||||
IsDirectory(PathBuf),
|
||||
|
||||
#[error("file not found: {}", .0.display())]
|
||||
NotFound(PathBuf),
|
||||
|
||||
#[error("file has not been read in this session; read it first: {}", .0.display())]
|
||||
NotRead(PathBuf),
|
||||
|
||||
#[error("file was modified externally after last read: {}", .0.display())]
|
||||
ExternallyModified(PathBuf),
|
||||
|
||||
#[error("string not found in file: {}", .path.display())]
|
||||
StringNotFound { path: PathBuf },
|
||||
|
||||
#[error(
|
||||
"string is not unique in file ({count} occurrences); pass replace_all=true or disambiguate: {}",
|
||||
.path.display()
|
||||
)]
|
||||
NotUnique { path: PathBuf, count: usize },
|
||||
|
||||
#[error("invalid argument: {0}")]
|
||||
InvalidArgument(String),
|
||||
|
||||
#[error("invalid regex: {0}")]
|
||||
InvalidRegex(String),
|
||||
|
||||
#[error("invalid glob pattern: {0}")]
|
||||
InvalidGlob(String),
|
||||
|
||||
#[error("I/O error at {}: {source}", .path.display())]
|
||||
Io {
|
||||
path: PathBuf,
|
||||
#[source]
|
||||
source: std::io::Error,
|
||||
},
|
||||
}
|
||||
|
||||
impl ToolsError {
|
||||
/// Helper to wrap an [`std::io::Error`] with the path it occurred on.
|
||||
pub fn io(path: impl Into<PathBuf>, source: std::io::Error) -> Self {
|
||||
Self::Io {
|
||||
path: path.into(),
|
||||
source,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ToolsError> for ToolError {
|
||||
fn from(err: ToolsError) -> Self {
|
||||
use ToolsError::*;
|
||||
match err {
|
||||
RelativePath(_)
|
||||
| OutOfScope(_)
|
||||
| IsDirectory(_)
|
||||
| NotRead(_)
|
||||
| ExternallyModified(_)
|
||||
| StringNotFound { .. }
|
||||
| NotUnique { .. }
|
||||
| InvalidArgument(_)
|
||||
| InvalidRegex(_)
|
||||
| InvalidGlob(_) => ToolError::InvalidArgument(err.to_string()),
|
||||
NotFound(_) => ToolError::ExecutionFailed(err.to_string()),
|
||||
Io { .. } => ToolError::ExecutionFailed(err.to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
254
crates/tools/src/glob.rs
Normal file
254
crates/tools/src/glob.rs
Normal file
|
|
@ -0,0 +1,254 @@
|
|||
//! `Glob` tool — recursive file search by glob pattern, sorted by mtime.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::time::SystemTime;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::ToolsError;
|
||||
use crate::scoped_fs::ScopedFs;
|
||||
|
||||
const DESCRIPTION: &str = "Recursively find files matching a glob pattern \
|
||||
(e.g. \"**/*.rs\"). Results are sorted by modification time, newest first, \
|
||||
and capped at 1000 entries. Hidden files are included. The `path` parameter \
|
||||
defaults to the scope root when omitted. Paths must be absolute.";
|
||||
|
||||
const RESULT_LIMIT: usize = 1000;
|
||||
|
||||
#[derive(Debug, Deserialize, schemars::JsonSchema)]
|
||||
pub(crate) struct GlobParams {
|
||||
/// Glob pattern, e.g. `"**/*.rs"`. Matched against paths relative to
|
||||
/// `path` (or the scope root if omitted).
|
||||
pub pattern: String,
|
||||
/// Absolute directory to search under. Defaults to the scope root.
|
||||
#[serde(default)]
|
||||
pub path: Option<PathBuf>,
|
||||
}
|
||||
|
||||
pub(crate) struct GlobTool {
|
||||
fs: ScopedFs,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for GlobTool {
|
||||
async fn execute(&self, input_json: &str) -> Result<ToolOutput, ToolError> {
|
||||
let params: GlobParams = serde_json::from_str(input_json)
|
||||
.map_err(|e| ToolError::InvalidArgument(format!("invalid Glob input: {e}")))?;
|
||||
|
||||
tracing::debug!(
|
||||
pattern = %params.pattern,
|
||||
path = ?params.path,
|
||||
"Glob"
|
||||
);
|
||||
|
||||
let base = params
|
||||
.path
|
||||
.clone()
|
||||
.unwrap_or_else(|| self.fs.scope().root().to_path_buf());
|
||||
let pattern = params.pattern.clone();
|
||||
|
||||
// ignore::Walk is synchronous; run it on a blocking thread so we
|
||||
// don't stall the runtime for large trees.
|
||||
let results = tokio::task::spawn_blocking(move || run_glob(&base, &pattern))
|
||||
.await
|
||||
.map_err(|e| ToolError::Internal(format!("spawn_blocking failed: {e}")))??;
|
||||
|
||||
let total = results.len();
|
||||
let (shown, truncated) = if total > RESULT_LIMIT {
|
||||
(&results[..RESULT_LIMIT], true)
|
||||
} else {
|
||||
(&results[..], false)
|
||||
};
|
||||
|
||||
if shown.is_empty() {
|
||||
return Ok(ToolOutput {
|
||||
summary: format!("No files found matching {}", params.pattern),
|
||||
content: None,
|
||||
});
|
||||
}
|
||||
|
||||
let mut body = String::new();
|
||||
for p in shown {
|
||||
body.push_str(&p.display().to_string());
|
||||
body.push('\n');
|
||||
}
|
||||
|
||||
let summary = if truncated {
|
||||
format!(
|
||||
"Found {total}+ files matching {} (truncated to {RESULT_LIMIT})",
|
||||
params.pattern
|
||||
)
|
||||
} else {
|
||||
format!("Found {total} file(s) matching {}", params.pattern)
|
||||
};
|
||||
|
||||
Ok(ToolOutput {
|
||||
summary,
|
||||
content: Some(body),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn run_glob(base: &Path, pattern: &str) -> Result<Vec<PathBuf>, ToolsError> {
|
||||
if !base.is_absolute() {
|
||||
return Err(ToolsError::RelativePath(base.to_path_buf()));
|
||||
}
|
||||
if !base.exists() {
|
||||
return Err(ToolsError::NotFound(base.to_path_buf()));
|
||||
}
|
||||
|
||||
let glob = globset::Glob::new(pattern)
|
||||
.map_err(|e| ToolsError::InvalidGlob(e.to_string()))?
|
||||
.compile_matcher();
|
||||
|
||||
// Glob is an explicit-pattern tool, so gitignore/hidden are *not* honored.
|
||||
let walker = ignore::WalkBuilder::new(base)
|
||||
.hidden(false)
|
||||
.git_ignore(false)
|
||||
.git_global(false)
|
||||
.git_exclude(false)
|
||||
.ignore(false)
|
||||
.parents(false)
|
||||
.follow_links(false)
|
||||
.build();
|
||||
|
||||
let mut hits: Vec<(PathBuf, SystemTime)> = Vec::new();
|
||||
for entry in walker.flatten() {
|
||||
let ft = match entry.file_type() {
|
||||
Some(ft) => ft,
|
||||
None => continue,
|
||||
};
|
||||
if !ft.is_file() {
|
||||
continue;
|
||||
}
|
||||
let rel = match entry.path().strip_prefix(base) {
|
||||
Ok(r) => r,
|
||||
Err(_) => continue,
|
||||
};
|
||||
if !glob.is_match(rel) {
|
||||
continue;
|
||||
}
|
||||
let mtime = entry
|
||||
.metadata()
|
||||
.ok()
|
||||
.and_then(|m| m.modified().ok())
|
||||
.unwrap_or(SystemTime::UNIX_EPOCH);
|
||||
hits.push((entry.path().to_path_buf(), mtime));
|
||||
}
|
||||
|
||||
hits.sort_by(|a, b| b.1.cmp(&a.1));
|
||||
Ok(hits.into_iter().map(|(p, _)| p).collect())
|
||||
}
|
||||
|
||||
/// Factory for the `Glob` tool.
|
||||
pub fn glob_tool(fs: ScopedFs) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(GlobParams);
|
||||
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
|
||||
let meta = ToolMeta::new("Glob")
|
||||
.description(DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(GlobTool { fs: fs.clone() });
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use manifest::Scope;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn setup() -> (TempDir, ScopedFs) {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let fs = ScopedFs::new(Scope::new(dir.path()).unwrap());
|
||||
(dir, fs)
|
||||
}
|
||||
|
||||
fn touch(path: &Path, content: &str) {
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent).unwrap();
|
||||
}
|
||||
std::fs::write(path, content).unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn glob_finds_matching_files() {
|
||||
let (dir, fs) = setup();
|
||||
touch(&dir.path().join("a.rs"), "");
|
||||
touch(&dir.path().join("sub/b.rs"), "");
|
||||
touch(&dir.path().join("sub/c.txt"), "");
|
||||
|
||||
let def = glob_tool(fs);
|
||||
let (meta, tool) = def();
|
||||
assert_eq!(meta.name, "Glob");
|
||||
|
||||
let inp = serde_json::json!({ "pattern": "**/*.rs" });
|
||||
let out = tool.execute(&inp.to_string()).await.unwrap();
|
||||
assert!(out.summary.contains("2 file(s)"));
|
||||
let body = out.content.unwrap();
|
||||
assert!(body.contains("a.rs"));
|
||||
assert!(body.contains("b.rs"));
|
||||
assert!(!body.contains("c.txt"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn glob_sorts_by_mtime_desc() {
|
||||
let (dir, fs) = setup();
|
||||
let older = dir.path().join("old.rs");
|
||||
let newer = dir.path().join("new.rs");
|
||||
touch(&older, "");
|
||||
touch(&newer, "");
|
||||
|
||||
filetime::set_file_mtime(&older, filetime::FileTime::from_unix_time(1_000, 0)).unwrap();
|
||||
filetime::set_file_mtime(&newer, filetime::FileTime::from_unix_time(2_000, 0)).unwrap();
|
||||
|
||||
let def = glob_tool(fs);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({ "pattern": "*.rs" });
|
||||
let out = tool.execute(&inp.to_string()).await.unwrap();
|
||||
let body = out.content.unwrap();
|
||||
let new_pos = body.find("new.rs").unwrap();
|
||||
let old_pos = body.find("old.rs").unwrap();
|
||||
assert!(new_pos < old_pos, "newer file should come first:\n{body}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn glob_empty_results() {
|
||||
let (_dir, fs) = setup();
|
||||
let def = glob_tool(fs);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({ "pattern": "**/*.nonexistent" });
|
||||
let out = tool.execute(&inp.to_string()).await.unwrap();
|
||||
assert!(out.summary.contains("No files"));
|
||||
assert!(out.content.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn glob_invalid_pattern() {
|
||||
let (_dir, fs) = setup();
|
||||
let def = glob_tool(fs);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({ "pattern": "[unterminated" });
|
||||
let err = tool.execute(&inp.to_string()).await.unwrap_err();
|
||||
assert!(matches!(err, ToolError::InvalidArgument(_)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn glob_honors_hidden_files() {
|
||||
let (dir, fs) = setup();
|
||||
touch(&dir.path().join(".hidden.rs"), "");
|
||||
touch(&dir.path().join("visible.rs"), "");
|
||||
|
||||
let def = glob_tool(fs);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({ "pattern": "*.rs" });
|
||||
let out = tool.execute(&inp.to_string()).await.unwrap();
|
||||
let body = out.content.unwrap();
|
||||
assert!(body.contains(".hidden.rs"));
|
||||
assert!(body.contains("visible.rs"));
|
||||
}
|
||||
}
|
||||
751
crates/tools/src/grep.rs
Normal file
751
crates/tools/src/grep.rs
Normal file
|
|
@ -0,0 +1,751 @@
|
|||
//! `Grep` tool — recursive regex search powered by ripgrep's component crates.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use grep_regex::RegexMatcherBuilder;
|
||||
use grep_searcher::sinks::UTF8 as UTF8Sink;
|
||||
use grep_searcher::{
|
||||
BinaryDetection, Searcher, SearcherBuilder, Sink, SinkContext, SinkMatch,
|
||||
};
|
||||
use ignore::WalkBuilder;
|
||||
use ignore::overrides::OverrideBuilder;
|
||||
use ignore::types::TypesBuilder;
|
||||
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::ToolsError;
|
||||
use crate::scoped_fs::ScopedFs;
|
||||
|
||||
const DESCRIPTION: &str = "Recursive regex search across files, powered by \
|
||||
ripgrep. Supports file filtering (`glob`, `type`), context lines, multiline \
|
||||
matching, and three output modes: `files_with_matches` (default), `content`, \
|
||||
and `count`. Honors .gitignore. Binary files are skipped. Paths must be \
|
||||
absolute.";
|
||||
|
||||
const DEFAULT_HEAD_LIMIT: usize = 250;
|
||||
|
||||
#[derive(Debug, Clone, Copy, Deserialize, schemars::JsonSchema, Default, PartialEq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub(crate) enum GrepOutputMode {
|
||||
#[default]
|
||||
FilesWithMatches,
|
||||
Content,
|
||||
Count,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, schemars::JsonSchema)]
|
||||
pub(crate) struct GrepParams {
|
||||
/// Regex pattern to search for.
|
||||
pub pattern: String,
|
||||
/// Absolute path to search under. Defaults to the scope root.
|
||||
#[serde(default)]
|
||||
pub path: Option<PathBuf>,
|
||||
/// Glob filter applied to candidate files, e.g. `"*.rs"`.
|
||||
#[serde(default)]
|
||||
pub glob: Option<String>,
|
||||
/// File type filter, e.g. `"rust"` or `"py"`. See ripgrep's default types.
|
||||
#[serde(default, rename = "type")]
|
||||
pub file_type: Option<String>,
|
||||
/// Output mode: `files_with_matches` (default), `content`, or `count`.
|
||||
#[serde(default)]
|
||||
pub output_mode: Option<GrepOutputMode>,
|
||||
/// Show line numbers in content mode. Defaults to true.
|
||||
#[serde(default, rename = "-n")]
|
||||
pub line_numbers: Option<bool>,
|
||||
/// Case-insensitive matching.
|
||||
#[serde(default, rename = "-i")]
|
||||
pub case_insensitive: bool,
|
||||
/// Trailing context lines after each match.
|
||||
#[serde(default, rename = "-A")]
|
||||
pub after: Option<usize>,
|
||||
/// Leading context lines before each match.
|
||||
#[serde(default, rename = "-B")]
|
||||
pub before: Option<usize>,
|
||||
/// Context lines before AND after each match (overrides -A/-B when set).
|
||||
#[serde(default, rename = "-C")]
|
||||
pub context: Option<usize>,
|
||||
/// Allow patterns to match across newlines.
|
||||
#[serde(default)]
|
||||
pub multiline: bool,
|
||||
/// Maximum number of output entries. Defaults to 250.
|
||||
#[serde(default)]
|
||||
pub head_limit: Option<usize>,
|
||||
/// Skip the first N output entries (pagination).
|
||||
#[serde(default)]
|
||||
pub offset: Option<usize>,
|
||||
}
|
||||
|
||||
pub(crate) struct GrepTool {
|
||||
fs: ScopedFs,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for GrepTool {
|
||||
async fn execute(&self, input_json: &str) -> Result<ToolOutput, ToolError> {
|
||||
let params: GrepParams = serde_json::from_str(input_json)
|
||||
.map_err(|e| ToolError::InvalidArgument(format!("invalid Grep input: {e}")))?;
|
||||
|
||||
tracing::debug!(
|
||||
pattern = %params.pattern,
|
||||
mode = ?params.output_mode,
|
||||
"Grep"
|
||||
);
|
||||
|
||||
let default_base = self.fs.scope().root().to_path_buf();
|
||||
let report =
|
||||
tokio::task::spawn_blocking(move || run_grep(default_base, params))
|
||||
.await
|
||||
.map_err(|e| ToolError::Internal(format!("spawn_blocking failed: {e}")))??;
|
||||
|
||||
Ok(report.render())
|
||||
}
|
||||
}
|
||||
|
||||
/// Factory for the `Grep` tool.
|
||||
pub fn grep_tool(fs: ScopedFs) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(GrepParams);
|
||||
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
|
||||
let meta = ToolMeta::new("Grep")
|
||||
.description(DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(GrepTool { fs: fs.clone() });
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Implementation
|
||||
// =============================================================================
|
||||
|
||||
struct ContentLine {
|
||||
path: PathBuf,
|
||||
line_number: Option<u64>,
|
||||
text: String,
|
||||
is_match: bool,
|
||||
}
|
||||
|
||||
struct GrepReport {
|
||||
mode: GrepOutputMode,
|
||||
show_line_numbers: bool,
|
||||
files: Vec<PathBuf>,
|
||||
counts: Vec<(PathBuf, usize)>,
|
||||
lines: Vec<ContentLine>,
|
||||
truncated: bool,
|
||||
head_limit: usize,
|
||||
}
|
||||
|
||||
impl GrepReport {
|
||||
fn render(self) -> ToolOutput {
|
||||
match self.mode {
|
||||
GrepOutputMode::FilesWithMatches => {
|
||||
if self.files.is_empty() {
|
||||
return ToolOutput {
|
||||
summary: "No files matched".into(),
|
||||
content: None,
|
||||
};
|
||||
}
|
||||
let mut body = String::new();
|
||||
for p in &self.files {
|
||||
body.push_str(&p.display().to_string());
|
||||
body.push('\n');
|
||||
}
|
||||
let mut summary = format!("Found matches in {} file(s)", self.files.len());
|
||||
if self.truncated {
|
||||
summary.push_str(&format!(" (truncated at {})", self.head_limit));
|
||||
}
|
||||
ToolOutput {
|
||||
summary,
|
||||
content: Some(body),
|
||||
}
|
||||
}
|
||||
GrepOutputMode::Count => {
|
||||
if self.counts.is_empty() {
|
||||
return ToolOutput {
|
||||
summary: "No files matched".into(),
|
||||
content: None,
|
||||
};
|
||||
}
|
||||
let total_lines: usize = self.counts.iter().map(|(_, n)| *n).sum();
|
||||
let mut body = String::new();
|
||||
for (p, n) in &self.counts {
|
||||
body.push_str(&format!("{}:{}\n", p.display(), n));
|
||||
}
|
||||
let mut summary = format!(
|
||||
"Found matches in {} file(s), {} total line(s)",
|
||||
self.counts.len(),
|
||||
total_lines
|
||||
);
|
||||
if self.truncated {
|
||||
summary.push_str(&format!(" (truncated at {})", self.head_limit));
|
||||
}
|
||||
ToolOutput {
|
||||
summary,
|
||||
content: Some(body),
|
||||
}
|
||||
}
|
||||
GrepOutputMode::Content => {
|
||||
if self.lines.is_empty() {
|
||||
return ToolOutput {
|
||||
summary: "No matches".into(),
|
||||
content: None,
|
||||
};
|
||||
}
|
||||
let match_count = self.lines.iter().filter(|l| l.is_match).count();
|
||||
let file_set: std::collections::BTreeSet<&Path> =
|
||||
self.lines.iter().map(|l| l.path.as_path()).collect();
|
||||
let mut body = String::new();
|
||||
for line in &self.lines {
|
||||
let sep = if line.is_match { ':' } else { '-' };
|
||||
if self.show_line_numbers {
|
||||
if let Some(n) = line.line_number {
|
||||
body.push_str(&format!(
|
||||
"{}{}{}{}{}\n",
|
||||
line.path.display(),
|
||||
sep,
|
||||
n,
|
||||
sep,
|
||||
line.text
|
||||
));
|
||||
continue;
|
||||
}
|
||||
}
|
||||
body.push_str(&format!(
|
||||
"{}{}{}\n",
|
||||
line.path.display(),
|
||||
sep,
|
||||
line.text
|
||||
));
|
||||
}
|
||||
let mut summary = format!(
|
||||
"{} matching line(s) in {} file(s)",
|
||||
match_count,
|
||||
file_set.len()
|
||||
);
|
||||
if self.truncated {
|
||||
summary.push_str(&format!(" (truncated at {})", self.head_limit));
|
||||
}
|
||||
ToolOutput {
|
||||
summary,
|
||||
content: Some(body),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn run_grep(default_base: PathBuf, p: GrepParams) -> Result<GrepReport, ToolsError> {
|
||||
let matcher = RegexMatcherBuilder::new()
|
||||
.case_insensitive(p.case_insensitive)
|
||||
.multi_line(p.multiline)
|
||||
.dot_matches_new_line(p.multiline)
|
||||
.build(&p.pattern)
|
||||
.map_err(|e| ToolsError::InvalidRegex(e.to_string()))?;
|
||||
|
||||
let (before, after) = match (p.before, p.after, p.context) {
|
||||
(_, _, Some(c)) => (c, c),
|
||||
(b, a, None) => (b.unwrap_or(0), a.unwrap_or(0)),
|
||||
};
|
||||
|
||||
let mut sb = SearcherBuilder::new();
|
||||
sb.binary_detection(BinaryDetection::quit(b'\x00'))
|
||||
.line_number(p.line_numbers.unwrap_or(true))
|
||||
.multi_line(p.multiline)
|
||||
.before_context(before)
|
||||
.after_context(after);
|
||||
let mut searcher = sb.build();
|
||||
|
||||
let base = p.path.unwrap_or(default_base);
|
||||
if !base.is_absolute() {
|
||||
return Err(ToolsError::RelativePath(base));
|
||||
}
|
||||
if !base.exists() {
|
||||
return Err(ToolsError::NotFound(base));
|
||||
}
|
||||
|
||||
let mut wb = WalkBuilder::new(&base);
|
||||
wb.hidden(true)
|
||||
.git_ignore(true)
|
||||
.git_global(true)
|
||||
.git_exclude(true)
|
||||
.ignore(true)
|
||||
.parents(true)
|
||||
.follow_links(false);
|
||||
|
||||
if let Some(t) = p.file_type.as_deref() {
|
||||
let mut tb = TypesBuilder::new();
|
||||
tb.add_defaults();
|
||||
tb.select(t);
|
||||
let types = tb
|
||||
.build()
|
||||
.map_err(|e| ToolsError::InvalidArgument(format!("invalid type {t}: {e}")))?;
|
||||
wb.types(types);
|
||||
}
|
||||
if let Some(g) = p.glob.as_deref() {
|
||||
let mut ob = OverrideBuilder::new(&base);
|
||||
ob.add(g).map_err(|e| ToolsError::InvalidGlob(e.to_string()))?;
|
||||
let ov = ob
|
||||
.build()
|
||||
.map_err(|e| ToolsError::InvalidGlob(e.to_string()))?;
|
||||
wb.overrides(ov);
|
||||
}
|
||||
|
||||
let mode = p.output_mode.unwrap_or_default();
|
||||
let head_limit = p.head_limit.unwrap_or(DEFAULT_HEAD_LIMIT);
|
||||
let offset = p.offset.unwrap_or(0);
|
||||
let show_line_numbers = p.line_numbers.unwrap_or(true);
|
||||
|
||||
let mut report = GrepReport {
|
||||
mode,
|
||||
show_line_numbers,
|
||||
files: Vec::new(),
|
||||
counts: Vec::new(),
|
||||
lines: Vec::new(),
|
||||
truncated: false,
|
||||
head_limit,
|
||||
};
|
||||
|
||||
// Per-mode walker state.
|
||||
let mut matching_files_seen: usize = 0;
|
||||
let mut matches_seen: usize = 0;
|
||||
|
||||
'walker: for entry in wb.build().flatten() {
|
||||
if !entry.file_type().map(|t| t.is_file()).unwrap_or(false) {
|
||||
continue;
|
||||
}
|
||||
let path = entry.path();
|
||||
|
||||
match mode {
|
||||
GrepOutputMode::FilesWithMatches => {
|
||||
let hit = scan_any_match(&mut searcher, &matcher, path)?;
|
||||
if !hit {
|
||||
continue;
|
||||
}
|
||||
if matching_files_seen >= offset {
|
||||
report.files.push(path.to_path_buf());
|
||||
if report.files.len() >= head_limit {
|
||||
report.truncated = true;
|
||||
break 'walker;
|
||||
}
|
||||
}
|
||||
matching_files_seen += 1;
|
||||
}
|
||||
GrepOutputMode::Count => {
|
||||
let count = scan_count(&mut searcher, &matcher, path)?;
|
||||
if count == 0 {
|
||||
continue;
|
||||
}
|
||||
if matching_files_seen >= offset {
|
||||
report.counts.push((path.to_path_buf(), count));
|
||||
if report.counts.len() >= head_limit {
|
||||
report.truncated = true;
|
||||
break 'walker;
|
||||
}
|
||||
}
|
||||
matching_files_seen += 1;
|
||||
}
|
||||
GrepOutputMode::Content => {
|
||||
let before_count = matches_seen;
|
||||
let mut sink = ContentSink {
|
||||
path: path.to_path_buf(),
|
||||
lines: &mut report.lines,
|
||||
matches_seen: &mut matches_seen,
|
||||
offset,
|
||||
head_limit,
|
||||
};
|
||||
searcher
|
||||
.search_path(&matcher, path, &mut sink)
|
||||
.map_err(|e| ToolsError::io(path, e))?;
|
||||
// If we hit head_limit during this file, stop walking.
|
||||
if matches_seen >= offset.saturating_add(head_limit) && matches_seen > before_count
|
||||
{
|
||||
report.truncated = true;
|
||||
break 'walker;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(report)
|
||||
}
|
||||
|
||||
fn scan_any_match(
|
||||
searcher: &mut Searcher,
|
||||
matcher: &grep_regex::RegexMatcher,
|
||||
path: &Path,
|
||||
) -> Result<bool, ToolsError> {
|
||||
let mut hit = false;
|
||||
let sink = UTF8Sink(|_, _| {
|
||||
hit = true;
|
||||
Ok(false) // stop searching this file immediately
|
||||
});
|
||||
searcher
|
||||
.search_path(matcher, path, sink)
|
||||
.map_err(|e| ToolsError::io(path, e))?;
|
||||
Ok(hit)
|
||||
}
|
||||
|
||||
fn scan_count(
|
||||
searcher: &mut Searcher,
|
||||
matcher: &grep_regex::RegexMatcher,
|
||||
path: &Path,
|
||||
) -> Result<usize, ToolsError> {
|
||||
let mut count = 0usize;
|
||||
let sink = UTF8Sink(|_, _| {
|
||||
count += 1;
|
||||
Ok(true)
|
||||
});
|
||||
searcher
|
||||
.search_path(matcher, path, sink)
|
||||
.map_err(|e| ToolsError::io(path, e))?;
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
struct ContentSink<'a> {
|
||||
path: PathBuf,
|
||||
lines: &'a mut Vec<ContentLine>,
|
||||
matches_seen: &'a mut usize,
|
||||
offset: usize,
|
||||
head_limit: usize,
|
||||
}
|
||||
|
||||
impl Sink for ContentSink<'_> {
|
||||
type Error = std::io::Error;
|
||||
|
||||
fn matched(
|
||||
&mut self,
|
||||
_searcher: &Searcher,
|
||||
mat: &SinkMatch<'_>,
|
||||
) -> Result<bool, Self::Error> {
|
||||
let idx = *self.matches_seen;
|
||||
*self.matches_seen += 1;
|
||||
|
||||
// Skip matches before offset.
|
||||
if idx < self.offset {
|
||||
return Ok(true);
|
||||
}
|
||||
// Stop searching this file once we've filled the head_limit.
|
||||
if idx >= self.offset.saturating_add(self.head_limit) {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let text = String::from_utf8_lossy(mat.bytes())
|
||||
.trim_end_matches('\n')
|
||||
.trim_end_matches('\r')
|
||||
.to_string();
|
||||
self.lines.push(ContentLine {
|
||||
path: self.path.clone(),
|
||||
line_number: mat.line_number(),
|
||||
text,
|
||||
is_match: true,
|
||||
});
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
fn context(
|
||||
&mut self,
|
||||
_searcher: &Searcher,
|
||||
ctx: &SinkContext<'_>,
|
||||
) -> Result<bool, Self::Error> {
|
||||
let seen = *self.matches_seen;
|
||||
if seen < self.offset {
|
||||
return Ok(true);
|
||||
}
|
||||
if seen >= self.offset.saturating_add(self.head_limit) {
|
||||
return Ok(false);
|
||||
}
|
||||
let text = String::from_utf8_lossy(ctx.bytes())
|
||||
.trim_end_matches('\n')
|
||||
.trim_end_matches('\r')
|
||||
.to_string();
|
||||
self.lines.push(ContentLine {
|
||||
path: self.path.clone(),
|
||||
line_number: ctx.line_number(),
|
||||
text,
|
||||
is_match: false,
|
||||
});
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Tests
|
||||
// =============================================================================
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use manifest::Scope;
|
||||
use std::fs;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn setup() -> (TempDir, ScopedFs) {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let fs = ScopedFs::new(Scope::new(dir.path()).unwrap());
|
||||
(dir, fs)
|
||||
}
|
||||
|
||||
fn touch(path: &Path, content: &str) {
|
||||
if let Some(parent) = path.parent() {
|
||||
fs::create_dir_all(parent).unwrap();
|
||||
}
|
||||
fs::write(path, content).unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grep_files_with_matches_default() {
|
||||
let (dir, fs) = setup();
|
||||
touch(&dir.path().join("a.txt"), "alpha\nbravo\n");
|
||||
touch(&dir.path().join("b.txt"), "charlie\n");
|
||||
|
||||
let def = grep_tool(fs);
|
||||
let (meta, tool) = def();
|
||||
assert_eq!(meta.name, "Grep");
|
||||
|
||||
let inp = serde_json::json!({ "pattern": "bravo" });
|
||||
let out = tool.execute(&inp.to_string()).await.unwrap();
|
||||
assert!(out.summary.contains("1 file"));
|
||||
assert!(out.content.unwrap().contains("a.txt"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grep_content_mode_with_line_numbers() {
|
||||
let (dir, fs) = setup();
|
||||
touch(&dir.path().join("a.txt"), "one\ntwo\nthree\n");
|
||||
|
||||
let def = grep_tool(fs);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({
|
||||
"pattern": "two",
|
||||
"output_mode": "content",
|
||||
});
|
||||
let out = tool.execute(&inp.to_string()).await.unwrap();
|
||||
let body = out.content.unwrap();
|
||||
assert!(body.contains(":2:two"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grep_count_mode() {
|
||||
let (dir, fs) = setup();
|
||||
touch(&dir.path().join("a.txt"), "x\nx\nx\n");
|
||||
touch(&dir.path().join("b.txt"), "x\ny\n");
|
||||
|
||||
let def = grep_tool(fs);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({
|
||||
"pattern": "x",
|
||||
"output_mode": "count",
|
||||
});
|
||||
let out = tool.execute(&inp.to_string()).await.unwrap();
|
||||
let body = out.content.unwrap();
|
||||
assert!(body.contains("a.txt:3"));
|
||||
assert!(body.contains("b.txt:1"));
|
||||
assert!(out.summary.contains("4 total"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grep_case_insensitive() {
|
||||
let (dir, fs) = setup();
|
||||
touch(&dir.path().join("a.txt"), "HELLO\n");
|
||||
|
||||
let def = grep_tool(fs);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({
|
||||
"pattern": "hello",
|
||||
"-i": true,
|
||||
"output_mode": "content",
|
||||
});
|
||||
let out = tool.execute(&inp.to_string()).await.unwrap();
|
||||
assert!(out.content.unwrap().contains("HELLO"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grep_context_lines() {
|
||||
let (dir, fs) = setup();
|
||||
touch(
|
||||
&dir.path().join("a.txt"),
|
||||
"line1\nline2\nMATCH\nline4\nline5\n",
|
||||
);
|
||||
|
||||
let def = grep_tool(fs);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({
|
||||
"pattern": "MATCH",
|
||||
"output_mode": "content",
|
||||
"-C": 1,
|
||||
});
|
||||
let out = tool.execute(&inp.to_string()).await.unwrap();
|
||||
let body = out.content.unwrap();
|
||||
// should contain: line2 (before context), MATCH, line4 (after context)
|
||||
assert!(body.contains("line2"));
|
||||
assert!(body.contains("MATCH"));
|
||||
assert!(body.contains("line4"));
|
||||
assert!(!body.contains("line1"));
|
||||
assert!(!body.contains("line5"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grep_multiline() {
|
||||
let (dir, fs) = setup();
|
||||
touch(
|
||||
&dir.path().join("a.txt"),
|
||||
"start\nfoo\nbar\nend\n",
|
||||
);
|
||||
|
||||
let def = grep_tool(fs);
|
||||
let (_, tool) = def();
|
||||
// Match across newlines: "foo" followed by "bar" on the next line
|
||||
let inp = serde_json::json!({
|
||||
"pattern": "foo[\\s\\S]*?bar",
|
||||
"multiline": true,
|
||||
"output_mode": "content",
|
||||
});
|
||||
let out = tool.execute(&inp.to_string()).await.unwrap();
|
||||
let body = out.content.unwrap();
|
||||
assert!(body.contains("foo"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grep_glob_filter() {
|
||||
let (dir, fs) = setup();
|
||||
touch(&dir.path().join("a.rs"), "target\n");
|
||||
touch(&dir.path().join("b.txt"), "target\n");
|
||||
|
||||
let def = grep_tool(fs);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({
|
||||
"pattern": "target",
|
||||
"glob": "*.rs",
|
||||
});
|
||||
let out = tool.execute(&inp.to_string()).await.unwrap();
|
||||
let body = out.content.unwrap();
|
||||
assert!(body.contains("a.rs"));
|
||||
assert!(!body.contains("b.txt"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grep_type_filter() {
|
||||
let (dir, fs) = setup();
|
||||
touch(&dir.path().join("a.rs"), "target\n");
|
||||
touch(&dir.path().join("b.py"), "target\n");
|
||||
|
||||
let def = grep_tool(fs);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({
|
||||
"pattern": "target",
|
||||
"type": "rust",
|
||||
});
|
||||
let out = tool.execute(&inp.to_string()).await.unwrap();
|
||||
let body = out.content.unwrap();
|
||||
assert!(body.contains("a.rs"));
|
||||
assert!(!body.contains("b.py"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grep_head_limit_truncates() {
|
||||
let (dir, fs) = setup();
|
||||
for i in 0..5 {
|
||||
touch(&dir.path().join(format!("f{i}.txt")), "x\n");
|
||||
}
|
||||
|
||||
let def = grep_tool(fs);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({
|
||||
"pattern": "x",
|
||||
"head_limit": 2,
|
||||
});
|
||||
let out = tool.execute(&inp.to_string()).await.unwrap();
|
||||
let body = out.content.unwrap();
|
||||
assert_eq!(body.lines().count(), 2);
|
||||
assert!(out.summary.contains("truncated at 2"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grep_offset_paginates() {
|
||||
let (dir, fs) = setup();
|
||||
// Create 5 files, all matching, deterministically named
|
||||
for i in 0..5 {
|
||||
touch(&dir.path().join(format!("f{i}.txt")), "x\n");
|
||||
}
|
||||
|
||||
let def = grep_tool(fs);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({
|
||||
"pattern": "x",
|
||||
"offset": 3,
|
||||
"head_limit": 10,
|
||||
});
|
||||
let out = tool.execute(&inp.to_string()).await.unwrap();
|
||||
let body = out.content.unwrap();
|
||||
// We skipped 3, so only 2 should remain.
|
||||
assert_eq!(body.lines().count(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grep_binary_files_are_skipped() {
|
||||
let (dir, fs) = setup();
|
||||
let mut bin = Vec::from(b"\x00\x01\x02needle\n".as_slice());
|
||||
bin.extend(b"more\n");
|
||||
fs::write(dir.path().join("a.bin"), bin).unwrap();
|
||||
touch(&dir.path().join("b.txt"), "needle\n");
|
||||
|
||||
let def = grep_tool(fs);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({ "pattern": "needle" });
|
||||
let out = tool.execute(&inp.to_string()).await.unwrap();
|
||||
let body = out.content.unwrap();
|
||||
assert!(body.contains("b.txt"));
|
||||
assert!(!body.contains("a.bin"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grep_invalid_regex() {
|
||||
let (_dir, fs) = setup();
|
||||
let def = grep_tool(fs);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({ "pattern": "(" });
|
||||
let err = tool.execute(&inp.to_string()).await.unwrap_err();
|
||||
assert!(matches!(err, ToolError::InvalidArgument(_)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grep_unknown_type() {
|
||||
let (_dir, fs) = setup();
|
||||
let def = grep_tool(fs);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({
|
||||
"pattern": "x",
|
||||
"type": "nonexistent",
|
||||
});
|
||||
let err = tool.execute(&inp.to_string()).await.unwrap_err();
|
||||
assert!(matches!(err, ToolError::InvalidArgument(_)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grep_no_matches() {
|
||||
let (dir, fs) = setup();
|
||||
touch(&dir.path().join("a.txt"), "nothing here\n");
|
||||
let def = grep_tool(fs);
|
||||
let (_, tool) = def();
|
||||
let inp = serde_json::json!({ "pattern": "zzz" });
|
||||
let out = tool.execute(&inp.to_string()).await.unwrap();
|
||||
assert_eq!(out.summary, "No files matched");
|
||||
assert!(out.content.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn grep_schema_contains_dash_keys() {
|
||||
// Sanity check: schemars must preserve the `-n`, `-A`, etc. keys
|
||||
// from serde(rename). If this fails we need to rename the fields.
|
||||
let schema = schemars::schema_for!(GrepParams);
|
||||
let json = serde_json::to_value(&schema).unwrap();
|
||||
let json_str = json.to_string();
|
||||
assert!(json_str.contains("\"-n\""), "schema missing -n: {json_str}");
|
||||
assert!(json_str.contains("\"-A\""), "schema missing -A: {json_str}");
|
||||
assert!(json_str.contains("\"-B\""), "schema missing -B: {json_str}");
|
||||
assert!(json_str.contains("\"-C\""), "schema missing -C: {json_str}");
|
||||
assert!(json_str.contains("\"-i\""), "schema missing -i: {json_str}");
|
||||
}
|
||||
}
|
||||
52
crates/tools/src/lib.rs
Normal file
52
crates/tools/src/lib.rs
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
//! Built-in tools for the Insomnia LLM agent.
|
||||
//!
|
||||
//! Implements Read / Write / Edit / Glob / Grep on top of the `llm-worker`
|
||||
//! `Tool` infrastructure. Filesystem access is mediated by two orthogonal
|
||||
//! concerns:
|
||||
//!
|
||||
//! - [`ScopedFs`] — pod-lifetime, expresses the write-block boundary for
|
||||
//! the current scope. Derived from the manifest and shareable across
|
||||
//! sessions.
|
||||
//! - [`ReadTracker`] — session-lifetime, enforces the "read before edit"
|
||||
//! policy via content hashes. Recreated fresh per session.
|
||||
//!
|
||||
//! The Pod layer owns both instances and passes them to
|
||||
//! [`builtin_tools`] when registering tools on a `Worker`.
|
||||
|
||||
pub mod error;
|
||||
pub mod read_tracker;
|
||||
pub mod scoped_fs;
|
||||
|
||||
mod edit;
|
||||
mod glob;
|
||||
mod grep;
|
||||
mod read;
|
||||
mod write;
|
||||
|
||||
pub use edit::edit_tool;
|
||||
pub use error::ToolsError;
|
||||
pub use glob::glob_tool;
|
||||
pub use grep::grep_tool;
|
||||
pub use read::read_tool;
|
||||
pub use read_tracker::ReadTracker;
|
||||
pub use scoped_fs::ScopedFs;
|
||||
pub use write::write_tool;
|
||||
|
||||
/// Register all builtin tools, wiring them to a shared `ScopedFs`
|
||||
/// (pod-lifetime) and `ReadTracker` (session-lifetime).
|
||||
///
|
||||
/// All returned factories share the same tracker instance so that
|
||||
/// `Read` / `Write` / `Edit` see a consistent history across tool
|
||||
/// invocations within a single session.
|
||||
pub fn builtin_tools(
|
||||
fs: ScopedFs,
|
||||
tracker: ReadTracker,
|
||||
) -> Vec<llm_worker::tool::ToolDefinition> {
|
||||
vec![
|
||||
read_tool(fs.clone(), tracker.clone()),
|
||||
write_tool(fs.clone(), tracker.clone()),
|
||||
edit_tool(fs.clone(), tracker.clone()),
|
||||
glob_tool(fs.clone()),
|
||||
grep_tool(fs),
|
||||
]
|
||||
}
|
||||
205
crates/tools/src/read.rs
Normal file
205
crates/tools/src/read.rs
Normal file
|
|
@ -0,0 +1,205 @@
|
|||
//! `Read` tool — read a text file with offset/limit, return line-numbered output.
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::read_tracker::ReadTracker;
|
||||
use crate::scoped_fs::ScopedFs;
|
||||
|
||||
const DESCRIPTION: &str = "Read a text file from the local filesystem. \
|
||||
Supports offset/limit for large files. Returns line-numbered output (1-based). \
|
||||
Directories cannot be read. The file must be read before Write or Edit can \
|
||||
modify it. Paths must be absolute.";
|
||||
|
||||
const DEFAULT_LIMIT: usize = 2000;
|
||||
|
||||
#[derive(Debug, Deserialize, schemars::JsonSchema)]
|
||||
pub(crate) struct ReadParams {
|
||||
/// Absolute path to the file.
|
||||
pub file_path: PathBuf,
|
||||
/// 0-based line offset from the start. Defaults to 0.
|
||||
#[serde(default)]
|
||||
pub offset: Option<usize>,
|
||||
/// Maximum number of lines to return. Defaults to 2000.
|
||||
#[serde(default)]
|
||||
pub limit: Option<usize>,
|
||||
}
|
||||
|
||||
pub(crate) struct ReadTool {
|
||||
fs: ScopedFs,
|
||||
tracker: ReadTracker,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ReadTool {
|
||||
async fn execute(&self, input_json: &str) -> Result<ToolOutput, ToolError> {
|
||||
let params: ReadParams = serde_json::from_str(input_json)
|
||||
.map_err(|e| ToolError::InvalidArgument(format!("invalid Read input: {e}")))?;
|
||||
let offset = params.offset.unwrap_or(0);
|
||||
let limit = params.limit.unwrap_or(DEFAULT_LIMIT).max(1);
|
||||
|
||||
tracing::debug!(
|
||||
path = %params.file_path.display(),
|
||||
offset,
|
||||
limit,
|
||||
"Read"
|
||||
);
|
||||
|
||||
let bytes = self.fs.read_bytes(¶ms.file_path)?;
|
||||
// Record the raw bytes under the read-history so subsequent Edit /
|
||||
// Write can detect external modification.
|
||||
self.tracker.record(¶ms.file_path, &bytes);
|
||||
|
||||
let text = String::from_utf8_lossy(&bytes).into_owned();
|
||||
let rendered = render_numbered(&text, offset, limit);
|
||||
|
||||
let summary = if rendered.truncated {
|
||||
format!(
|
||||
"Read {} line(s) [{}..{}] of {} from {}",
|
||||
rendered.line_count,
|
||||
offset + 1,
|
||||
offset + rendered.line_count,
|
||||
rendered.total_lines,
|
||||
params.file_path.display()
|
||||
)
|
||||
} else {
|
||||
format!(
|
||||
"Read {} line(s) from {}",
|
||||
rendered.line_count,
|
||||
params.file_path.display()
|
||||
)
|
||||
};
|
||||
|
||||
Ok(ToolOutput {
|
||||
summary,
|
||||
content: Some(rendered.body),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct Rendered {
|
||||
body: String,
|
||||
line_count: usize,
|
||||
total_lines: usize,
|
||||
truncated: bool,
|
||||
}
|
||||
|
||||
/// Format a slice of lines from `text` with `cat -n` style 1-based line
|
||||
/// numbers. Pure function — no I/O, no history touching.
|
||||
fn render_numbered(text: &str, offset: usize, limit: usize) -> Rendered {
|
||||
let all_lines: Vec<&str> = text.lines().collect();
|
||||
let total_lines = all_lines.len();
|
||||
let start = offset.min(total_lines);
|
||||
let end = start.saturating_add(limit).min(total_lines);
|
||||
let slice = &all_lines[start..end];
|
||||
let line_count = slice.len();
|
||||
|
||||
use std::fmt::Write as _;
|
||||
let mut body = String::with_capacity(text.len().saturating_add(line_count * 8));
|
||||
for (i, line) in slice.iter().enumerate() {
|
||||
let lineno = start + i + 1;
|
||||
let _ = writeln!(&mut body, "{:>6}\t{}", lineno, line);
|
||||
}
|
||||
|
||||
Rendered {
|
||||
body,
|
||||
line_count,
|
||||
total_lines,
|
||||
truncated: start > 0 || end < total_lines,
|
||||
}
|
||||
}
|
||||
|
||||
/// Factory for the `Read` tool.
|
||||
pub fn read_tool(fs: ScopedFs, tracker: ReadTracker) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(ReadParams);
|
||||
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
|
||||
let meta = ToolMeta::new("Read")
|
||||
.description(DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(ReadTool {
|
||||
fs: fs.clone(),
|
||||
tracker: tracker.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use manifest::Scope;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn setup() -> (TempDir, ScopedFs, ReadTracker) {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let fs = ScopedFs::new(Scope::new(dir.path()).unwrap());
|
||||
(dir, fs, ReadTracker::new())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_tool_basic_records_history() {
|
||||
let (dir, fs, tracker) = setup();
|
||||
let file = dir.path().join("a.txt");
|
||||
std::fs::write(&file, "alpha\nbeta\ngamma\n").unwrap();
|
||||
|
||||
let def = read_tool(fs, tracker.clone());
|
||||
let (meta, tool) = def();
|
||||
assert_eq!(meta.name, "Read");
|
||||
|
||||
let input = serde_json::json!({ "file_path": file.to_str().unwrap() });
|
||||
let out = tool.execute(&input.to_string()).await.unwrap();
|
||||
assert!(out.summary.contains("Read 3 line(s)"));
|
||||
let body = out.content.unwrap();
|
||||
assert!(body.contains(" 1\talpha"));
|
||||
assert!(body.contains(" 3\tgamma"));
|
||||
|
||||
// History recorded
|
||||
assert!(tracker.has(&file));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_tool_offset_limit() {
|
||||
let (dir, fs, tracker) = setup();
|
||||
let file = dir.path().join("a.txt");
|
||||
std::fs::write(&file, "1\n2\n3\n4\n5\n").unwrap();
|
||||
|
||||
let def = read_tool(fs, tracker);
|
||||
let (_, tool) = def();
|
||||
let input = serde_json::json!({
|
||||
"file_path": file.to_str().unwrap(),
|
||||
"offset": 1,
|
||||
"limit": 2,
|
||||
});
|
||||
let out = tool.execute(&input.to_string()).await.unwrap();
|
||||
assert!(out.summary.contains("[2..3] of 5"));
|
||||
let body = out.content.unwrap();
|
||||
assert!(body.contains(" 2\t2"));
|
||||
assert!(body.contains(" 3\t3"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_tool_missing_file() {
|
||||
let (dir, fs, tracker) = setup();
|
||||
let def = read_tool(fs, tracker);
|
||||
let (_, tool) = def();
|
||||
let input = serde_json::json!({
|
||||
"file_path": dir.path().join("nope.txt").to_str().unwrap()
|
||||
});
|
||||
let err = tool.execute(&input.to_string()).await.unwrap_err();
|
||||
assert!(matches!(err, ToolError::ExecutionFailed(_)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_tool_bad_json() {
|
||||
let (_dir, fs, tracker) = setup();
|
||||
let def = read_tool(fs, tracker);
|
||||
let (_, tool) = def();
|
||||
let err = tool.execute("not json").await.unwrap_err();
|
||||
assert!(matches!(err, ToolError::InvalidArgument(_)));
|
||||
}
|
||||
}
|
||||
213
crates/tools/src/read_tracker.rs
Normal file
213
crates/tools/src/read_tracker.rs
Normal file
|
|
@ -0,0 +1,213 @@
|
|||
//! Read-before-edit policy tracker for the builtin file-manipulation tools.
|
||||
//!
|
||||
//! A `ReadTracker` records a SHA-256 hash of each file's contents at the
|
||||
//! moment it was observed via the `Read` tool, and lets `Write` / `Edit`
|
||||
//! later verify that the file has not been externally modified since then.
|
||||
//!
|
||||
//! # Lifetime
|
||||
//!
|
||||
//! A `ReadTracker` is **session-scoped**: the Pod layer creates a fresh
|
||||
//! instance at the start of each agent session and discards it when the
|
||||
//! session ends. The `ScopedFs` write boundary, by contrast, is
|
||||
//! pod-lifetime (derived from the manifest). The two are orthogonal and
|
||||
//! the Pod wires them together when registering builtin tools.
|
||||
//!
|
||||
//! ```no_run
|
||||
//! # use manifest::Scope;
|
||||
//! # use tools::{ScopedFs, ReadTracker, builtin_tools};
|
||||
//! let scope = Scope::new("/workspace").unwrap();
|
||||
//! let fs = ScopedFs::new(scope); // pod lifetime
|
||||
//! let tracker = ReadTracker::new(); // session lifetime
|
||||
//! let defs = builtin_tools(fs, tracker);
|
||||
//! ```
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::error::ToolsError;
|
||||
|
||||
/// Fixed-size content hash recorded per file.
|
||||
type ContentHash = [u8; 32];
|
||||
|
||||
fn hash_bytes(bytes: &[u8]) -> ContentHash {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(bytes);
|
||||
hasher.finalize().into()
|
||||
}
|
||||
|
||||
/// Canonical-path keyed record of which files have been observed and at
|
||||
/// what content hash.
|
||||
///
|
||||
/// Cheap to clone: internally an `Arc<Mutex<HashMap<...>>>`, so sharing a
|
||||
/// `ReadTracker` across every builtin tool in a session is effectively
|
||||
/// free and keeps their views consistent.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ReadTracker {
|
||||
inner: Arc<Mutex<HashMap<PathBuf, ContentHash>>>,
|
||||
}
|
||||
|
||||
impl ReadTracker {
|
||||
/// Create an empty tracker. Typically called once per session.
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Record that `path` has been observed with the given content bytes.
|
||||
///
|
||||
/// Called by the `Read` tool after a successful read, and by the
|
||||
/// `Write` / `Edit` tools after a successful modification (so that
|
||||
/// subsequent edits see a clean history).
|
||||
pub fn record(&self, path: &Path, bytes: &[u8]) {
|
||||
let key = canonicalize_or_owned(path);
|
||||
let hash = hash_bytes(bytes);
|
||||
self.inner
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.insert(key, hash);
|
||||
}
|
||||
|
||||
/// Verify that `path` was previously recorded and its current bytes
|
||||
/// match the recorded hash.
|
||||
///
|
||||
/// - If the path has no history entry, returns [`ToolsError::NotRead`].
|
||||
/// - If the current content hashes differ from the recorded value,
|
||||
/// returns [`ToolsError::ExternallyModified`].
|
||||
pub fn verify(&self, path: &Path, current_bytes: &[u8]) -> Result<(), ToolsError> {
|
||||
let key = canonicalize_or_owned(path);
|
||||
let guard = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let recorded = guard
|
||||
.get(&key)
|
||||
.ok_or_else(|| ToolsError::NotRead(path.to_path_buf()))?;
|
||||
let current = hash_bytes(current_bytes);
|
||||
if *recorded != current {
|
||||
return Err(ToolsError::ExternallyModified(path.to_path_buf()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Returns true if `path` has a history entry. Test-only.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn has(&self, path: &Path) -> bool {
|
||||
let key = canonicalize_or_owned(path);
|
||||
self.inner
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.contains_key(&key)
|
||||
}
|
||||
|
||||
/// Number of distinct files in the history. Test-only.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn len(&self) -> usize {
|
||||
self.inner
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.len()
|
||||
}
|
||||
}
|
||||
|
||||
fn canonicalize_or_owned(path: &Path) -> PathBuf {
|
||||
path.canonicalize().unwrap_or_else(|_| path.to_path_buf())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::fs;
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[test]
|
||||
fn record_then_verify_clean_ok() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let file = dir.path().join("a.txt");
|
||||
fs::write(&file, b"hello").unwrap();
|
||||
|
||||
let tracker = ReadTracker::new();
|
||||
tracker.record(&file, b"hello");
|
||||
assert!(tracker.has(&file));
|
||||
assert_eq!(tracker.len(), 1);
|
||||
tracker.verify(&file, b"hello").unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn verify_without_record_returns_not_read() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let file = dir.path().join("a.txt");
|
||||
fs::write(&file, b"x").unwrap();
|
||||
|
||||
let tracker = ReadTracker::new();
|
||||
let err = tracker.verify(&file, b"x").unwrap_err();
|
||||
assert!(matches!(err, ToolsError::NotRead(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn verify_mismatch_returns_externally_modified() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let file = dir.path().join("a.txt");
|
||||
fs::write(&file, b"original").unwrap();
|
||||
|
||||
let tracker = ReadTracker::new();
|
||||
tracker.record(&file, b"original");
|
||||
let err = tracker.verify(&file, b"tampered").unwrap_err();
|
||||
assert!(matches!(err, ToolsError::ExternallyModified(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn record_overwrites_previous_hash() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let file = dir.path().join("a.txt");
|
||||
fs::write(&file, b"v1").unwrap();
|
||||
|
||||
let tracker = ReadTracker::new();
|
||||
tracker.record(&file, b"v1");
|
||||
tracker.record(&file, b"v2");
|
||||
tracker.verify(&file, b"v2").unwrap();
|
||||
assert!(tracker.verify(&file, b"v1").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn canonical_keys_collapse_symlink_variants() {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::symlink;
|
||||
let dir = TempDir::new().unwrap();
|
||||
let real = dir.path().join("real.txt");
|
||||
fs::write(&real, b"data").unwrap();
|
||||
let link = dir.path().join("link.txt");
|
||||
symlink(&real, &link).unwrap();
|
||||
|
||||
let tracker = ReadTracker::new();
|
||||
tracker.record(&real, b"data");
|
||||
// Looking up via the symlink should hit the same entry.
|
||||
tracker.verify(&link, b"data").unwrap();
|
||||
// Exactly one entry.
|
||||
assert_eq!(tracker.len(), 1);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clone_shares_state() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let file = dir.path().join("a.txt");
|
||||
fs::write(&file, b"x").unwrap();
|
||||
|
||||
let t1 = ReadTracker::new();
|
||||
let t2 = t1.clone();
|
||||
t1.record(&file, b"x");
|
||||
t2.verify(&file, b"x").unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_bytes_hash_stable() {
|
||||
let tracker = ReadTracker::new();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let file = dir.path().join("empty.txt");
|
||||
fs::write(&file, b"").unwrap();
|
||||
|
||||
tracker.record(&file, b"");
|
||||
tracker.verify(&file, b"").unwrap();
|
||||
assert!(tracker.verify(&file, b"x").is_err());
|
||||
}
|
||||
}
|
||||
258
crates/tools/src/scoped_fs.rs
Normal file
258
crates/tools/src/scoped_fs.rs
Normal file
|
|
@ -0,0 +1,258 @@
|
|||
//! Scope-aware filesystem primitive.
|
||||
//!
|
||||
//! `ScopedFs` represents **only** the write-block boundary: it knows a
|
||||
//! [`manifest::Scope`] and refuses writes outside of it. It carries no
|
||||
//! per-session state and is cheap to clone (pod-lifetime, reusable across
|
||||
//! sessions). The read-before-edit policy lives separately in
|
||||
//! [`crate::ReadTracker`].
|
||||
//!
|
||||
//! Reads are unrestricted by design (see `tickets/builtin-tools.md`).
|
||||
|
||||
use std::io::Write as _;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use manifest::Scope;
|
||||
|
||||
use crate::error::ToolsError;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct ScopedFsInner {
|
||||
scope: Scope,
|
||||
}
|
||||
|
||||
/// Scope-aware filesystem handle. Clone-cheap (`Arc` inside).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ScopedFs {
|
||||
inner: Arc<ScopedFsInner>,
|
||||
}
|
||||
|
||||
/// Outcome of a [`ScopedFs::write`] call.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct WriteOutcome {
|
||||
pub bytes_written: usize,
|
||||
pub created: bool,
|
||||
}
|
||||
|
||||
impl ScopedFs {
|
||||
/// Create a new [`ScopedFs`] wrapping the given [`Scope`].
|
||||
pub fn new(scope: Scope) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(ScopedFsInner { scope }),
|
||||
}
|
||||
}
|
||||
|
||||
/// The underlying [`Scope`].
|
||||
pub fn scope(&self) -> &Scope {
|
||||
&self.inner.scope
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Read — unrestricted
|
||||
// =========================================================================
|
||||
|
||||
/// Read the full contents of `path` as raw bytes.
|
||||
///
|
||||
/// Follows symlinks. Rejects directories, relative paths, and missing
|
||||
/// files. No scope check.
|
||||
pub fn read_bytes(&self, path: &Path) -> Result<Vec<u8>, ToolsError> {
|
||||
if !path.is_absolute() {
|
||||
return Err(ToolsError::RelativePath(path.to_path_buf()));
|
||||
}
|
||||
let meta = std::fs::metadata(path).map_err(|e| match e.kind() {
|
||||
std::io::ErrorKind::NotFound => ToolsError::NotFound(path.to_path_buf()),
|
||||
_ => ToolsError::io(path, e),
|
||||
})?;
|
||||
if meta.is_dir() {
|
||||
return Err(ToolsError::IsDirectory(path.to_path_buf()));
|
||||
}
|
||||
std::fs::read(path).map_err(|e| ToolsError::io(path, e))
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Write — scope-checked, atomic
|
||||
// =========================================================================
|
||||
|
||||
/// Atomically write `content` to `path`, creating or overwriting it.
|
||||
///
|
||||
/// - `path` must be absolute and inside the scope (delegates to
|
||||
/// [`Scope::contains`]).
|
||||
/// - Missing parent directories inside the scope are created.
|
||||
/// - The actual write uses a sibling tempfile + `persist`, so the
|
||||
/// target file transitions atomically between states.
|
||||
///
|
||||
/// This method does **not** consult any read history. Callers that
|
||||
/// want the "must read before overwrite" policy should verify with a
|
||||
/// [`ReadTracker`](crate::ReadTracker) beforehand.
|
||||
pub fn write(&self, path: &Path, content: &[u8]) -> Result<WriteOutcome, ToolsError> {
|
||||
if !path.is_absolute() {
|
||||
return Err(ToolsError::RelativePath(path.to_path_buf()));
|
||||
}
|
||||
if !self.inner.scope.contains(path) {
|
||||
return Err(ToolsError::OutOfScope(path.to_path_buf()));
|
||||
}
|
||||
|
||||
// Reject existing directory targets.
|
||||
match std::fs::metadata(path) {
|
||||
Ok(meta) if meta.is_dir() => {
|
||||
return Err(ToolsError::IsDirectory(path.to_path_buf()));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
let existed = path.exists();
|
||||
|
||||
let parent = path.parent().ok_or_else(|| {
|
||||
ToolsError::InvalidArgument(format!(
|
||||
"path has no parent directory: {}",
|
||||
path.display()
|
||||
))
|
||||
})?;
|
||||
if !parent.as_os_str().is_empty() && !parent.exists() {
|
||||
std::fs::create_dir_all(parent).map_err(|e| ToolsError::io(parent, e))?;
|
||||
}
|
||||
|
||||
let tmp_parent: &Path = if parent.as_os_str().is_empty() {
|
||||
Path::new(".")
|
||||
} else {
|
||||
parent
|
||||
};
|
||||
let mut tmp = tempfile::NamedTempFile::new_in(tmp_parent)
|
||||
.map_err(|e| ToolsError::io(tmp_parent, e))?;
|
||||
tmp.write_all(content).map_err(|e| ToolsError::io(path, e))?;
|
||||
tmp.as_file()
|
||||
.sync_all()
|
||||
.map_err(|e| ToolsError::io(path, e))?;
|
||||
tmp.persist(path)
|
||||
.map_err(|e| ToolsError::io(path, e.error))?;
|
||||
|
||||
Ok(WriteOutcome {
|
||||
bytes_written: content.len(),
|
||||
created: !existed,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Tests
|
||||
// =============================================================================
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::fs;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn make_fs(dir: &TempDir) -> ScopedFs {
|
||||
ScopedFs::new(Scope::new(dir.path()).unwrap())
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------
|
||||
// read_bytes
|
||||
// -------------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn read_bytes_returns_content() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let fs = make_fs(&dir);
|
||||
let file = dir.path().join("a.txt");
|
||||
fs::write(&file, b"abc").unwrap();
|
||||
assert_eq!(fs.read_bytes(&file).unwrap(), b"abc");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_bytes_rejects_relative() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let fs = make_fs(&dir);
|
||||
let err = fs.read_bytes(Path::new("rel.txt")).unwrap_err();
|
||||
assert!(matches!(err, ToolsError::RelativePath(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_bytes_rejects_directory() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let fs = make_fs(&dir);
|
||||
let err = fs.read_bytes(dir.path()).unwrap_err();
|
||||
assert!(matches!(err, ToolsError::IsDirectory(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_bytes_rejects_missing() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let fs = make_fs(&dir);
|
||||
let err = fs.read_bytes(&dir.path().join("nope.txt")).unwrap_err();
|
||||
assert!(matches!(err, ToolsError::NotFound(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_bytes_allows_paths_outside_scope() {
|
||||
// Reads are unrestricted — scope only gates writes.
|
||||
let dir = TempDir::new().unwrap();
|
||||
let outside = TempDir::new().unwrap();
|
||||
let outside_file = outside.path().join("x.txt");
|
||||
fs::write(&outside_file, b"hi").unwrap();
|
||||
|
||||
let scoped = make_fs(&dir);
|
||||
assert_eq!(scoped.read_bytes(&outside_file).unwrap(), b"hi");
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------
|
||||
// write
|
||||
// -------------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn write_creates_new_file() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let fs = make_fs(&dir);
|
||||
let file = dir.path().join("new.txt");
|
||||
let out = fs.write(&file, b"hello").unwrap();
|
||||
assert!(out.created);
|
||||
assert_eq!(out.bytes_written, 5);
|
||||
assert_eq!(fs::read(&file).unwrap(), b"hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_overwrites_existing() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let fs = make_fs(&dir);
|
||||
let file = dir.path().join("a.txt");
|
||||
fs::write(&file, b"old").unwrap();
|
||||
let out = fs.write(&file, b"new").unwrap();
|
||||
assert!(!out.created);
|
||||
assert_eq!(fs::read(&file).unwrap(), b"new");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_rejects_out_of_scope() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let outside = TempDir::new().unwrap();
|
||||
let fs = make_fs(&dir);
|
||||
let err = fs.write(&outside.path().join("x"), b"x").unwrap_err();
|
||||
assert!(matches!(err, ToolsError::OutOfScope(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_rejects_relative_path() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let fs = make_fs(&dir);
|
||||
let err = fs.write(Path::new("rel.txt"), b"x").unwrap_err();
|
||||
assert!(matches!(err, ToolsError::RelativePath(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_creates_missing_parents_inside_scope() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let fs = make_fs(&dir);
|
||||
let nested = dir.path().join("a/b/c/deep.txt");
|
||||
fs.write(&nested, b"x").unwrap();
|
||||
assert_eq!(fs::read(&nested).unwrap(), b"x");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_rejects_directory_target() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let fs = make_fs(&dir);
|
||||
let err = fs.write(dir.path(), b"x").unwrap_err();
|
||||
assert!(matches!(err, ToolsError::IsDirectory(_)));
|
||||
}
|
||||
}
|
||||
202
crates/tools/src/write.rs
Normal file
202
crates/tools/src/write.rs
Normal file
|
|
@ -0,0 +1,202 @@
|
|||
//! `Write` tool — create or overwrite a file.
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::read_tracker::ReadTracker;
|
||||
use crate::scoped_fs::ScopedFs;
|
||||
|
||||
const DESCRIPTION: &str = "Create a new file or overwrite an existing one with \
|
||||
the given content. Missing parent directories within scope are created \
|
||||
automatically. Existing files must have been read first (via the Read tool) \
|
||||
in this session. Paths must be absolute.";
|
||||
|
||||
#[derive(Debug, Deserialize, schemars::JsonSchema)]
|
||||
pub(crate) struct WriteParams {
|
||||
/// Absolute path to the file.
|
||||
pub file_path: PathBuf,
|
||||
/// Full content to write. Overwrites any existing content.
|
||||
pub content: String,
|
||||
}
|
||||
|
||||
pub(crate) struct WriteTool {
|
||||
fs: ScopedFs,
|
||||
tracker: ReadTracker,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for WriteTool {
|
||||
async fn execute(&self, input_json: &str) -> Result<ToolOutput, ToolError> {
|
||||
let params: WriteParams = serde_json::from_str(input_json)
|
||||
.map_err(|e| ToolError::InvalidArgument(format!("invalid Write input: {e}")))?;
|
||||
|
||||
tracing::debug!(
|
||||
path = %params.file_path.display(),
|
||||
bytes = params.content.len(),
|
||||
"Write"
|
||||
);
|
||||
|
||||
// Policy check: if the target already exists, it must have been
|
||||
// observed by the Read tool (via the tracker) and its current
|
||||
// contents must match the recorded hash.
|
||||
if params.file_path.exists() {
|
||||
let current = self.fs.read_bytes(¶ms.file_path)?;
|
||||
self.tracker.verify(¶ms.file_path, ¤t)?;
|
||||
}
|
||||
|
||||
let outcome = self.fs.write(¶ms.file_path, params.content.as_bytes())?;
|
||||
|
||||
// Refresh the history entry to reflect the newly-written content,
|
||||
// so a subsequent Edit / Write can proceed without a re-read.
|
||||
self.tracker
|
||||
.record(¶ms.file_path, params.content.as_bytes());
|
||||
|
||||
let summary = format!(
|
||||
"{} {} ({} bytes)",
|
||||
if outcome.created { "Created" } else { "Overwrote" },
|
||||
params.file_path.display(),
|
||||
outcome.bytes_written
|
||||
);
|
||||
Ok(ToolOutput {
|
||||
summary,
|
||||
content: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Factory for the `Write` tool.
|
||||
pub fn write_tool(fs: ScopedFs, tracker: ReadTracker) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(WriteParams);
|
||||
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
|
||||
let meta = ToolMeta::new("Write")
|
||||
.description(DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(WriteTool {
|
||||
fs: fs.clone(),
|
||||
tracker: tracker.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::read::read_tool;
|
||||
use manifest::Scope;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn setup() -> (TempDir, ScopedFs, ReadTracker) {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let fs = ScopedFs::new(Scope::new(dir.path()).unwrap());
|
||||
(dir, fs, ReadTracker::new())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_creates_new_file_without_read() {
|
||||
let (dir, fs, tracker) = setup();
|
||||
let def = write_tool(fs, tracker);
|
||||
let (meta, tool) = def();
|
||||
assert_eq!(meta.name, "Write");
|
||||
|
||||
let file = dir.path().join("new.txt");
|
||||
let input = serde_json::json!({
|
||||
"file_path": file.to_str().unwrap(),
|
||||
"content": "hello\n",
|
||||
});
|
||||
let out = tool.execute(&input.to_string()).await.unwrap();
|
||||
assert!(out.summary.contains("Created"));
|
||||
assert_eq!(std::fs::read_to_string(&file).unwrap(), "hello\n");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_existing_requires_prior_read() {
|
||||
let (dir, fs, tracker) = setup();
|
||||
let file = dir.path().join("a.txt");
|
||||
std::fs::write(&file, "old").unwrap();
|
||||
|
||||
let def = write_tool(fs, tracker);
|
||||
let (_, tool) = def();
|
||||
let input = serde_json::json!({
|
||||
"file_path": file.to_str().unwrap(),
|
||||
"content": "new",
|
||||
});
|
||||
let err = tool.execute(&input.to_string()).await.unwrap_err();
|
||||
assert!(matches!(err, ToolError::InvalidArgument(_)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_existing_after_read_succeeds() {
|
||||
let (dir, fs, tracker) = setup();
|
||||
let file = dir.path().join("a.txt");
|
||||
std::fs::write(&file, "old\n").unwrap();
|
||||
|
||||
let read_def = read_tool(fs.clone(), tracker.clone());
|
||||
let (_, reader) = read_def();
|
||||
let read_in = serde_json::json!({ "file_path": file.to_str().unwrap() });
|
||||
reader.execute(&read_in.to_string()).await.unwrap();
|
||||
|
||||
let write_def = write_tool(fs, tracker);
|
||||
let (_, writer) = write_def();
|
||||
let write_in = serde_json::json!({
|
||||
"file_path": file.to_str().unwrap(),
|
||||
"content": "new\n",
|
||||
});
|
||||
let out = writer.execute(&write_in.to_string()).await.unwrap();
|
||||
assert!(out.summary.contains("Overwrote"));
|
||||
assert_eq!(std::fs::read_to_string(&file).unwrap(), "new\n");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_detects_external_modification_via_hash() {
|
||||
let (dir, fs, tracker) = setup();
|
||||
let file = dir.path().join("a.txt");
|
||||
std::fs::write(&file, "v1").unwrap();
|
||||
|
||||
// Read records hash of "v1".
|
||||
let read_def = read_tool(fs.clone(), tracker.clone());
|
||||
let (_, reader) = read_def();
|
||||
reader
|
||||
.execute(&serde_json::json!({ "file_path": file.to_str().unwrap() }).to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// External process overwrites with a different content.
|
||||
std::fs::write(&file, "tampered").unwrap();
|
||||
|
||||
let write_def = write_tool(fs, tracker);
|
||||
let (_, writer) = write_def();
|
||||
let err = writer
|
||||
.execute(
|
||||
&serde_json::json!({
|
||||
"file_path": file.to_str().unwrap(),
|
||||
"content": "new",
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
let msg = format!("{err}");
|
||||
assert!(msg.contains("modified externally"), "{msg}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_rejects_out_of_scope() {
|
||||
let (_dir, fs, tracker) = setup();
|
||||
let outside = TempDir::new().unwrap();
|
||||
|
||||
let def = write_tool(fs, tracker);
|
||||
let (_, tool) = def();
|
||||
let input = serde_json::json!({
|
||||
"file_path": outside.path().join("x.txt").to_str().unwrap(),
|
||||
"content": "x",
|
||||
});
|
||||
let err = tool.execute(&input.to_string()).await.unwrap_err();
|
||||
assert!(matches!(err, ToolError::InvalidArgument(_)));
|
||||
}
|
||||
}
|
||||
242
crates/tools/tests/edge_cases.rs
Normal file
242
crates/tools/tests/edge_cases.rs
Normal file
|
|
@ -0,0 +1,242 @@
|
|||
//! Edge-case regression tests that should stay green.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use llm_worker::tool::{Tool, ToolDefinition};
|
||||
use manifest::Scope;
|
||||
use serde_json::json;
|
||||
use tempfile::TempDir;
|
||||
use tools::{ReadTracker, ScopedFs, builtin_tools};
|
||||
|
||||
struct Registry {
|
||||
entries: Vec<(llm_worker::tool::ToolMeta, Arc<dyn Tool>)>,
|
||||
}
|
||||
|
||||
impl Registry {
|
||||
fn new(defs: Vec<ToolDefinition>) -> Self {
|
||||
Self {
|
||||
entries: defs.into_iter().map(|f| f()).collect(),
|
||||
}
|
||||
}
|
||||
fn get(&self, name: &str) -> Arc<dyn Tool> {
|
||||
self.entries
|
||||
.iter()
|
||||
.find(|(m, _)| m.name == name)
|
||||
.map(|(_, t)| Arc::clone(t))
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
fn setup() -> (TempDir, Registry) {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let fs = ScopedFs::new(Scope::new(dir.path()).unwrap());
|
||||
let tracker = ReadTracker::new();
|
||||
(dir, Registry::new(builtin_tools(fs, tracker)))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unicode_path_and_content() {
|
||||
let (dir, reg) = setup();
|
||||
let file = dir.path().join("日本語ファイル.txt");
|
||||
let content = "こんにちは 🦀 世界\nabc\n";
|
||||
|
||||
let write = reg.get("Write");
|
||||
write
|
||||
.execute(
|
||||
&json!({
|
||||
"file_path": file.to_str().unwrap(),
|
||||
"content": content,
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let read = reg.get("Read");
|
||||
let out = read
|
||||
.execute(
|
||||
&json!({ "file_path": file.to_str().unwrap() })
|
||||
.to_string(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let body = out.content.unwrap();
|
||||
assert!(body.contains("🦀"));
|
||||
assert!(body.contains("こんにちは"));
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn symlink_to_outside_scope_is_rejected_for_write() {
|
||||
use std::os::unix::fs::symlink;
|
||||
|
||||
let (dir, reg) = setup();
|
||||
let outside = TempDir::new().unwrap();
|
||||
let outside_target = outside.path().join("secret.txt");
|
||||
std::fs::write(&outside_target, "secret").unwrap();
|
||||
|
||||
// Create a symlink inside the scope pointing to the outside file.
|
||||
let link = dir.path().join("linked.txt");
|
||||
symlink(&outside_target, &link).unwrap();
|
||||
|
||||
// Read tool must work against the symlink (read is unrestricted).
|
||||
let read = reg.get("Read");
|
||||
read.execute(
|
||||
&json!({ "file_path": link.to_str().unwrap() }).to_string(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Write through the symlink must be rejected because canonicalization
|
||||
// resolves it to outside the scope.
|
||||
let write = reg.get("Write");
|
||||
let err = write
|
||||
.execute(
|
||||
&json!({
|
||||
"file_path": link.to_str().unwrap(),
|
||||
"content": "overwritten",
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
let msg = format!("{err}");
|
||||
assert!(
|
||||
msg.contains("outside allowed scope"),
|
||||
"symlink escape not rejected: {msg}"
|
||||
);
|
||||
// Outside file must not have been touched.
|
||||
assert_eq!(std::fs::read_to_string(&outside_target).unwrap(), "secret");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_file_read_and_edit() {
|
||||
let (dir, reg) = setup();
|
||||
let file = dir.path().join("empty.txt");
|
||||
std::fs::write(&file, "").unwrap();
|
||||
|
||||
let read = reg.get("Read");
|
||||
let out = read
|
||||
.execute(&json!({ "file_path": file.to_str().unwrap() }).to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(out.summary.contains("0 line"));
|
||||
|
||||
// Edit on empty file must produce StringNotFound
|
||||
let edit = reg.get("Edit");
|
||||
let err = edit
|
||||
.execute(
|
||||
&json!({
|
||||
"file_path": file.to_str().unwrap(),
|
||||
"old_string": "foo",
|
||||
"new_string": "bar",
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
let msg = format!("{err}");
|
||||
assert!(msg.contains("not found"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn very_long_single_line() {
|
||||
let (dir, reg) = setup();
|
||||
let file = dir.path().join("long.txt");
|
||||
let big: String = "x".repeat(1024 * 1024); // 1 MiB, no newlines
|
||||
std::fs::write(&file, &big).unwrap();
|
||||
|
||||
let read = reg.get("Read");
|
||||
let out = read
|
||||
.execute(&json!({ "file_path": file.to_str().unwrap() }).to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
// Should return exactly 1 line
|
||||
assert!(out.summary.contains("1 line"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn relative_path_is_rejected() {
|
||||
let (_dir, reg) = setup();
|
||||
let read = reg.get("Read");
|
||||
let err = read
|
||||
.execute(&json!({ "file_path": "relative.txt" }).to_string())
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(format!("{err}").contains("absolute"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn directory_target_is_rejected_for_read() {
|
||||
let (dir, reg) = setup();
|
||||
let read = reg.get("Read");
|
||||
let err = read
|
||||
.execute(&json!({ "file_path": dir.path().to_str().unwrap() }).to_string())
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(format!("{err}").contains("directory"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn deeply_nested_new_file_is_created() {
|
||||
let (dir, reg) = setup();
|
||||
let deep = dir.path().join("a/b/c/d/e/deep.txt");
|
||||
let write = reg.get("Write");
|
||||
write
|
||||
.execute(
|
||||
&json!({
|
||||
"file_path": deep.to_str().unwrap(),
|
||||
"content": "deep\n",
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(std::fs::read_to_string(&deep).unwrap(), "deep\n");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn replace_preserves_unicode() {
|
||||
let (dir, reg) = setup();
|
||||
let file = dir.path().join("u.txt");
|
||||
std::fs::write(&file, "🦀 rust 🦀\n").unwrap();
|
||||
|
||||
let read = reg.get("Read");
|
||||
read.execute(&json!({ "file_path": file.to_str().unwrap() }).to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let edit = reg.get("Edit");
|
||||
edit.execute(
|
||||
&json!({
|
||||
"file_path": file.to_str().unwrap(),
|
||||
"old_string": "rust",
|
||||
"new_string": "ラスト",
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(std::fs::read_to_string(&file).unwrap(), "🦀 ラスト 🦀\n");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grep_handles_unicode_pattern() {
|
||||
let (dir, reg) = setup();
|
||||
let file = dir.path().join("u.txt");
|
||||
std::fs::write(&file, "English\n日本語\nрусский\n").unwrap();
|
||||
|
||||
let grep = reg.get("Grep");
|
||||
let out = grep
|
||||
.execute(
|
||||
&json!({
|
||||
"pattern": "日本語",
|
||||
"output_mode": "content",
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let body = out.content.unwrap();
|
||||
assert!(body.contains("日本語"));
|
||||
}
|
||||
274
crates/tools/tests/integration.rs
Normal file
274
crates/tools/tests/integration.rs
Normal file
|
|
@ -0,0 +1,274 @@
|
|||
//! Cross-tool integration tests exercising `builtin_tools()` end-to-end.
|
||||
//!
|
||||
//! `ToolServerHandle::register_tool` / `flush_pending` are `pub(crate)` in
|
||||
//! llm-worker, so from here we exercise the factories directly — the same
|
||||
//! code path that `flush_pending()` runs at production time.
|
||||
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use llm_worker::tool::{Tool, ToolDefinition, ToolMeta};
|
||||
use manifest::Scope;
|
||||
use serde_json::json;
|
||||
use tempfile::TempDir;
|
||||
use tools::{ReadTracker, ScopedFs, builtin_tools};
|
||||
|
||||
struct Registry {
|
||||
entries: Vec<(ToolMeta, Arc<dyn Tool>)>,
|
||||
}
|
||||
|
||||
impl Registry {
|
||||
fn new(defs: Vec<ToolDefinition>) -> Self {
|
||||
let entries = defs.into_iter().map(|f| f()).collect();
|
||||
Self { entries }
|
||||
}
|
||||
|
||||
fn get(&self, name: &str) -> Arc<dyn Tool> {
|
||||
self.entries
|
||||
.iter()
|
||||
.find(|(m, _)| m.name == name)
|
||||
.map(|(_, t)| Arc::clone(t))
|
||||
.unwrap_or_else(|| panic!("tool not found: {name}"))
|
||||
}
|
||||
|
||||
fn names(&self) -> Vec<&str> {
|
||||
self.entries.iter().map(|(m, _)| m.name.as_str()).collect()
|
||||
}
|
||||
}
|
||||
|
||||
fn setup() -> (TempDir, Registry) {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let fs = ScopedFs::new(Scope::new(dir.path()).unwrap());
|
||||
let tracker = ReadTracker::new();
|
||||
let reg = Registry::new(builtin_tools(fs, tracker));
|
||||
(dir, reg)
|
||||
}
|
||||
|
||||
async fn call(tool: &Arc<dyn Tool>, input: serde_json::Value) -> llm_worker::tool::ToolOutput {
|
||||
tool.execute(&input.to_string())
|
||||
.await
|
||||
.expect("tool execution failed")
|
||||
}
|
||||
|
||||
async fn call_err(
|
||||
tool: &Arc<dyn Tool>,
|
||||
input: serde_json::Value,
|
||||
) -> llm_worker::tool::ToolError {
|
||||
tool.execute(&input.to_string())
|
||||
.await
|
||||
.expect_err("expected error")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builtin_tools_registers_all_five() {
|
||||
let (_dir, reg) = setup();
|
||||
let mut names = reg.names();
|
||||
names.sort();
|
||||
assert_eq!(names, vec!["Edit", "Glob", "Grep", "Read", "Write"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn meta_has_description_and_schema() {
|
||||
let (_dir, reg) = setup();
|
||||
for (meta, _) in ®.entries {
|
||||
assert!(!meta.description.is_empty(), "{} missing description", meta.name);
|
||||
// Input schema must be a JSON object
|
||||
assert!(
|
||||
meta.input_schema.is_object(),
|
||||
"{} input_schema is not an object",
|
||||
meta.name
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_then_edit_then_read_roundtrip() {
|
||||
let (dir, reg) = setup();
|
||||
let file = dir.path().join("a.txt");
|
||||
std::fs::write(&file, "hello world\n").unwrap();
|
||||
let p = file.to_str().unwrap();
|
||||
|
||||
let read = reg.get("Read");
|
||||
let edit = reg.get("Edit");
|
||||
|
||||
// Read
|
||||
let r = call(&read, json!({ "file_path": p })).await;
|
||||
assert!(r.content.unwrap().contains("hello world"));
|
||||
|
||||
// Edit (unique replacement)
|
||||
let e = call(
|
||||
&edit,
|
||||
json!({
|
||||
"file_path": p,
|
||||
"old_string": "world",
|
||||
"new_string": "universe",
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
assert!(e.summary.contains("1 replacement"));
|
||||
assert_eq!(std::fs::read_to_string(&file).unwrap(), "hello universe\n");
|
||||
|
||||
// Re-read reflects the change
|
||||
let r2 = call(&read, json!({ "file_path": p })).await;
|
||||
assert!(r2.content.unwrap().contains("hello universe"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_then_grep_finds_content() {
|
||||
let (dir, reg) = setup();
|
||||
let write = reg.get("Write");
|
||||
let grep = reg.get("Grep");
|
||||
|
||||
let file = dir.path().join("notes.txt");
|
||||
call(
|
||||
&write,
|
||||
json!({
|
||||
"file_path": file.to_str().unwrap(),
|
||||
"content": "alpha\nNEEDLE\nomega\n",
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let g = call(
|
||||
&grep,
|
||||
json!({
|
||||
"pattern": "NEEDLE",
|
||||
"output_mode": "content",
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let body = g.content.unwrap();
|
||||
assert!(body.contains("notes.txt"));
|
||||
assert!(body.contains("NEEDLE"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn glob_finds_written_files() {
|
||||
let (dir, reg) = setup();
|
||||
let write = reg.get("Write");
|
||||
let glob = reg.get("Glob");
|
||||
|
||||
for name in ["one.md", "two.md", "three.txt"] {
|
||||
call(
|
||||
&write,
|
||||
json!({
|
||||
"file_path": dir.path().join(name).to_str().unwrap(),
|
||||
"content": "x",
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
let g = call(&glob, json!({ "pattern": "*.md" })).await;
|
||||
let body = g.content.unwrap();
|
||||
assert!(body.contains("one.md"));
|
||||
assert!(body.contains("two.md"));
|
||||
assert!(!body.contains("three.txt"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn out_of_scope_write_is_rejected() {
|
||||
let (_dir, reg) = setup();
|
||||
let outside = TempDir::new().unwrap();
|
||||
let write = reg.get("Write");
|
||||
|
||||
let err = call_err(
|
||||
&write,
|
||||
json!({
|
||||
"file_path": outside.path().join("x.txt").to_str().unwrap(),
|
||||
"content": "x",
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
// ToolsError::OutOfScope → ToolError::InvalidArgument
|
||||
let msg = format!("{err}");
|
||||
assert!(msg.contains("outside allowed scope"), "unexpected: {msg}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_to_existing_without_read_fails() {
|
||||
let (dir, reg) = setup();
|
||||
let file = dir.path().join("exists.txt");
|
||||
std::fs::write(&file, "preexisting").unwrap();
|
||||
|
||||
let write = reg.get("Write");
|
||||
let err = call_err(
|
||||
&write,
|
||||
json!({
|
||||
"file_path": file.to_str().unwrap(),
|
||||
"content": "new",
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let msg = format!("{err}");
|
||||
assert!(msg.contains("has not been read"), "unexpected: {msg}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shared_scoped_fs_across_tools() {
|
||||
// The key invariant: all builtin tools share the same ScopedFs instance,
|
||||
// so read-history set by Read is visible to Edit and Write.
|
||||
let (dir, reg) = setup();
|
||||
let file = dir.path().join("shared.txt");
|
||||
std::fs::write(&file, "one\n").unwrap();
|
||||
|
||||
let read = reg.get("Read");
|
||||
let write = reg.get("Write");
|
||||
|
||||
// Read via Read tool
|
||||
call(&read, json!({ "file_path": file.to_str().unwrap() })).await;
|
||||
// Write via Write tool — must succeed because the shared ScopedFs has the read
|
||||
call(
|
||||
&write,
|
||||
json!({
|
||||
"file_path": file.to_str().unwrap(),
|
||||
"content": "two\n",
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(std::fs::read_to_string(&file).unwrap(), "two\n");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn edit_requires_read_across_tools() {
|
||||
let (dir, reg) = setup();
|
||||
let file = dir.path().join("a.txt");
|
||||
std::fs::write(&file, "foo\n").unwrap();
|
||||
|
||||
let edit = reg.get("Edit");
|
||||
// No prior Read — Edit should fail
|
||||
let err = call_err(
|
||||
&edit,
|
||||
json!({
|
||||
"file_path": file.to_str().unwrap(),
|
||||
"old_string": "foo",
|
||||
"new_string": "bar",
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let msg = format!("{err}");
|
||||
assert!(msg.contains("has not been read"), "unexpected: {msg}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn deterministic_tool_order_is_registration_order() {
|
||||
let (_dir, reg) = setup();
|
||||
// Registration order from builtin_tools(): Read, Write, Edit, Glob, Grep
|
||||
let names: Vec<&str> = reg.entries.iter().map(|(m, _)| m.name.as_str()).collect();
|
||||
assert_eq!(names, vec!["Read", "Write", "Edit", "Glob", "Grep"]);
|
||||
}
|
||||
|
||||
// Regression: tool name capitalization matches Claude Code reference
|
||||
#[test]
|
||||
fn tool_names_match_reference_spec() {
|
||||
let (_dir, reg) = setup();
|
||||
for expected in ["Read", "Write", "Edit", "Glob", "Grep"] {
|
||||
assert!(
|
||||
reg.entries.iter().any(|(m, _)| m.name == expected),
|
||||
"missing tool {expected}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Sanity: unused Path import guard
|
||||
const _: fn() -> &'static Path = || Path::new("/");
|
||||
|
|
@ -2,141 +2,313 @@
|
|||
|
||||
## 背景
|
||||
|
||||
`Pod::compact()` とその周辺機構(CompactInterceptor, CompactState, Controller 統合)は
|
||||
実装済み。挙動の詳細に未決定事項が残っており、要約品質にも改善余地がある。
|
||||
`Pod::compact()` とその周辺機構は実装済み。
|
||||
要約品質、保護単位、compact 後のコンテキスト構築に改善が必要。
|
||||
|
||||
## 前提チケット
|
||||
|
||||
- [token-counter.md](token-counter.md) — LlmClient に Tokenizer 導入。retained_tokens / auto-read budget がこれに依存
|
||||
- [tool-output-referenced-files.md](tool-output-referenced-files.md) — ToolOutput にファイル追跡フィールド追加。デフォルトリファレンスがこれに依存
|
||||
|
||||
---
|
||||
|
||||
## 1. 要約入力の改善
|
||||
## 要件
|
||||
|
||||
### 現状の問題
|
||||
### R1: 一貫した振る舞い
|
||||
- システムプロンプトは不変
|
||||
- compact 前後でユーザーが違和感を覚えない
|
||||
- 「何を知っていて何を忘れたか」が自然であること
|
||||
|
||||
`build_summary_prompt()` が全 Item をフラットにテキスト化して LLM に投げている。
|
||||
### R2: 直近の記憶の確実性
|
||||
- 直近 N トークン分の会話をそのまま保持(Prune 済み = summary only の状態で計測)
|
||||
- **トークン数ベース** で保護量を決める(ターン単位ではない)
|
||||
- 自走エージェントは1ターン内で多数のリクエストを回す
|
||||
- ターン単位だと保護量がターン長に依存してしまう
|
||||
|
||||
1. **データが多すぎる**: ToolResult の content(ファイル内容、grep 結果等)を含めている
|
||||
2. **単一関心事の前提**: "Original Task" が1つだけ。タスク切り替わりに対応できない
|
||||
### R3: Auto-Read + リファレンス
|
||||
- compact 後の最初のターンで、タスク遂行に必要なファイルが既に読まれている
|
||||
- 2段階: **Read**(全文/範囲をコンテキストに注入)と **Reference**(「読んだことがある」とだけ伝える)
|
||||
- compact worker が「続行に必要なファイル」を判断して指定する
|
||||
|
||||
### Phase 1: 入力データの削減
|
||||
### R4: マルチタスク対応
|
||||
- セッション中に一貫した課題に取り組んでいないものとする
|
||||
- **完了タスク**: 簡潔に。注意点・発覚した事実だけ
|
||||
- **進行中タスク**: サマリ + 現状 + 次のステップを十分に
|
||||
|
||||
`build_summary_prompt` で渡すデータを絞る:
|
||||
---
|
||||
|
||||
```rust
|
||||
fn build_summary_prompt(items: &[Item]) -> String {
|
||||
for item in items {
|
||||
match item {
|
||||
Item::ToolResult { summary, .. } => {
|
||||
// content を含めない。summary だけ
|
||||
lines.push(format!("[ToolResult] {summary}"));
|
||||
}
|
||||
Item::ToolCall { name, .. } => {
|
||||
// arguments を含めない。ツール名だけ
|
||||
lines.push(format!("[ToolCall] {name}"));
|
||||
}
|
||||
Item::Reasoning { .. } => {
|
||||
// skip(内部思考は要約に不要)
|
||||
}
|
||||
// User/Assistant のテキストはそのまま
|
||||
}
|
||||
}
|
||||
## 閾値の修正(重要)
|
||||
|
||||
現状の実装は閾値の大小関係が意図と逆になっている。修正する。
|
||||
|
||||
### 正しい方針
|
||||
|
||||
- **post-run (タスク区切り) = 早めの閾値**: タスクの区切りで先を見越して compact
|
||||
- **mid-turn (pre_llm_request) = 遅めの閾値**: ターン中は最終防衛ラインとして、遅くなっても止まらないよう
|
||||
|
||||
```
|
||||
manifest.compact_threshold → post_run_threshold (基本ライン, 早め)
|
||||
turn_threshold = post_run_threshold * 9 / 8 (safety net, 遅め)
|
||||
```
|
||||
|
||||
### 影響箇所
|
||||
|
||||
- `crates/pod/src/compact_state.rs`
|
||||
- フィールド名と初期化を入れ替え: `manifest compact_threshold` は `post_run_threshold` に代入
|
||||
- `turn_threshold` は `post_run_threshold * 9 / 8` として導出
|
||||
- テストの `assert_eq!(state.post_run_threshold, 90_000)` を逆転(`turn_threshold = 90_000`, `post_run_threshold = 80_000` が正)
|
||||
- `crates/pod/src/compact_interceptor.rs` — そのまま(`exceeds_turn` を呼ぶだけ)
|
||||
- `crates/pod/src/pod.rs:371` の `exceeds_post_run` 判定 — そのまま
|
||||
- `docs/compaction.md` — 「ターン間は早めの閾値」の記述を逆に修正
|
||||
|
||||
---
|
||||
|
||||
## compact 後の history 構造
|
||||
|
||||
全て system message(`Item::Message { role: System }`)として注入。
|
||||
|
||||
```
|
||||
[system prompt] ← 不変 (R1)
|
||||
[system: 構造化要約] ← R4: compact worker の出力
|
||||
[system: auto-read ファイル群] ← R3: read_required の結果
|
||||
[system: リファレンス一覧] ← R3: reference の結果
|
||||
[直近 N トークン分の生の会話] ← R2: pruned 状態で保持
|
||||
```
|
||||
|
||||
system message で統一する理由:
|
||||
- LLM に「システムから提供された前提情報」として認識させる
|
||||
- fake ユーザーメッセージや fake ToolCall を作らない
|
||||
- 要約もファイルも同じ role で自然に並ぶ
|
||||
|
||||
### auto-read の system message 例
|
||||
|
||||
```
|
||||
[Auto-read file: src/main.rs:42-142]
|
||||
fn main() {
|
||||
let config = Config::load();
|
||||
...
|
||||
}
|
||||
```
|
||||
|
||||
### Phase 2: 要約フォーマットの改善
|
||||
|
||||
タスク切り替わりを反映する:
|
||||
### リファレンスの system message 例
|
||||
|
||||
```
|
||||
## Tasks
|
||||
### Task 1: (最初のユーザー指示)
|
||||
- 完了した作業
|
||||
- 判明した事実
|
||||
|
||||
### Task 2: (次のユーザー指示)
|
||||
- 完了した作業
|
||||
- 判明した事実
|
||||
|
||||
## Current State
|
||||
- (変更されたファイル、残タスク)
|
||||
[Referenced files — read before compaction, contents not included]
|
||||
- src/config.rs (read during task setup)
|
||||
- tests/integration_test.rs (read during test implementation)
|
||||
Use read_file to access current contents if needed.
|
||||
```
|
||||
|
||||
### Phase 3: マルチターン要約 Worker
|
||||
|
||||
1リクエストで全部読ませるのではなく、要約 Worker にツールを持たせて自律的に要約させる。
|
||||
|
||||
```
|
||||
要約 Worker:
|
||||
system: 「セッションログを読んで構造化要約を生成せよ」
|
||||
ツール: read_session_segment(offset, limit)
|
||||
```
|
||||
|
||||
利点:
|
||||
- 巨大セッションでもコンテキストに収まる
|
||||
- Worker が自分で「重要/不要」を判断できる
|
||||
- タスク切り替わりを検出し、関心事ごとに要約できる
|
||||
|
||||
builtin-tools チケットとの依存あり。
|
||||
|
||||
---
|
||||
|
||||
## 2. 挙動の未決定事項
|
||||
## R2: トークンベースの保護
|
||||
|
||||
### 現在の挙動
|
||||
|
||||
**トリガー(2段階):**
|
||||
1. ターン間 (CompactInterceptor): `input_tokens > turn_threshold` → Yield → compact + resume
|
||||
2. run 後 (Controller): `input_tokens > post_run_threshold (×9/8)` → best-effort compact
|
||||
|
||||
**安全機構:**
|
||||
- サーキットブレーカー: 3回連続失敗で無効化
|
||||
- Thrash 検出: compact 直後に再び閾値超過 → CompactThrash エラー
|
||||
- Yield 前の永続化: persist_turn を compact の前に実行
|
||||
|
||||
### 2-1. Yield のタイミング精度
|
||||
|
||||
現状: `pre_llm_request` でチェック = ターンの切れ目でしか発火しない。
|
||||
1ターン内でツール呼び出しが多く、途中でコンテキストが膨らむケースは次のターンまで待つ。
|
||||
|
||||
検討:
|
||||
- ツール実行後にもチェックする?(`post_tool_call` で Yield 相当の処理)
|
||||
- 現状の「ターン切れ目のみ」で十分か?
|
||||
|
||||
### 2-2. 閾値の導出
|
||||
|
||||
- `turn_threshold` = マニフェストの `compact_threshold` そのまま
|
||||
- `post_run_threshold` = `turn_threshold * 9 / 8`(≈ 112.5%)
|
||||
|
||||
9/8 の根拠はない(安全マージン)。マニフェストで個別指定可能にする?
|
||||
|
||||
### 2-3. Prune と Compact の相互作用
|
||||
現状の `retained_turns` を `retained_tokens` に変更。
|
||||
|
||||
```
|
||||
pre_llm_request:
|
||||
1. PruneHook(content を除去)
|
||||
2. CompactInterceptor(トークン数チェック)
|
||||
history (全て pruned 済み = summary only):
|
||||
[...古い部分...] [...直近 N トークン分...]
|
||||
↓ ↓
|
||||
要約対象 そのまま新 history に載せる
|
||||
```
|
||||
|
||||
Prune はリクエストコンテキストのみ操作し、`last_input_tokens` は前回の LLM レスポンスの値。
|
||||
Prune の効果は `last_input_tokens` に反映されず、Compact の判断には影響しない。
|
||||
→ Prune で十分に縮んでも Compact が走る可能性がある。保守的で実害は小さい。
|
||||
- Prune 済みの history に対して `LlmClient::tokenizer()` でトークン数を推定
|
||||
- 末尾から逆順に数えて N トークン分の位置で切る
|
||||
- ターン境界は無視。アイテム単位で切る
|
||||
|
||||
### 2-4. compact 中のクライアント通知
|
||||
```toml
|
||||
[compaction]
|
||||
compact_threshold = 80000
|
||||
retained_tokens = 8000 # ← retained_turns から変更
|
||||
```
|
||||
|
||||
compact は LLM 呼び出しを伴う。この間 Controller は Pod を占有。
|
||||
`AlreadyRunning` エラーで弾かれる。Protocol チケットの `CompactStart`/`CompactDone` で対応。
|
||||
token-counter チケットが前提。
|
||||
|
||||
### 2-5. 復元時の挙動
|
||||
---
|
||||
|
||||
`Outcome::Yielded` で記録されたセッションを restore した場合:
|
||||
- `last_run_interrupted = true` で復元
|
||||
- Pod は resume 可能(通常の interrupted セッションと同じ)
|
||||
- compact 後の新セッションが存在する場合、どちらを restore するかは呼び出し側の責任
|
||||
- `compacted_from` で辿れる
|
||||
## R3: Auto-Read + リファレンス
|
||||
|
||||
### デフォルトリファレンスの抽出
|
||||
|
||||
Pod は `ToolOutput.referenced_files` を `HookInterceptor::post_tool_call` で観察し、
|
||||
LRU 的な履歴バッファに積む(→ tool-output-referenced-files チケット)。
|
||||
Compact 時は先頭 5 件を compact worker のデフォルトリファレンスとして渡す。
|
||||
|
||||
### compact worker のツール
|
||||
|
||||
```
|
||||
read_file(path, offset?, limit?) — ファイルを読んで判断するため
|
||||
mark_read_required(path, offset?, limit?) — auto-read 対象(内容をコンテキストに載せる)
|
||||
add_reference(path) — リファレンス追加(内容は載せない)
|
||||
write_summary(text) — 構造化要約を出力/上書き(上書き可)
|
||||
```
|
||||
|
||||
`write_summary` は**上書き可**。マルチターンで「下書き → 追加 read → 書き直し」の順序が自然に動く。
|
||||
最終的に直近の呼び出しが採用される。ガードは「一度も呼ばれていない」時のみ。
|
||||
|
||||
### フロー
|
||||
|
||||
1. Pod が referenced_files バッファから先頭 5 件を抽出(デフォルトリファレンス)
|
||||
2. compact worker のプロンプトに含める:
|
||||
|
||||
```
|
||||
以下のファイルがリファレンスとして指定されています。
|
||||
全て読んで、タスク続行に必要なものを mark_read_required で指定してください。
|
||||
リファレンスを追加したい場合は add_reference で追加できます。
|
||||
```
|
||||
|
||||
3. compact worker が read_file で全ファイルを読み、判断:
|
||||
- 必要なファイル → `mark_read_required(path, offset?, limit?)`
|
||||
- 不要だがコンテキストとして有用 → リファレンスのまま残す
|
||||
- 追加のリファレンス → `add_reference(path)`
|
||||
4. `write_summary` で構造化要約を出力(最後のが採用される)
|
||||
5. ターン終了時に summary が一度も書かれていない or read_required が空(かつファイル操作履歴がある場合)→ 追加プロンプトで促す
|
||||
|
||||
### Auto-Read の Budget 管理
|
||||
|
||||
compact worker が `mark_read_required` を無制限に呼ぶとコンテキストが膨張する。
|
||||
共有 budget で制御:
|
||||
|
||||
```toml
|
||||
[compaction]
|
||||
auto_read_budget = 8000 # 合計トークン上限
|
||||
```
|
||||
|
||||
- `mark_read_required` のツール結果で残量を返す:
|
||||
`"Marked. Budget: 4200/8000 tokens remaining"`
|
||||
- 50% 以下になったら次のツール結果に system reminder を append:
|
||||
`"Budget half consumed. Consider calling write_summary soon."`
|
||||
- 100% 超過で Err:
|
||||
`"Error: auto-read budget exhausted (8000 tokens). Remove an existing mark or use add_reference instead."`
|
||||
- compact worker が判断して自分で調整できる(Err は即中断ではない)
|
||||
|
||||
token-counter チケットが前提(budget の計測に `estimate_text` が要る)。
|
||||
|
||||
### compact worker の暴走抑止
|
||||
|
||||
Turn/request 数ではなく、compact worker の累計入力トークンで上限を設ける:
|
||||
|
||||
```toml
|
||||
[compaction]
|
||||
compact_worker_max_input_tokens = 50000
|
||||
```
|
||||
|
||||
超えたら compact worker を強制終了。`CompactState::record_compact_failure()` 経由で
|
||||
サーキットブレーカーの自然な経路に乗る。
|
||||
|
||||
---
|
||||
|
||||
## R4: 要約の内容と品質
|
||||
|
||||
### 出力方法
|
||||
|
||||
compact worker が `write_summary(text)` ツールで出力する(上書き可)。
|
||||
最後のテキスト出力ではなくツールにする理由:
|
||||
- マルチターンで read_file → 判断 → 要約の順序が自由
|
||||
- 要約を書いた後にさらにリファレンスを追加できる
|
||||
- 「要約を書いていない」のガードが mark_read_required と同じパターンで検出可能
|
||||
|
||||
### 含めるべき内容
|
||||
|
||||
コードスニペットは auto-read に任せる。要約に求めるのは:
|
||||
1. **何を、なぜやったか** — 意思決定の記録。具体的な型名・関数名で言及
|
||||
2. **ユーザーの指示・フィードバックの原文** — ニュアンス保持。重要なもののみ
|
||||
3. **発生した問題と解決策** — 同じ轍を踏まない
|
||||
4. **今どこにいて次に何をするか** — compact 前後の一貫性 (R1)
|
||||
|
||||
含めないもの:
|
||||
- コードの全文(auto-read が担う)
|
||||
- 変更の diff(git がある)
|
||||
- 中間のやりとりの詳細(最終結論だけ)
|
||||
|
||||
### フォーマット(5セクション、1000-2000 トークン目安)
|
||||
|
||||
```
|
||||
## Completed Tasks
|
||||
### (タスク名)
|
||||
- 完了した作業(具体的な型名・ファイル名で)
|
||||
- 注意点 / 発覚した事実
|
||||
|
||||
### (タスク名)
|
||||
- ...
|
||||
|
||||
## Active Task
|
||||
### (タスク名)
|
||||
- 目標
|
||||
- 現状(何が済んで何が未着手か)
|
||||
- 次のステップ
|
||||
|
||||
## Key Decisions
|
||||
- (判断内容) — (理由)
|
||||
- ...
|
||||
|
||||
## User Directives
|
||||
- 「(ユーザー発言の原文)」 — 重要な指示・フィードバックのみ
|
||||
- ...
|
||||
|
||||
## Current Work
|
||||
(直前に何をしていたか。2-3行)
|
||||
```
|
||||
|
||||
各セクションの目安量:
|
||||
- Completed Tasks: 各タスク 2-3 行 × タスク数
|
||||
- Active Task: 5-10 行
|
||||
- Key Decisions: 各 1-2 行
|
||||
- User Directives: 重要な発言のみ原文引用
|
||||
- Current Work: 2-3 行
|
||||
|
||||
### 要約の入力
|
||||
|
||||
pruned history から:
|
||||
- ToolResult は summary のみ(content 除去)
|
||||
- ToolCall は名前のみ(arguments 除去)
|
||||
- Reasoning は除去
|
||||
|
||||
---
|
||||
|
||||
## 挙動の未決定事項
|
||||
|
||||
### Yield のタイミング精度
|
||||
|
||||
現状 `pre_llm_request`(リクエストの切れ目)でのみチェック。
|
||||
1ターン内でツール呼び出しが多く途中でコンテキストが膨らむケースは次のリクエストまで待つ。
|
||||
|
||||
検討: `post_tool_call` でもチェックする?
|
||||
|
||||
### 閾値の比率
|
||||
|
||||
- `post_run_threshold` = マニフェストの `compact_threshold`
|
||||
- `turn_threshold` = `post_run_threshold * 9 / 8`(≈ 112.5%)
|
||||
|
||||
9/8 の根拠はない(安全マージン)。要調整。
|
||||
|
||||
### Prune と Compact の相互作用
|
||||
|
||||
Prune はリクエストコンテキストのみ操作、`last_input_tokens` は前回の LLM レスポンスの値。
|
||||
Prune の効果は閾値判断に反映されない。保守的(compact しすぎる方向)で実害は小さい。
|
||||
|
||||
### compact 中のクライアント通知
|
||||
|
||||
Protocol チケットの `CompactStart`/`CompactDone` で対応。
|
||||
|
||||
### 復元時の挙動
|
||||
|
||||
`Outcome::Yielded` で記録されたセッションは `last_run_interrupted = true` で復元。
|
||||
compact 後の新セッションが存在する場合、どちらを restore するかは呼び出し側の責任。
|
||||
`compacted_from` で辿れる。
|
||||
|
||||
---
|
||||
|
||||
## 実装順序
|
||||
|
||||
1. Phase 1(content/arguments/reasoning 除去)→ `build_summary_prompt` の変更のみ
|
||||
2. 挙動の未決定事項 → 実運用でのフィードバックを元に判断
|
||||
3. Phase 2(フォーマット改善)→ チューニング
|
||||
4. Phase 3(マルチターン要約)→ builtin-tools 後
|
||||
0. **[前提] token-counter** — LlmClient に Tokenizer
|
||||
0. **[前提] tool-output-referenced-files** — ToolOutput + Pod の LRU バッファ
|
||||
1. **閾値逆転の修正** — `compact_state.rs` のフィールド入れ替え、テスト修正、docs 更新
|
||||
2. **要約入力の削減** — `build_summary_prompt` から content/arguments/reasoning を除去
|
||||
3. **retained_tokens 化** — retained_turns → retained_tokens に変更。マニフェスト設定追加
|
||||
4. **compact worker のツール化** — read_file + mark_read_required + add_reference + write_summary (上書き可)
|
||||
5. **Auto-Read + リファレンス** — デフォルト5ファイル抽出 (referenced_files バッファから)、compact worker による選定、system message での注入
|
||||
6. **Auto-Read Budget** — `mark_read_required` のトークン会計、残量通知、超過エラー
|
||||
7. **compact worker の累計入力トークン制限** — `compact_worker_max_input_tokens`
|
||||
8. **要約フォーマット** — タスク分類の要約プロンプト調整
|
||||
9. **ガード** — write_summary 未呼び出し or mark_read_required 空時の追加プロンプト
|
||||
|
|
|
|||
67
tickets/token-counter.md
Normal file
67
tickets/token-counter.md
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
# LlmClient へ Tokenizer の導入
|
||||
|
||||
## 背景
|
||||
|
||||
現状、トークン数の推定は `len / 4` の荒い近似でしかできていない。
|
||||
Compact 改善で以下が全て正確なトークン数に依存するため、LlmClient 層で
|
||||
トークナイザを提供する仕組みが必要になる。
|
||||
|
||||
## 動機
|
||||
|
||||
正確なトークン数が必要になる箇所:
|
||||
|
||||
- Prune の `min_savings` 判定(節約見込みの事前推定)
|
||||
- Compact 後の `retained_tokens` 切り出し(直近 N トークンの保護)
|
||||
- Compact worker の auto-read budget 判定
|
||||
- Protocol の UI 向けトークン表示(将来)
|
||||
|
||||
`CompactState::last_input_tokens` は LLM レスポンスから得られる実測値なので
|
||||
これには影響しない(閾値比較だけなら Tokenizer なしで回る)。ただし
|
||||
retained_tokens の切り出しは事前にローカルで計算する必要がある。
|
||||
|
||||
## 方針
|
||||
|
||||
`LlmClient` trait に同期的な `Tokenizer` 取得口を追加する。非同期 API は使わない
|
||||
(Prune/Compact の判定フローを async にしたくない)。
|
||||
|
||||
```rust
|
||||
pub trait Tokenizer: Send + Sync {
|
||||
fn estimate_text(&self, s: &str) -> u64;
|
||||
fn estimate_items(&self, items: &[Item]) -> u64;
|
||||
}
|
||||
|
||||
pub trait LlmClient {
|
||||
// 既存メソッド...
|
||||
fn tokenizer(&self) -> Arc<dyn Tokenizer>;
|
||||
}
|
||||
```
|
||||
|
||||
- provider ごとに実装。精度は近似で十分(±10% 程度)
|
||||
- OpenAI/Anthropic 系は `tiktoken-rs` + BPE テーブルベースの近似
|
||||
- Gemini も近似で対応。厳密値が要る場面があれば後から `count_tokens` API 呼び出しに置き換え可
|
||||
- `estimate_items` は `Item` の variant を舐めて text/tool name/arguments を足し込む単純実装
|
||||
|
||||
## 設計ポイント
|
||||
|
||||
- **同期 API**: Prune hook や Compact 判定の中で使う。`async fn` を増やさない
|
||||
- **Arc 返却**: Worker/Pod で使い回せるよう共有参照
|
||||
- **近似で十分**: 正確性より呼び出しコストの低さを優先。実測値 (`last_input_tokens`) との二重化で補正される
|
||||
- **Client ごとの実装**: provider 固有のトークナイザ差異はここで吸収
|
||||
|
||||
## 実装対象
|
||||
|
||||
- `llm-worker/src/llm_client/tokenizer.rs` (新規) — `Tokenizer` trait 定義
|
||||
- `llm-worker/src/llm_client/types.rs` or `LlmClient` trait — `fn tokenizer(&self) -> Arc<dyn Tokenizer>`
|
||||
- provider 実装:
|
||||
- `providers/anthropic.rs`
|
||||
- `providers/openai.rs`
|
||||
- `providers/gemini.rs`
|
||||
- 依存追加: `tiktoken-rs` (`cargo add` で)
|
||||
|
||||
## 依存
|
||||
|
||||
- なし(これ自体が前提チケット)
|
||||
|
||||
## ブロックする後続
|
||||
|
||||
- [compact-improvements.md](compact-improvements.md) — retained_tokens, auto-read budget が依存
|
||||
76
tickets/tool-output-referenced-files.md
Normal file
76
tickets/tool-output-referenced-files.md
Normal file
|
|
@ -0,0 +1,76 @@
|
|||
# ToolOutput.referenced_files: ツールが参照したファイルの追跡
|
||||
|
||||
## 背景
|
||||
|
||||
Compact 実行時、「タスク続行に必要なファイル」を compact worker に提示するため
|
||||
のデフォルトリファレンス(過去に読み書きされたファイル一覧)が必要。
|
||||
builtin-tools は実装済みだが、Pod 側でファイル操作を追跡する口がない。
|
||||
|
||||
ツール名ベースのヒューリスティック検出("read" / "edit" / "write" の名前で判別)は
|
||||
脆弱でユーザー定義ツールに対応できない。各ツールが自己申告する形にする。
|
||||
|
||||
## 方針
|
||||
|
||||
`ToolOutput` に optional な `referenced_files: Vec<PathBuf>` を追加し、ツール自身が
|
||||
「この実行で触ったファイル」を申告する。Pod の `HookInterceptor::post_tool_call` で
|
||||
観察し、Pod 内の LRU 的な履歴に積む。Compact 時はその履歴から先頭 N 件を compact
|
||||
worker のデフォルトリファレンスとして渡す。
|
||||
|
||||
## 実装
|
||||
|
||||
### llm-worker: ToolOutput 拡張
|
||||
|
||||
```rust
|
||||
pub struct ToolOutput {
|
||||
pub summary: String,
|
||||
pub content: Option<String>,
|
||||
pub referenced_files: Vec<PathBuf>, // NEW. 空なら未申告
|
||||
}
|
||||
```
|
||||
|
||||
- `From<String>` 実装では空の Vec を入れる
|
||||
- 既存のツール実装は変更不要(デフォルトで空)
|
||||
- `ToolResultInfo` (`post_tool_call` Hook が受ける型)にも伝播させて Pod から見えるようにする
|
||||
|
||||
### builtin-tools: 各ツールで埋める
|
||||
|
||||
既に実装済みの以下のツールに `referenced_files` を埋める変更を入れる:
|
||||
|
||||
| ツール | 申告するファイル |
|
||||
|--------|------------------|
|
||||
| Read | 読んだファイルパス |
|
||||
| Write | 書いたファイルパス |
|
||||
| Edit | 編集したファイルパス |
|
||||
| Glob | (積極的には埋めない — マッチしたファイル全部は多すぎる) |
|
||||
| Grep | (同上) |
|
||||
| Bash | (不可能。空のまま) |
|
||||
|
||||
Glob/Grep は「発見」しかしていないので referenced とは扱わない。
|
||||
|
||||
### Pod: 追跡と取り出し
|
||||
|
||||
`Pod` 内に履歴バッファ:
|
||||
|
||||
```rust
|
||||
referenced_files: VecDeque<PathBuf>, // 最新 N 件, 重複排除
|
||||
```
|
||||
|
||||
- `HookInterceptor::post_tool_call` で `ToolResultInfo.referenced_files` を受け取り、
|
||||
既存エントリを先頭に移動 or 新規追加(LRU 的な挙動)
|
||||
- 容量上限(例: 20)を超えたら末尾を落とす
|
||||
- Compact 開始時に先頭 5 件を取り出して compact worker に渡す
|
||||
|
||||
## 設計ポイント
|
||||
|
||||
- **ツールが自己申告**: 外部から名前で判別しない。拡張性のため
|
||||
- **Vec<PathBuf> なので 0〜複数件**: 複数ファイルに触るツール (例: 大量置換) にも対応
|
||||
- **空 Vec 許容**: Bash や Glob のような「追跡不能 or 不適切」なツールはそのまま空
|
||||
- **Pod が LRU バッファを持つ**: Compact 時の抽出を O(1) に近づける。Worker 層は関知しない
|
||||
|
||||
## 依存
|
||||
|
||||
- builtin-tools が実装済み前提(チケット: [builtin-tools.md](builtin-tools.md) 完了後)
|
||||
|
||||
## ブロックする後続
|
||||
|
||||
- [compact-improvements.md](compact-improvements.md) — デフォルトリファレンスの抽出がこのフィールドに依存
|
||||
Loading…
Reference in New Issue
Block a user