use std::borrow::Cow;
use std::io::Read;
use std::path::{Path, PathBuf};
use serde_json::Value;
const TOKENSAVE_RESEARCH_BLOCK_REASON: &str = "STOP: Use tokensave MCP tools \
(tokensave_context, tokensave_search, tokensave_callees, tokensave_callers, \
tokensave_impact, tokensave_files, tokensave_affected) instead of agents for \
code research. Tokensave is faster and more precise for symbol relationships, \
call paths, and code structure. Only use agents for code exploration if you \
have already tried tokensave and it cannot answer the question.";
const MAX_PATTERN_LEN: usize = 200;
const CODE_EXTENSIONS: &[&str] = &[
"rs", "go", "java", "scala", "sc", "ts", "tsx", "mts", "cts", "js", "jsx", "mjs", "cjs", "py",
"pyi", "pyw", "c", "h", "cpp", "cc", "cxx", "c++", "hpp", "hh", "hxx", "h++", "inl", "ipp",
"tcc", "kt", "kts", "cs", "csx", "swift", "dart", "pas", "pp", "dpr", "php", "phtml", "rb", "rake", "gemspec", "sh", "bash", "zsh",
"proto", "ps1", "psm1", "psd1", "nix", "vb", "vbs", "lua", "zig", "m", "mm", "pl", "pm", "bat", "cmd", "f", "f90", "f95", "f03", "for", "ftn",
"cbl", "cob", "cpy", "bas", "v", "vh", "sv", "svh",
];
const CODE_DIRS: &[&str] = &[
"src", "lib", "tests", "test", "crates", "app", "internal", "pkg", "cmd", "include",
];
const CODE_TYPE_FILTERS: &[&str] = &[
"rust",
"go",
"py",
"python",
"ts",
"typescript",
"js",
"javascript",
"java",
"scala",
"kt",
"kotlin",
"c",
"cpp",
"cxx",
"swift",
"cs",
"csharp",
"dart",
"rb",
"ruby",
"php",
"lua",
"zig",
"perl",
"pascal",
"vb",
"vbnet",
"nix",
"bash",
"sh",
"shell",
"proto",
"powershell",
"ps1",
"fortran",
"cobol",
"objc",
"objective-c",
"basic",
];
#[derive(Debug, Clone, Default)]
pub struct HookEnv {
pub in_tokensave_project: bool,
pub disable_grep_hook: bool,
pub project_root: Option<PathBuf>,
}
impl HookEnv {
pub fn from_runtime() -> Self {
Self::from_runtime_at(std::env::current_dir().ok().as_deref())
}
pub fn from_runtime_at(cwd: Option<&Path>) -> Self {
let project_root = cwd.and_then(crate::config::discover_project_root);
let disable_grep_hook = std::env::var("TOKENSAVE_DISABLE_GREP_HOOK")
.is_ok_and(|v| !v.is_empty() && v != "0" && !v.eq_ignore_ascii_case("false"));
Self {
in_tokensave_project: project_root.is_some(),
disable_grep_hook,
project_root,
}
}
fn for_event(&self, event: &Value) -> Self {
let Some(cwd) = event_cwd(event) else {
return self.clone();
};
let project_root = crate::config::discover_project_root(&cwd);
Self {
in_tokensave_project: project_root.is_some(),
disable_grep_hook: self.disable_grep_hook,
project_root,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PatternShape {
BareSymbol,
WordBoundary,
Alternation,
Definition,
}
pub fn hook_pre_tool_use() {
let raw = read_stdin_to_string();
let decision = if raw.trim().is_empty() {
evaluate_hook_decision(&std::env::var("TOOL_INPUT").unwrap_or_default())
} else {
evaluate_claude_pre_tool_use(&raw)
};
if decision.is_empty() {
println!("{}", build_allow_message());
} else {
println!("{decision}");
}
}
pub fn evaluate_claude_pre_tool_use(raw: &str) -> String {
evaluate_claude_pre_tool_use_with_env(raw, &HookEnv::from_runtime())
}
pub fn evaluate_claude_pre_tool_use_with_env(raw: &str, env: &HookEnv) -> String {
let event = serde_json::from_str::<serde_json::Value>(raw).ok();
let env = event
.as_ref()
.map_or_else(|| env.clone(), |event| env.for_event(event));
let tool_input = event
.and_then(|v| v.get("tool_input").cloned())
.map_or_else(|| raw.to_string(), |ti| ti.to_string());
evaluate_hook_decision_with_env(&tool_input, &env)
}
pub fn evaluate_hook_decision(tool_input: &str) -> String {
evaluate_hook_decision_with_env(tool_input, &HookEnv::from_runtime())
}
pub fn evaluate_hook_decision_with_env(tool_input: &str, env: &HookEnv) -> String {
match evaluate_hook_decision_core(tool_input, env) {
Some(reason) => build_block_message(&reason),
None => String::new(),
}
}
fn evaluate_hook_decision_core(tool_input: &str, env: &HookEnv) -> Option<String> {
let parsed: serde_json::Value =
serde_json::from_str(tool_input).unwrap_or_else(|_| serde_json::json!({}));
if env.in_tokensave_project && !env.disable_grep_hook {
let subagent = parsed
.get("subagent_type")
.and_then(|v| v.as_str())
.filter(|s| !s.is_empty());
if subagent == Some("Explore") {
return Some(TOKENSAVE_RESEARCH_BLOCK_REASON.to_string());
}
if subagent.is_none() {
if let Some(prompt) = parsed.get("prompt").and_then(|v| v.as_str()) {
if is_code_research_prompt(prompt) {
return Some(TOKENSAVE_RESEARCH_BLOCK_REASON.to_string());
}
}
}
}
if parsed.get("pattern").is_some() {
if let Some(reason) = evaluate_grep_tool_input(&parsed, env) {
return Some(reason);
}
if let Some(reason) = evaluate_glob_tool_input(&parsed, env) {
return Some(reason);
}
}
if let Some(command) = parsed.get("command").and_then(|v| v.as_str()) {
if let Some(reason) = evaluate_bash_command(command, env) {
return Some(reason);
}
if let Some(reason) = evaluate_find_command(command, env) {
return Some(reason);
}
}
None
}
fn build_allow_message() -> String {
serde_json::json!({ "permission": "allow" }).to_string()
}
fn build_block_message(reason: &str) -> String {
serde_json::json!({
"permission": "deny",
"user_message": reason,
"agent_message": reason,
"hookSpecificOutput": {
"hookEventName": "PreToolUse",
"permissionDecision": "deny",
"permissionDecisionReason": reason,
}
})
.to_string()
}
fn evaluate_grep_tool_input(parsed: &Value, env: &HookEnv) -> Option<String> {
if !env.in_tokensave_project || env.disable_grep_hook {
return None;
}
let pattern = parsed.get("pattern").and_then(|v| v.as_str())?;
if pattern.is_empty() || pattern.len() > MAX_PATTERN_LEN {
return None;
}
if parsed.get("output_mode").and_then(|v| v.as_str()) != Some("content") {
return None;
}
let path = parsed.get("path").and_then(|v| v.as_str()).unwrap_or("");
let glob = parsed
.get("glob")
.or_else(|| parsed.get("glob_pattern"))
.and_then(|v| v.as_str())
.unwrap_or("");
let ty = parsed.get("type").and_then(|v| v.as_str()).unwrap_or("");
if !target_looks_like_code(path, &[glob], ty, env) {
return None;
}
let shape = classify_symbol_pattern(pattern)?;
Some(redirect_message("Grep", pattern, shape))
}
fn is_inert_command(segment: &str, before_search: bool) -> bool {
const INERT: [&str; 5] = ["echo", "pwd", "ls", "cat", "printf"];
const EXIT_STATUS_ONLY: [&str; 2] = ["true", ":"];
let rest = strip_command_prefixes(segment.trim()).rest;
let head = rest.split_whitespace().next().unwrap_or("");
INERT.contains(&head) || (before_search && EXIT_STATUS_ONLY.contains(&head))
}
fn split_top_level_segments(command: &str) -> Option<Vec<&str>> {
if command.contains("$(") || command.contains('`') || command.contains('\n') {
return None;
}
let mut segments: Vec<&str> = Vec::new();
let mut start = 0usize;
let mut in_single = false;
let mut in_double = false;
let mut chars = command.char_indices().peekable();
while let Some((i, c)) = chars.next() {
if in_single {
if c == '\'' {
in_single = false;
}
continue;
}
if in_double {
if c == '\\' {
chars.next();
} else if c == '"' {
in_double = false;
}
continue;
}
match c {
'\'' => in_single = true,
'"' => in_double = true,
'\\' if !cfg!(windows) => {
chars.next();
}
'(' | ')' | '`' | '<' | '>' => return None,
'&' | '|' => {
if chars.peek().map(|&(_, next)| next) != Some(c) {
return None;
}
chars.next();
segments.push(&command[start..i]);
start = i + 2;
}
';' => {
segments.push(&command[start..i]);
start = i + 1;
}
_ => {}
}
}
if in_single || in_double {
return None;
}
segments.push(&command[start..]);
let segments: Vec<&str> = segments
.into_iter()
.map(str::trim)
.filter(|s| !s.is_empty())
.collect();
if segments.len() < 2 {
None
} else {
Some(segments)
}
}
fn evaluate_bash_command(command: &str, env: &HookEnv) -> Option<String> {
if !env.in_tokensave_project || env.disable_grep_hook {
return None;
}
if let Some(reason) = evaluate_bash_segment(command, env) {
return Some(reason);
}
let segments = split_top_level_segments(command)?;
let mut reason = None;
for segment in &segments {
if let Some(found) = evaluate_bash_segment(segment, env) {
reason.get_or_insert(found);
} else if !is_inert_command(segment, reason.is_none()) {
return None;
}
}
reason
}
fn evaluate_bash_segment(command: &str, env: &HookEnv) -> Option<String> {
let stripped = strip_command_prefixes(command.trim());
if stripped.disables_hook {
return None;
}
let inv = extract_grep_invocation(command)?;
if inv.pattern.is_empty() || inv.pattern.len() > MAX_PATTERN_LEN {
return None;
}
let target = inv.targets.first().map_or("", String::as_str);
let target =
if let (Some(cd_path), Some(root)) = (stripped.cd_target, env.project_root.as_deref()) {
let cd_path = unquote(cd_path);
let cd_path = unescape_shell_backslashes(cd_path);
let cd_path = expand_home_prefix(&cd_path, crate::agents::home_dir().as_deref())
.unwrap_or_else(|| PathBuf::from(cd_path.as_ref()));
match classify_path_within_project(&cd_path.to_string_lossy(), Some(root)) {
Containment::Outside => return None,
Containment::Inside => {
let cd_base = root
.join(&cd_path)
.canonicalize()
.unwrap_or_else(|_| root.join(&cd_path));
let effective = if target.is_empty() || target == "." || target == "./" {
cd_base
} else {
cd_base.join(target)
};
Cow::Owned(effective.to_string_lossy().into_owned())
}
Containment::Unknown => Cow::Borrowed(target),
}
} else {
Cow::Borrowed(target)
};
let globs: Vec<&str> = inv.globs.iter().map(String::as_str).collect();
if !target_looks_like_code(
target.as_ref(),
&globs,
inv.ty.as_deref().unwrap_or(""),
env,
) {
return None;
}
let shape = classify_symbol_pattern(&inv.pattern)?;
Some(redirect_message("Bash grep", &inv.pattern, shape))
}
fn evaluate_find_command(command: &str, env: &HookEnv) -> Option<String> {
if !env.in_tokensave_project || env.disable_grep_hook {
return None;
}
let stripped = strip_command_prefixes(command.trim());
if stripped.disables_hook {
return None;
}
let inv = extract_find_invocation(command)?;
let targets =
if let (Some(cd_path), Some(root)) = (stripped.cd_target, env.project_root.as_deref()) {
let cd_path = unquote(cd_path);
let cd_path = unescape_shell_backslashes(cd_path);
let cd_path = expand_home_prefix(&cd_path, crate::agents::home_dir().as_deref())
.unwrap_or_else(|| PathBuf::from(cd_path.as_ref()));
match classify_path_within_project(&cd_path.to_string_lossy(), Some(root)) {
Containment::Outside => return None,
Containment::Inside => {
let cd_base = root
.join(&cd_path)
.canonicalize()
.unwrap_or_else(|_| root.join(&cd_path));
if inv.targets.is_empty() {
vec![cd_base.to_string_lossy().into_owned()]
} else {
inv.targets
.iter()
.map(|t| {
if t.is_empty() || t == "." || t == "./" {
cd_base.to_string_lossy().into_owned()
} else {
cd_base.join(t).to_string_lossy().into_owned()
}
})
.collect()
}
}
Containment::Unknown => inv.targets,
}
} else {
inv.targets
};
if !targets.is_empty()
&& !targets
.iter()
.all(|target| target_looks_like_code(target, &[], "", env))
{
return None;
}
if !inv
.globs
.iter()
.all(|glob| classify_glob_target(glob) == Some(true))
{
return None;
}
Some(files_redirect_message("Bash find", &inv.globs.join(", ")))
}
fn evaluate_glob_tool_input(parsed: &Value, env: &HookEnv) -> Option<String> {
if !env.in_tokensave_project || env.disable_grep_hook {
return None;
}
if parsed.get("output_mode").is_some()
|| parsed.get("glob").is_some()
|| parsed.get("glob_pattern").is_some()
|| parsed.get("type").is_some()
{
return None;
}
let pattern = parsed.get("pattern").and_then(|v| v.as_str())?;
if pattern.is_empty() || pattern.len() > MAX_PATTERN_LEN {
return None;
}
if !pattern.contains('*') && !pattern.contains('?') {
return None;
}
if classify_glob_target(pattern) != Some(true) {
return None;
}
let path = parsed.get("path").and_then(|v| v.as_str()).unwrap_or("");
if !target_looks_like_code(path, &[], "", env) {
return None;
}
Some(files_redirect_message("Glob", pattern))
}
fn files_redirect_message(tool_label: &str, pattern: &str) -> String {
format!(
"STOP: This {tool_label} searches a tokensave-indexed project for files matching \
`{pattern}`. Use tokensave_files(pattern=\"{pattern}\") instead — it answers from the \
index, honors the project's ignore rules, and covers non-code artifacts (specs, \
schemas, fixtures) as well as source. To override for this one call, set \
TOKENSAVE_DISABLE_GREP_HOOK=1 in the shell."
)
}
fn redirect_message(tool_label: &str, pattern: &str, shape: PatternShape) -> String {
let suggestion = match shape {
PatternShape::BareSymbol | PatternShape::WordBoundary => {
"tokensave_search (definition) or tokensave_callers_for (usages)"
}
PatternShape::Definition => "tokensave_search (definition)",
PatternShape::Alternation => {
"tokensave_signature_search (multiple names at once) or repeated tokensave_search calls"
}
};
format!(
"STOP: This {tool_label} targets a code file in a tokensave-indexed project and the \
pattern `{pattern}` looks like a symbol name. Use {suggestion} instead — symbol-indexed \
lookups are faster and more accurate than text grep. To override for this one call, set \
TOKENSAVE_DISABLE_GREP_HOOK=1 in the shell."
)
}
fn classify_symbol_pattern(pattern: &str) -> Option<PatternShape> {
let mut p = pattern;
let mut had_wb = false;
if let Some(rest) = p.strip_prefix("\\b") {
if let Some(rest2) = rest.strip_suffix("\\b") {
p = rest2;
had_wb = true;
}
}
let normalized = p.replace("\\|", "|");
let parts: Vec<&str> = normalized.split('|').collect();
if !parts.iter().all(|s| is_pure_identifier(s)) {
return classify_definition_pattern(p);
}
match (parts.len(), had_wb) {
(1, true) => Some(PatternShape::WordBoundary),
(1, false) => Some(PatternShape::BareSymbol),
_ => Some(PatternShape::Alternation),
}
}
const DEFINITION_KEYWORDS: &[&str] = &[
"def",
"class",
"fn",
"func",
"function",
"struct",
"enum",
"trait",
"interface",
"impl",
"type",
"module",
"package",
];
fn classify_definition_pattern(pattern: &str) -> Option<PatternShape> {
let mut rest = pattern.trim();
rest = rest.strip_prefix('^').unwrap_or(rest);
rest = rest.trim_start();
let mut had_keyword = false;
let mut had_paren = false;
for kw in DEFINITION_KEYWORDS {
if let Some(tail) = rest.strip_prefix(kw) {
if tail.starts_with(|c: char| c.is_whitespace()) {
rest = tail.trim_start();
had_keyword = true;
break;
}
}
}
if let Some(head) = rest.strip_suffix('(') {
rest = head.strip_suffix('\\').unwrap_or(head);
had_paren = true;
}
rest = rest.trim_end();
if !(had_keyword || had_paren) || !is_pure_identifier(rest) {
return None;
}
if had_keyword {
Some(PatternShape::Definition)
} else {
Some(PatternShape::BareSymbol)
}
}
fn is_pure_identifier(s: &str) -> bool {
let mut chars = s.chars();
let Some(first) = chars.next() else {
return false;
};
if !(first.is_ascii_alphabetic() || first == '_') {
return false;
}
chars.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == ':')
}
fn target_looks_like_code(path: &str, globs: &[&str], ty: &str, env: &HookEnv) -> bool {
let mut known_inside = false;
if !path.is_empty() && env.project_root.is_some() {
let raw = path.trim_matches(|c: char| c.is_whitespace() || c == '"' || c == '\'');
match classify_path_within_project(raw, env.project_root.as_deref()) {
Containment::Outside => return false,
Containment::Inside => {
if path_is_config_excluded(raw, env.project_root.as_deref()) {
return false;
}
known_inside = true;
}
Containment::Unknown => {}
}
}
if !ty.is_empty() {
return CODE_TYPE_FILTERS.contains(&ty.to_ascii_lowercase().as_str());
}
let glob_verdicts: Vec<bool> = globs
.iter()
.filter_map(|g| classify_glob_target(g))
.collect();
if glob_verdicts.iter().any(|is_code| !is_code) {
return false;
}
if !glob_verdicts.is_empty() {
return true;
}
let raw = if path.is_empty() {
globs.first().copied().unwrap_or("")
} else {
path
};
let trimmed = raw.trim_matches(|c: char| c.is_whitespace() || c == '"' || c == '\'');
if trimmed.is_empty() || trimmed == "." || trimmed == "./" {
return true;
}
if path_is_project_root(trimmed, env.project_root.as_deref()) {
return true;
}
let file_part = trimmed
.trim_end_matches(std::path::is_separator)
.rsplit(std::path::is_separator)
.next()
.unwrap_or(trimmed);
if let Some(idx) = file_part.rfind('.') {
let after_dot = &file_part[idx + 1..];
let ext: String = after_dot
.chars()
.take_while(|c| c.is_ascii_alphanumeric() || *c == '+')
.collect::<String>()
.to_ascii_lowercase();
if !ext.is_empty() {
return CODE_EXTENSIONS.contains(&ext.as_str());
}
}
if known_inside && dir_holds_code_files(trimmed) {
return true;
}
let last = trimmed
.trim_end_matches(std::path::is_separator)
.rsplit(std::path::is_separator)
.next()
.unwrap_or("");
CODE_DIRS.contains(&last)
}
const CODE_FILE_SCAN_BUDGET: usize = 2_000;
fn dir_holds_code_files(path: &str) -> bool {
let start = PathBuf::from(path);
if start.is_file() {
return true;
}
let mut queue = std::collections::VecDeque::from([start]);
let mut seen = 0usize;
while let Some(dir) = queue.pop_front() {
let Ok(entries) = std::fs::read_dir(&dir) else {
continue;
};
for entry in entries.flatten() {
seen += 1;
if seen > CODE_FILE_SCAN_BUDGET {
return false;
}
let name = entry.file_name();
let name = name.to_string_lossy();
if name.starts_with('.') {
continue;
}
match entry.file_type() {
Ok(ft) if ft.is_dir() => queue.push_back(entry.path()),
Ok(ft) if ft.is_file() => {
let ext = name
.rsplit_once('.')
.map(|(_, e)| e.to_ascii_lowercase())
.unwrap_or_default();
if !ext.is_empty() && CODE_EXTENSIONS.contains(&ext.as_str()) {
return true;
}
}
_ => {}
}
}
}
false
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Containment {
Inside,
Outside,
Unknown,
}
fn classify_path_containment_with_home(
raw: &str,
project_root: Option<&Path>,
home: Option<&Path>,
) -> Containment {
let Some(root) = project_root else {
return Containment::Unknown;
};
let Some(target) = expand_home_prefix(raw, home) else {
return Containment::Unknown;
};
if target.as_os_str().is_empty() || target == Path::new(".") || target == Path::new("./") {
return Containment::Inside;
}
let resolved = if target.is_absolute() {
target
} else {
root.join(target)
};
match (resolved.canonicalize(), root.canonicalize()) {
(Ok(target), Ok(root)) => {
if target.starts_with(&root) {
Containment::Inside
} else {
Containment::Outside
}
}
_ => Containment::Unknown,
}
}
fn path_is_config_excluded(raw: &str, project_root: Option<&Path>) -> bool {
let Some(root) = project_root else {
return false;
};
let Ok(config) = crate::config::load_config(root) else {
return false;
};
path_is_config_excluded_with(raw, root, &config, crate::agents::home_dir().as_deref())
}
fn path_is_config_excluded_with(
raw: &str,
root: &Path,
config: &crate::config::TokenSaveConfig,
home: Option<&Path>,
) -> bool {
let Some(target) = expand_home_prefix(raw, home) else {
return false;
};
let resolved = if target.is_absolute() {
target
} else {
root.join(target)
};
let (Ok(target), Ok(root)) = (resolved.canonicalize(), root.canonicalize()) else {
return false;
};
let Ok(relative) = target.strip_prefix(&root) else {
return false;
};
let relative = relative.to_string_lossy().replace('\\', "/");
if relative.is_empty() {
return false;
}
crate::config::is_excluded(&relative, config)
|| crate::config::is_excluded_dir(&relative, config)
}
fn classify_path_within_project(raw: &str, project_root: Option<&Path>) -> Containment {
classify_path_containment_with_home(raw, project_root, crate::agents::home_dir().as_deref())
}
#[cfg(test)]
fn path_is_within_project_with_home(
raw: &str,
project_root: Option<&Path>,
home: Option<&Path>,
) -> bool {
matches!(
classify_path_containment_with_home(raw, project_root, home),
Containment::Inside
)
}
fn path_is_project_root(raw: &str, project_root: Option<&Path>) -> bool {
path_is_project_root_with_home(raw, project_root, crate::agents::home_dir().as_deref())
}
fn path_is_project_root_with_home(
raw: &str,
project_root: Option<&Path>,
home: Option<&Path>,
) -> bool {
let Some(root) = project_root else {
return false;
};
let Some(target) = expand_home_prefix(raw, home) else {
return false;
};
if !target.is_absolute() {
return false;
}
match (target.canonicalize(), root.canonicalize()) {
(Ok(target), Ok(root)) => target == root,
_ => false,
}
}
fn expand_home_prefix(raw: &str, home: Option<&Path>) -> Option<PathBuf> {
let trimmed = raw.trim_matches(|c: char| c.is_whitespace() || c == '"' || c == '\'');
if trimmed.is_empty() {
return None;
}
if trimmed == "~" {
return home.map(Path::to_path_buf);
}
if let Some(rest) = trimmed.strip_prefix("~/") {
return home.map(|h| h.join(rest));
}
Some(PathBuf::from(trimmed))
}
fn classify_glob_target(glob: &str) -> Option<bool> {
let trimmed = glob.trim_matches(|c: char| c.is_whitespace() || c == '"' || c == '\'');
if trimmed.is_empty() {
return None;
}
let file_glob = trimmed.rsplit('/').next().unwrap_or(trimmed);
let extensions = if file_glob.ends_with('}') {
let brace_start = file_glob.rfind(".{")?;
let values = &file_glob[brace_start + 2..file_glob.len() - 1];
let extensions = values
.split(',')
.map(str::trim)
.map(|ext| {
(!ext.is_empty() && ext.chars().all(|c| c.is_ascii_alphanumeric() || c == '+'))
.then(|| ext.to_ascii_lowercase())
})
.collect::<Option<Vec<_>>>()?;
extensions
} else {
let idx = file_glob.rfind('.')?;
let ext = file_glob[idx + 1..]
.chars()
.take_while(|c| c.is_ascii_alphanumeric() || *c == '+')
.collect::<String>()
.to_ascii_lowercase();
if ext.is_empty() {
return None;
}
vec![ext]
};
(!extensions.is_empty()).then(|| {
extensions
.iter()
.all(|ext| CODE_EXTENSIONS.contains(&ext.as_str()))
})
}
#[derive(Debug, Default)]
struct GrepInvocation {
pattern: String,
targets: Vec<String>,
globs: Vec<String>,
ty: Option<String>,
}
#[derive(Debug, Default)]
struct FindInvocation {
globs: Vec<String>,
targets: Vec<String>,
}
fn has_chained_command(command: &str) -> bool {
let mut in_single = false;
let mut in_double = false;
let mut chars = command.chars().peekable();
while let Some(c) = chars.next() {
if in_single {
if c == '\'' {
in_single = false;
}
} else if in_double {
if c == '\\' {
chars.next();
} else if c == '"' {
in_double = false;
} else if c == '`' || (c == '$' && chars.peek() == Some(&'(')) {
return true;
}
} else {
match c {
'\'' => in_single = true,
'"' => in_double = true,
'\\' if !cfg!(windows) => {
chars.next();
}
'$' if chars.peek() == Some(&'(') => return true,
'2' if chars.peek() == Some(&'>') => {
chars.next();
if chars.peek() == Some(&'>') {
chars.next();
}
}
'&' | '|' | ';' | '`' | '>' => return true,
_ => {}
}
}
}
false
}
fn extract_find_invocation(command: &str) -> Option<FindInvocation> {
let rest = strip_command_prefixes(command.trim()).rest;
if has_chained_command(rest) {
return None;
}
let (is_find, after_tool) = rest
.strip_prefix("find ")
.map(|after| (true, after))
.or_else(|| rest.strip_prefix("fd ").map(|after| (false, after)))?;
let mut inv = FindInvocation::default();
let mut iter = shell_split(after_tool).into_iter().peekable();
while let Some(tok) = iter.next() {
match tok.as_str() {
"-name" | "-iname" if is_find => {
let Some(glob) = iter.next() else { continue };
inv.globs.push(glob);
}
"-e" | "--extension" if !is_find => {
let Some(ext) = iter.next() else { continue };
inv.globs.push(format!("*.{ext}"));
}
"-g" | "--glob" if !is_find => {
let Some(glob) = iter.next() else { continue };
inv.globs.push(glob);
}
"-type" | "-maxdepth" | "-mindepth" if is_find => {
iter.next();
}
"-print" | "-print0" | "-follow" if is_find => {}
"-t" | "--type" | "-d" | "--max-depth" | "-E" | "--exclude" if !is_find => {
iter.next();
}
_ if is_find && tok.starts_with('-') => return None,
_ if tok.starts_with('-') => {}
_ => inv.targets.push(tok),
}
}
if !is_find && !inv.targets.is_empty() && inv.globs.is_empty() {
return None;
}
(!inv.globs.is_empty()).then_some(inv)
}
fn extract_grep_invocation(command: &str) -> Option<GrepInvocation> {
let rest = strip_command_prefixes(command.trim()).rest;
if has_chained_command(rest) {
return None;
}
let after_tool = ["grep ", "rg ", "ag "]
.iter()
.find_map(|prefix| rest.strip_prefix(prefix))?;
let tokens = shell_split(after_tool);
let mut pattern: Option<String> = None;
let mut targets: Vec<String> = Vec::new();
let mut globs: Vec<String> = Vec::new();
let mut ty: Option<String> = None;
let mut iter = tokens.into_iter().peekable();
while let Some(tok) = iter.next() {
if tok.starts_with('-') {
if (tok == "-e" || tok == "--regexp") && pattern.is_none() {
if let Some(p) = iter.next() {
pattern = Some(p);
}
} else if let Some(p) = tok.strip_prefix("--regexp=") {
if pattern.is_none() {
pattern = Some(p.to_string());
}
} else if tok == "--include" || tok == "-g" || tok == "--glob" || tok == "--iglob" {
if let Some(g) = iter.next() {
globs.push(g);
}
} else if let Some(g) = tok
.strip_prefix("--include=")
.or_else(|| tok.strip_prefix("--glob="))
.or_else(|| tok.strip_prefix("--iglob="))
{
globs.push(g.to_string());
} else if tok == "-t" || tok == "--type" {
if let Some(t) = iter.next() {
ty.get_or_insert(t);
}
} else if let Some(t) = tok.strip_prefix("--type=") {
ty.get_or_insert(t.to_string());
} else if tok == "--exclude"
|| tok == "--exclude-dir"
|| tok == "-T"
|| tok == "--type-not"
{
iter.next();
}
continue;
}
if pattern.is_none() {
pattern = Some(tok);
} else {
targets.push(tok);
}
}
Some(GrepInvocation {
pattern: pattern?,
targets,
globs,
ty,
})
}
struct StrippedCommand<'a> {
rest: &'a str,
disables_hook: bool,
cd_target: Option<&'a str>,
}
fn strip_command_prefixes(command: &str) -> StrippedCommand<'_> {
let mut rest = command.trim_start();
let mut disables_hook = false;
let mut cd_target: Option<&str> = None;
loop {
let mut advanced = false;
for prefix in ["rtk ", "sudo ", "time ", "nice "] {
if let Some(after) = rest.strip_prefix(prefix) {
rest = after.trim_start();
advanced = true;
}
}
if let Some((name, value, after)) = parse_leading_env_assignment(rest) {
if name == "TOKENSAVE_DISABLE_GREP_HOOK" {
disables_hook = disable_value_is_truthy(unquote(value));
}
rest = after.trim_start();
advanced = true;
}
if let Some((cd_arg, after)) = strip_leading_cd(rest) {
if cd_target.is_none() {
cd_target = Some(cd_arg);
}
rest = after.trim_start();
advanced = true;
}
if !advanced {
return StrippedCommand {
rest,
disables_hook,
cd_target,
};
}
}
}
fn disable_value_is_truthy(value: &str) -> bool {
!value.is_empty() && value != "0" && !value.eq_ignore_ascii_case("false")
}
fn unquote(v: &str) -> &str {
let b = v.as_bytes();
if b.len() >= 2 && (b[0] == b'\'' || b[0] == b'"') && b[b.len() - 1] == b[0] {
&v[1..v.len() - 1]
} else {
v
}
}
fn unescape_shell_backslashes(s: &str) -> Cow<'_, str> {
if cfg!(windows) {
return Cow::Borrowed(s);
}
let mut out = String::with_capacity(s.len());
let mut chars = s.chars().peekable();
let mut changed = false;
while let Some(c) = chars.next() {
if c == '\\' {
if let Some(next) = chars.next() {
out.push(next);
changed = true;
continue;
}
}
out.push(c);
}
if changed {
Cow::Owned(out)
} else {
Cow::Borrowed(s)
}
}
fn parse_leading_env_assignment(s: &str) -> Option<(&str, &str, &str)> {
let mut chars = s.char_indices();
let (_, first) = chars.next()?;
if !(first.is_ascii_alphabetic() || first == '_') {
return None;
}
let mut eq_pos = None;
for (idx, c) in chars {
if c == '=' {
eq_pos = Some(idx);
break;
}
if !(c.is_ascii_alphanumeric() || c == '_') {
return None;
}
}
let eq = eq_pos?;
let value_start = eq + 1;
let mut in_single = false;
let mut in_double = false;
for (idx, c) in s[value_start..].char_indices() {
match c {
'\'' if !in_double => in_single = !in_single,
'"' if !in_single => in_double = !in_double,
c if c.is_whitespace() && !in_single && !in_double => {
let value = &s[value_start..value_start + idx];
return Some((&s[..eq], value, &s[value_start + idx..]));
}
_ => {}
}
}
None
}
fn strip_leading_cd(s: &str) -> Option<(&str, &str)> {
let after = s.strip_prefix("cd")?;
if !after.starts_with(char::is_whitespace) {
return None;
}
let mut in_single = false;
let mut in_double = false;
let mut iter = s.char_indices().peekable();
while let Some((idx, c)) = iter.next() {
if in_single {
if c == '\'' {
in_single = false;
}
continue;
}
if in_double {
if c == '"' {
in_double = false;
}
continue;
}
match c {
'\'' => in_single = true,
'"' => in_double = true,
'|' => return None,
';' => {
let cd_arg = extract_cd_argument(s)?;
return Some((cd_arg, &s[idx + c.len_utf8()..]));
}
'&' => {
if let Some(&(idx2, '&')) = iter.peek() {
let cd_arg = extract_cd_argument(s)?;
return Some((cd_arg, &s[idx2 + 1..]));
}
return None;
}
_ => {}
}
}
None
}
fn extract_cd_argument(s: &str) -> Option<&str> {
let after = s.strip_prefix("cd")?.trim_start();
if after.is_empty() {
return None;
}
let mut in_single = false;
let mut in_double = false;
let mut iter = after.char_indices().peekable();
let mut last_end = after.len();
while let Some((idx, c)) = iter.next() {
if in_single {
if c == '\'' {
in_single = false;
}
continue;
}
if in_double {
if c == '"' {
in_double = false;
}
continue;
}
match c {
'\'' => in_single = true,
'"' => in_double = true,
';' | '|' => {
last_end = idx;
break;
}
'&' => {
if let Some(&(_, '&')) = iter.peek() {
last_end = idx;
break;
}
}
_ => {}
}
}
let path = after[..last_end].trim_end();
if path.is_empty() {
None
} else {
Some(path)
}
}
fn shell_split(s: &str) -> Vec<String> {
let mut out: Vec<String> = Vec::new();
let mut cur = String::new();
let mut in_single = false;
let mut in_double = false;
let mut chars = s.chars().peekable();
while let Some(c) = chars.next() {
if in_single {
if c == '\'' {
in_single = false;
} else {
cur.push(c);
}
} else if in_double {
if c == '"' {
in_double = false;
} else if c == '\\' {
if let Some(&next) = chars.peek() {
if matches!(next, '"' | '\\' | '$' | '`') {
chars.next();
cur.push(next);
continue;
}
}
cur.push(c);
} else {
cur.push(c);
}
} else {
match c {
'\'' => in_single = true,
'"' => in_double = true,
'\\' => {
if cfg!(windows) {
cur.push(c);
} else if let Some(next) = chars.next() {
cur.push(next);
}
}
'|' | ';' | '&' | '>' | '<' => break,
c if c.is_whitespace() => {
if !cur.is_empty() {
out.push(std::mem::take(&mut cur));
}
}
c => cur.push(c),
}
}
}
if !cur.is_empty() {
out.push(cur);
}
out
}
fn is_code_research_prompt(prompt: &str) -> bool {
let lower = prompt.to_ascii_lowercase();
let exploration_patterns = [
"explore",
"codebase structure",
"codebase architecture",
"codebase overview",
"source files contents",
"read every",
"full contents",
"entire codebase",
"architecture and structure",
"call graph",
"call path",
"call chain",
"symbol relat",
"symbol lookup",
"who calls",
"callers of",
"callees of",
];
exploration_patterns.iter().any(|pat| lower.contains(pat))
}
pub fn hook_kiro_pre_tool_use() -> i32 {
let event = read_stdin_to_string();
if let Some(reason) = evaluate_kiro_pre_tool_use_with_env(&event, &HookEnv::from_runtime()) {
eprintln!("{reason}");
2
} else {
0
}
}
pub fn evaluate_kiro_pre_tool_use(event_json: &str) -> Option<&'static str> {
evaluate_kiro_pre_tool_use_with_env(event_json, &HookEnv::from_runtime())
}
pub fn evaluate_kiro_pre_tool_use_with_env(
event_json: &str,
env: &HookEnv,
) -> Option<&'static str> {
if env.disable_grep_hook {
return None;
}
let parsed: Value = serde_json::from_str(event_json).ok()?;
if !env.for_event(&parsed).in_tokensave_project {
return None;
}
let tool_name = parsed.get("tool_name").and_then(Value::as_str)?;
if !is_kiro_delegation_tool(tool_name) {
return None;
}
if kiro_event_has_research_text(parsed.get("tool_input").unwrap_or(&Value::Null)) {
Some(TOKENSAVE_RESEARCH_BLOCK_REASON)
} else {
None
}
}
fn is_kiro_delegation_tool(tool_name: &str) -> bool {
matches!(tool_name, "delegate" | "subagent" | "use_subagent")
}
pub fn hook_droid_pre_tool_use() -> i32 {
let event = read_stdin_to_string();
if let Some(reason) = evaluate_droid_pre_tool_use(&event) {
eprintln!("{reason}");
2
} else {
0
}
}
pub fn evaluate_droid_pre_tool_use(raw: &str) -> Option<String> {
evaluate_droid_pre_tool_use_with_env(raw, &HookEnv::from_runtime())
}
pub fn evaluate_droid_pre_tool_use_with_env(raw: &str, env: &HookEnv) -> Option<String> {
let Ok(event) = serde_json::from_str::<Value>(raw) else {
return evaluate_hook_decision_core(raw, env);
};
let is_task = event.get("tool_name").and_then(Value::as_str) == Some("Task");
let env = env.for_event(&event);
let mut tool_input = event.get("tool_input").cloned().unwrap_or(event);
let subagent_type = tool_input.get("subagent_type").and_then(Value::as_str);
if is_task && subagent_type != Some("explorer") {
return None;
}
if subagent_type == Some("explorer") {
tool_input["subagent_type"] = Value::String("Explore".to_string());
}
evaluate_hook_decision_core(&tool_input.to_string(), &env)
}
fn kiro_event_has_research_text(value: &Value) -> bool {
let mut text = Vec::new();
collect_kiro_task_strings(value, &mut text);
if text.is_empty() {
collect_strings(value, &mut text);
}
text.iter().any(|s| is_code_research_prompt(s))
}
fn collect_kiro_task_strings<'a>(value: &'a Value, out: &mut Vec<&'a str>) {
match value {
Value::Object(map) => {
for (key, child) in map {
let key = key.to_ascii_lowercase();
if key.contains("prompt")
|| key.contains("task")
|| key.contains("query")
|| key.contains("instruction")
|| key.contains("message")
|| key.contains("description")
{
collect_strings(child, out);
} else {
collect_kiro_task_strings(child, out);
}
}
}
Value::Array(items) => {
for item in items {
collect_kiro_task_strings(item, out);
}
}
Value::String(s) => out.push(s),
_ => {}
}
}
fn collect_strings<'a>(value: &'a Value, out: &mut Vec<&'a str>) {
match value {
Value::String(s) => out.push(s),
Value::Array(items) => {
for item in items {
collect_strings(item, out);
}
}
Value::Object(map) => {
for child in map.values() {
collect_strings(child, out);
}
}
_ => {}
}
}
pub async fn hook_prompt_submit() {
let project_path = crate::config::resolve_path(None);
if let Ok(cg) = crate::tokensave::TokenSave::open(&project_path).await {
let _ = cg.reset_local_counter().await;
}
}
pub async fn hook_kiro_prompt_submit() -> i32 {
let event = read_stdin_to_string();
reset_counter_for_kiro_event(&event).await;
0
}
pub async fn hook_kiro_post_tool_use() -> i32 {
let event = read_stdin_to_string();
match sync_for_kiro_event(&event).await {
Ok(()) => 0,
Err(e) => {
eprintln!("tokensave sync failed: {e}");
1
}
}
}
async fn reset_counter_for_kiro_event(event_json: &str) {
let Some(project_root) = kiro_project_root(event_json) else {
return;
};
if let Ok(cg) = crate::tokensave::TokenSave::open(&project_root).await {
let _ = cg.reset_local_counter().await;
}
}
async fn sync_for_kiro_event(event_json: &str) -> crate::errors::Result<()> {
let Some(project_root) = kiro_project_root(event_json) else {
return Ok(());
};
let cg = crate::tokensave::TokenSave::open(&project_root).await?;
match cg.sync().await {
Ok(_) | Err(crate::errors::TokenSaveError::SyncLock { .. }) => Ok(()),
Err(e) => Err(e),
}
}
fn kiro_project_root(event_json: &str) -> Option<PathBuf> {
let cwd = kiro_event_cwd(event_json).or_else(|| std::env::current_dir().ok())?;
crate::config::discover_project_root(&cwd)
}
fn kiro_event_cwd(event_json: &str) -> Option<PathBuf> {
event_cwd(&serde_json::from_str::<Value>(event_json).ok()?)
}
fn event_cwd(event: &Value) -> Option<PathBuf> {
let path = Path::new(event.get("cwd").and_then(Value::as_str)?.trim());
(path.is_absolute() && path.is_dir()).then(|| path.to_path_buf())
}
fn read_stdin_to_string() -> String {
let mut input = String::new();
let _ = std::io::stdin().read_to_string(&mut input);
input
}
pub async fn hook_stop() {
let Some(gdb) = crate::global_db::GlobalDb::open().await else {
return;
};
let stats = crate::accounting::parser::ingest_claude_only(&gdb).await;
if stats.turns_inserted == 0 {
return;
}
let project_path = crate::config::resolve_path(None);
let tokens_saved = if let Ok(cg) = crate::tokensave::TokenSave::open(&project_path).await {
cg.get_tokens_saved().await.unwrap_or(0)
} else {
0
};
let efficiency = if tokens_saved + stats.tokens_consumed > 0 {
(tokens_saved as f64 / (tokens_saved + stats.tokens_consumed) as f64) * 100.0
} else {
0.0
};
let saved_str = crate::display::format_token_count(tokens_saved);
if stats.cost_usd >= 0.001 {
eprintln!(
"\x1b[36mSession: ${:.2} spent | {saved_str} saved | {efficiency:.0}% efficiency\x1b[0m",
stats.cost_usd
);
}
}
#[cfg(test)]
mod cursor_decision_tests {
#![allow(clippy::unwrap_used, clippy::expect_used)]
use super::{build_allow_message, build_block_message};
use serde_json::Value;
#[test]
fn allow_message_carries_cursor_permission_field() {
let v: Value = serde_json::from_str(&build_allow_message()).unwrap();
assert_eq!(v["permission"].as_str(), Some("allow"));
}
#[test]
fn block_message_is_cross_harness() {
let v: Value = serde_json::from_str(&build_block_message("use tokensave instead")).unwrap();
assert_eq!(v["permission"].as_str(), Some("deny"));
assert_eq!(v["user_message"].as_str(), Some("use tokensave instead"));
assert_eq!(v["agent_message"].as_str(), Some("use tokensave instead"));
assert_eq!(
v["hookSpecificOutput"]["permissionDecision"].as_str(),
Some("deny")
);
assert_eq!(
v["hookSpecificOutput"]["permissionDecisionReason"].as_str(),
Some("use tokensave instead")
);
}
}
#[cfg(test)]
mod project_root_target_tests {
#![allow(clippy::unwrap_used, clippy::expect_used)]
use super::{path_is_project_root_with_home, path_is_within_project_with_home};
fn indexed_project() -> (tempfile::TempDir, std::path::PathBuf) {
let tmp = tempfile::tempdir().unwrap();
let root = tmp.path().join("project");
std::fs::create_dir_all(&root).unwrap();
(tmp, root)
}
#[test]
fn tilde_alone_resolves_to_the_project_root() {
let (_tmp, root) = indexed_project();
assert!(path_is_project_root_with_home(
"~",
Some(&root),
Some(&root)
));
}
#[test]
fn tilde_subpath_resolves_to_the_project_root() {
let (tmp, root) = indexed_project();
assert!(path_is_project_root_with_home(
"~/project",
Some(&root),
Some(tmp.path())
));
}
#[test]
fn tilde_subpath_below_the_root_is_not_the_root() {
let (tmp, root) = indexed_project();
std::fs::create_dir_all(root.join("src")).unwrap();
assert!(!path_is_project_root_with_home(
"~/project/src",
Some(&root),
Some(tmp.path())
));
}
#[test]
fn tilde_without_a_known_home_fails_open() {
let (_tmp, root) = indexed_project();
assert!(!path_is_project_root_with_home("~", Some(&root), None));
}
#[test]
fn other_users_home_is_not_expanded() {
let (tmp, root) = indexed_project();
assert!(!path_is_project_root_with_home(
"~someone/project",
Some(&root),
Some(tmp.path())
));
}
#[test]
fn quoted_absolute_root_still_matches() {
let (_tmp, root) = indexed_project();
let quoted = format!("\"{}\"", root.display());
assert!(path_is_project_root_with_home(
"ed,
Some(&root),
Some(&root)
));
}
#[test]
fn relative_targets_are_left_to_the_other_rules() {
let (tmp, root) = indexed_project();
assert!(!path_is_project_root_with_home(
"project",
Some(&root),
Some(tmp.path())
));
}
#[test]
fn within_project_absolute_root_matches() {
let (_tmp, root) = indexed_project();
assert!(path_is_within_project_with_home(
root.to_str().unwrap(),
Some(&root),
None
));
}
#[test]
fn within_project_absolute_subdir_matches() {
let (_tmp, root) = indexed_project();
std::fs::create_dir_all(root.join("src")).unwrap();
let subdir = root.join("src");
assert!(path_is_within_project_with_home(
subdir.to_str().unwrap(),
Some(&root),
None
));
}
#[test]
fn within_project_absolute_outside_does_not_match() {
let tmp = tempfile::tempdir().unwrap();
let root = tmp.path().join("project");
let other = tmp.path().join("other");
std::fs::create_dir_all(&root).unwrap();
std::fs::create_dir_all(&other).unwrap();
assert!(!path_is_within_project_with_home(
other.to_str().unwrap(),
Some(&root),
None
));
}
#[test]
fn within_project_relative_subdir_resolves_against_root() {
let (_tmp, root) = indexed_project();
std::fs::create_dir_all(root.join("src")).unwrap();
assert!(path_is_within_project_with_home("src", Some(&root), None));
}
#[test]
fn within_project_relative_outside_does_not_match() {
let tmp = tempfile::tempdir().unwrap();
let root = tmp.path().join("project");
let other = tmp.path().join("other");
std::fs::create_dir_all(&root).unwrap();
std::fs::create_dir_all(&other).unwrap();
assert!(!path_is_within_project_with_home(
"../other",
Some(&root),
None
));
}
#[test]
fn within_project_tilde_resolves_to_project() {
let (tmp, root) = indexed_project();
assert!(path_is_within_project_with_home(
"~/project",
Some(&root),
Some(tmp.path())
));
}
#[test]
fn within_project_unresolvable_path_returns_false() {
let (_tmp, root) = indexed_project();
assert!(!path_is_within_project_with_home(
"/nonexistent/path/that/does/not/exist",
Some(&root),
None
));
}
#[test]
fn within_project_no_root_returns_false() {
let (_tmp, _root) = indexed_project();
assert!(!path_is_within_project_with_home("/anywhere", None, None));
}
#[test]
fn within_project_other_users_home_is_not_expanded() {
let (tmp, root) = indexed_project();
assert!(!path_is_within_project_with_home(
"~someone/project",
Some(&root),
Some(tmp.path())
));
}
}