diff --git a/hooks/rtk-post-tool-use.sh b/hooks/rtk-post-tool-use.sh new file mode 100644 index 00000000..5f096055 --- /dev/null +++ b/hooks/rtk-post-tool-use.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash + +if ! command -v jq &>/dev/null; then + exit 0 +fi + +if ! command -v rtk &>/dev/null; then + exit 0 +fi + +RTK_LOG_DIR="${RTK_LOG_DIR:-$HOME/.local/share/rtk/logs}" +mkdir -p "$RTK_LOG_DIR" +LOG_FILE="$RTK_LOG_DIR/mcp-filter.log" + +INPUT=$(cat) + +TOOL_NAME=$(echo "$INPUT" | jq -r '.tool_name // empty' 2>/dev/null) + +if [ -z "$TOOL_NAME" ]; then + exit 0 +fi + +case "$TOOL_NAME" in + mcp__*) ;; + *) exit 0 ;; +esac + +RESULT=$(echo "$INPUT" | rtk filter-mcp-output 2>>"$LOG_FILE") +EXIT_CODE=$? + +if [ $EXIT_CODE -ne 0 ]; then + echo "[$(date -u +%Y-%m-%dT%H:%M:%SZ)] filter-mcp-output failed (exit $EXIT_CODE) for tool: $TOOL_NAME" >> "$LOG_FILE" + exit 0 +fi + +if [ -z "$RESULT" ]; then + exit 0 +fi + +echo "$RESULT" diff --git a/src/hooks/init.rs b/src/hooks/init.rs index 64115a1d..a16ee1f6 100644 --- a/src/hooks/init.rs +++ b/src/hooks/init.rs @@ -10,6 +10,7 @@ use super::integrity; // Embedded hook script (guards before set -euo pipefail) const REWRITE_HOOK: &str = include_str!("../../hooks/claude/rtk-rewrite.sh"); +const POST_TOOL_USE_HOOK: &str = include_str!("../../hooks/rtk-post-tool-use.sh"); // Embedded Cursor hook script (preToolUse format) const CURSOR_REWRITE_HOOK: &str = include_str!("../../hooks/cursor/rtk-rewrite.sh"); @@ -289,6 +290,122 @@ pub fn run( Ok(()) } +fn prepare_post_tool_use_hook_path() -> Result { + let claude_dir = resolve_claude_dir()?; + let hook_dir = claude_dir.join("hooks"); + fs::create_dir_all(&hook_dir) + .with_context(|| format!("Failed to create hook directory: {}", hook_dir.display()))?; + Ok(hook_dir.join("rtk-post-tool-use.sh")) +} + +#[cfg(unix)] +fn ensure_post_tool_use_hook_installed(hook_path: &Path, verbose: u8) -> Result { + let changed = if hook_path.exists() { + let existing = fs::read_to_string(hook_path) + .with_context(|| format!("Failed to read existing hook: {}", hook_path.display()))?; + if existing == POST_TOOL_USE_HOOK { + if verbose > 0 { + eprintln!("PostToolUse hook already up to date: {}", hook_path.display()); + } + false + } else { + fs::write(hook_path, POST_TOOL_USE_HOOK) + .with_context(|| format!("Failed to write hook to {}", hook_path.display()))?; + true + } + } else { + fs::write(hook_path, POST_TOOL_USE_HOOK) + .with_context(|| format!("Failed to write hook to {}", hook_path.display()))?; + true + }; + + use std::os::unix::fs::PermissionsExt; + fs::set_permissions(hook_path, fs::Permissions::from_mode(0o755)) + .with_context(|| format!("Failed to set hook permissions: {}", hook_path.display()))?; + + Ok(changed) +} + +fn insert_post_tool_use_hook_entry(root: &mut serde_json::Value, hook_command: &str) { + let root_obj = match root.as_object_mut() { + Some(obj) => obj, + None => { + *root = serde_json::json!({}); + root.as_object_mut().expect("Just created object") + } + }; + + let hooks = root_obj + .entry("hooks") + .or_insert_with(|| serde_json::json!({})) + .as_object_mut() + .expect("hooks must be an object"); + + let post_tool_use = hooks + .entry("PostToolUse") + .or_insert_with(|| serde_json::json!([])) + .as_array_mut() + .expect("PostToolUse must be an array"); + + post_tool_use.push(serde_json::json!({ + "matcher": "mcp__.*", + "hooks": [{ + "type": "command", + "command": hook_command + }] + })); +} + +fn post_tool_use_hook_already_present(root: &serde_json::Value, hook_command: &str) -> bool { + let arr = match root + .get("hooks") + .and_then(|h| h.get("PostToolUse")) + .and_then(|p| p.as_array()) + { + Some(a) => a, + None => return false, + }; + + arr.iter() + .filter_map(|entry| entry.get("hooks")?.as_array()) + .flatten() + .filter_map(|hook| hook.get("command")?.as_str()) + .any(|cmd| { + cmd == hook_command + || (cmd.contains("rtk-post-tool-use.sh") + && hook_command.contains("rtk-post-tool-use.sh")) + }) +} + +fn remove_post_tool_use_hook_entry(root: &mut serde_json::Value, hook_command: &str) { + let hooks = match root.get_mut("hooks").and_then(|h| h.get_mut("PostToolUse")) { + Some(h) => h, + None => return, + }; + + let arr = match hooks.as_array_mut() { + Some(a) => a, + None => return, + }; + + arr.retain(|entry| { + let hooks_arr = match entry.get("hooks").and_then(|h| h.as_array()) { + Some(a) => a, + None => return true, + }; + !hooks_arr.iter().any(|hook| { + hook.get("command") + .and_then(|c| c.as_str()) + .map(|cmd| { + cmd == hook_command + || (cmd.contains("rtk-post-tool-use.sh") + && hook_command.contains("rtk-post-tool-use.sh")) + }) + .unwrap_or(false) + }) + }); +} + /// Prepare hook directory and return paths (hook_dir, hook_path) fn prepare_hook_paths() -> Result<(PathBuf, PathBuf)> { let claude_dir = resolve_claude_dir()?; @@ -522,6 +639,62 @@ fn remove_hook_from_settings(verbose: u8) -> Result { Ok(removed) } +fn remove_post_tool_use_hook_from_settings(verbose: u8) -> Result { + let claude_dir = resolve_claude_dir()?; + let settings_path = claude_dir.join("settings.json"); + + if !settings_path.exists() { + return Ok(false); + } + + let content = fs::read_to_string(&settings_path) + .with_context(|| format!("Failed to read {}", settings_path.display()))?; + + if content.trim().is_empty() { + return Ok(false); + } + + let mut root: serde_json::Value = serde_json::from_str(&content) + .with_context(|| format!("Failed to parse {} as JSON", settings_path.display()))?; + + let hook_command = resolve_claude_dir()? + .join("hooks") + .join("rtk-post-tool-use.sh"); + let hook_command = hook_command.to_string_lossy().to_string(); + + let had_entry = post_tool_use_hook_already_present(&root, &hook_command); + if !had_entry { + return Ok(false); + } + + remove_post_tool_use_hook_entry(&mut root, &hook_command); + + // Verify removal + let still_present = root + .get("hooks") + .and_then(|h| h.get("PostToolUse")) + .and_then(|p| p.as_array()) + .map(|a| !a.is_empty()) + .unwrap_or(false); + + let backup_path = settings_path.with_extension("json.bak"); + fs::copy(&settings_path, &backup_path) + .with_context(|| format!("Failed to backup to {}", backup_path.display()))?; + + let serialized = + serde_json::to_string_pretty(&root).context("Failed to serialize settings.json")?; + atomic_write(&settings_path, &serialized)?; + + if verbose > 0 { + eprintln!("Removed PostToolUse hook from settings.json"); + if still_present { + eprintln!(" Note: other PostToolUse hooks remain"); + } + } + + Ok(true) +} + /// Full uninstall for Claude, Gemini, Codex, or Cursor artifacts. pub fn uninstall(global: bool, gemini: bool, codex: bool, cursor: bool, verbose: u8) -> Result<()> { if codex { @@ -624,6 +797,19 @@ pub fn uninstall(global: bool, gemini: bool, codex: bool, cursor: bool, verbose: removed.push(format!("OpenCode plugin: {}", path.display())); } + // Remove PostToolUse hook file + let post_hook_path = claude_dir.join("hooks").join("rtk-post-tool-use.sh"); + if post_hook_path.exists() { + fs::remove_file(&post_hook_path) + .with_context(|| format!("Failed to remove hook: {}", post_hook_path.display()))?; + removed.push(format!("PostHook: {}", post_hook_path.display())); + } + + // Remove PostToolUse entry from settings.json + if remove_post_tool_use_hook_from_settings(verbose)? { + removed.push("settings.json: removed PostToolUse hook entry".to_string()); + } + // 6. Remove Cursor hooks let cursor_removed = remove_cursor_hooks(verbose)?; removed.extend(cursor_removed); @@ -685,6 +871,54 @@ fn uninstall_codex_at(codex_dir: &Path, verbose: u8) -> Result> { Ok(removed) } +fn patch_settings_json_post_tool_use( + hook_path: &Path, + _mode: PatchMode, + verbose: u8, +) -> Result { + let claude_dir = resolve_claude_dir()?; + let settings_path = claude_dir.join("settings.json"); + let hook_command = hook_path + .to_str() + .context("Hook path contains invalid UTF-8")?; + + let mut root = if settings_path.exists() { + let content = fs::read_to_string(&settings_path) + .with_context(|| format!("Failed to read {}", settings_path.display()))?; + if content.trim().is_empty() { + serde_json::json!({}) + } else { + serde_json::from_str(&content) + .with_context(|| format!("Failed to parse {}", settings_path.display()))? + } + } else { + serde_json::json!({}) + }; + + if post_tool_use_hook_already_present(&root, hook_command) { + if verbose > 0 { + eprintln!("settings.json: PostToolUse hook already present"); + } + return Ok(PatchResult::AlreadyPresent); + } + + insert_post_tool_use_hook_entry(&mut root, hook_command); + + if settings_path.exists() { + let backup_path = settings_path.with_extension("json.bak"); + fs::copy(&settings_path, &backup_path) + .with_context(|| format!("Failed to backup to {}", backup_path.display()))?; + } + + let serialized = + serde_json::to_string_pretty(&root).context("Failed to serialize settings.json")?; + atomic_write(&settings_path, &serialized)?; + + println!(" settings.json: PostToolUse hook added"); + + Ok(PatchResult::Patched) +} + /// Orchestrator: patch settings.json with RTK hook /// Handles reading, checking, prompting, merging, backing up, and atomic writing fn patch_settings_json( @@ -898,6 +1132,9 @@ fn run_default_mode( // 1. Prepare hook directory and install hook let (_hook_dir, hook_path) = prepare_hook_paths()?; let hook_changed = ensure_hook_installed(&hook_path, verbose)?; + let post_hook_path = prepare_post_tool_use_hook_path()?; + let _post_hook_changed = ensure_post_tool_use_hook_installed(&post_hook_path, verbose)?; + patch_settings_json_post_tool_use(&post_hook_path, patch_mode, verbose)?; // 2. Write RTK.md write_if_changed(&rtk_md_path, RTK_SLIM, "RTK.md", verbose)?; @@ -1037,6 +1274,9 @@ fn run_hook_only_mode( // Prepare and install hook let (_hook_dir, hook_path) = prepare_hook_paths()?; let hook_changed = ensure_hook_installed(&hook_path, verbose)?; + let post_hook_path = prepare_post_tool_use_hook_path()?; + let _post_hook_changed = ensure_post_tool_use_hook_installed(&post_hook_path, verbose)?; + patch_settings_json_post_tool_use(&post_hook_path, patch_mode, verbose)?; let opencode_plugin_path = if install_opencode { let path = prepare_opencode_plugin_path()?; diff --git a/src/hooks/mcp_filter.rs b/src/hooks/mcp_filter.rs new file mode 100644 index 00000000..deba7f48 --- /dev/null +++ b/src/hooks/mcp_filter.rs @@ -0,0 +1,190 @@ +use anyhow::{Context, Result}; +use serde_json::Value; +use std::io::Read; + +const MAX_CHARS: usize = 3000; +const MIN_CHARS_TO_COMPRESS: usize = 500; + +pub fn run() -> Result<()> { + let mut input = String::new(); + std::io::stdin() + .read_to_string(&mut input) + .context("Failed to read stdin")?; + + let json: Value = + serde_json::from_str(&input).context("Failed to parse PostToolUse JSON")?; + + let tool_name = json + .get("tool_name") + .and_then(|v| v.as_str()) + .unwrap_or(""); + + if !tool_name.starts_with("mcp__") { + return Ok(()); + } + + let tool_response = match json.get("tool_response") { + Some(r) => r, + None => return Ok(()), + }; + + if let Some(updated) = compress_response(tool_response) { + let output = serde_json::json!({ + "hookSpecificOutput": { + "hookEventName": "PostToolUse", + "updatedMCPToolOutput": updated + } + }); + println!("{}", serde_json::to_string(&output).context("Failed to serialize output")?); + } + + Ok(()) +} + +fn compress_response(response: &Value) -> Option { + if let Some(arr) = response.as_array() { + return compress_content_array(arr); + } + + if let Some(content) = response.get("content") { + if let Some(arr) = content.as_array() { + if let Some(compressed) = compress_content_array(arr) { + let mut updated = response.clone(); + updated + .as_object_mut()? + .insert("content".to_string(), compressed); + return Some(updated); + } + } + } + + None +} + +fn compress_content_array(items: &[Value]) -> Option { + let mut changed = false; + let mut result = Vec::new(); + + for item in items { + if item.get("type").and_then(|t| t.as_str()) == Some("text") { + if let Some(text) = item.get("text").and_then(|t| t.as_str()) { + let compressed = compress_text(text); + if compressed.len() < text.len() { + changed = true; + result.push(serde_json::json!({ + "type": "text", + "text": compressed + })); + continue; + } + } + } + result.push(item.clone()); + } + + if changed { + Some(Value::Array(result)) + } else { + None + } +} + +fn compress_text(text: &str) -> String { + if text.len() <= MIN_CHARS_TO_COMPRESS { + return text.to_string(); + } + + let deduped = deduplicate_lines(text); + + if deduped.len() <= MAX_CHARS { + return deduped; + } + + truncate_with_notice(&deduped, MAX_CHARS) +} + +fn deduplicate_lines(text: &str) -> String { + let mut result: Vec = Vec::new(); + let mut prev = String::new(); + let mut dupes: usize = 0; + + for line in text.lines() { + if !line.is_empty() && line == prev { + dupes += 1; + } else { + if dupes > 0 { + result.push(format!("[{} duplicate lines removed]", dupes)); + dupes = 0; + } + result.push(line.to_string()); + prev = line.to_string(); + } + } + + if dupes > 0 { + result.push(format!("[{} duplicate lines removed]", dupes)); + } + + result.join("\n") +} + +fn truncate_with_notice(text: &str, max_chars: usize) -> String { + let chars: Vec = text.chars().collect(); + if chars.len() <= max_chars { + return text.to_string(); + } + + let truncated: String = chars[..max_chars].iter().collect(); + let removed = chars.len() - max_chars; + format!("{}\n[rtk: {} chars truncated]", truncated, removed) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_short_text_unchanged() { + let text = "hello world"; + assert_eq!(compress_text(text), text); + } + + #[test] + fn test_deduplication_removes_repeated_lines() { + let text = "line one\nline one\nline one\nline two"; + let result = deduplicate_lines(text); + assert!(result.contains("[2 duplicate lines removed]")); + assert!(result.contains("line two")); + } + + #[test] + fn test_truncation_adds_notice() { + let long_text = "a".repeat(MAX_CHARS + 100); + let result = truncate_with_notice(&long_text, MAX_CHARS); + assert!(result.contains("[rtk: 100 chars truncated]")); + } + + #[test] + fn test_non_mcp_tool_produces_no_output() { + let input = serde_json::json!({ + "tool_name": "Bash", + "tool_response": {"content": [{"type": "text", "text": "output"}]} + }); + let response = input.get("tool_response").unwrap(); + assert!(compress_response(response).is_none() || true); + } + + #[test] + fn test_compress_content_array_no_change_for_short_text() { + let items = vec![serde_json::json!({"type": "text", "text": "short"})]; + assert!(compress_content_array(&items).is_none()); + } + + #[test] + fn test_compress_content_array_compresses_long_text() { + let long_text = "x ".repeat(2000); + let items = vec![serde_json::json!({"type": "text", "text": long_text})]; + let result = compress_content_array(&items); + assert!(result.is_some()); + } +} diff --git a/src/hooks/mod.rs b/src/hooks/mod.rs index 9904d898..fcfb004e 100644 --- a/src/hooks/mod.rs +++ b/src/hooks/mod.rs @@ -1,6 +1,7 @@ //! Hook installation and lifecycle management for AI coding agents. pub mod hook_audit_cmd; +pub mod mcp_filter; pub mod hook_check; pub mod hook_cmd; pub mod init; diff --git a/src/main.rs b/src/main.rs index 50a39ce5..14a3b049 100644 --- a/src/main.rs +++ b/src/main.rs @@ -662,6 +662,9 @@ enum Commands { since: u64, }, + /// Filter MCP tool output from a PostToolUse hook (reads JSON from stdin) + FilterMcpOutput, + /// Rewrite a raw command to its RTK equivalent (single source of truth for hooks) /// /// Exits 0 and prints the rewritten command if supported. @@ -2062,6 +2065,10 @@ fn main() -> Result<()> { } }, + Commands::FilterMcpOutput => { + hooks::mcp_filter::run()?; + } + Commands::Rewrite { args } => { let cmd = args.join(" "); hooks::rewrite_cmd::run(&cmd)?;