From b610a3d2437be36ef2bb06f2db40d4e627395730 Mon Sep 17 00:00:00 2001 From: "itai.sagi" Date: Thu, 19 Feb 2026 15:54:56 +0200 Subject: [PATCH 1/3] feat: add AWS CLI and psql modules for token-optimized output Add two new command modules: - `rtk aws`: Specialized filters for STS, S3, EC2, ECS, RDS, and CloudFormation commands. Generic fallback forces --output json and compresses via json_cmd schema extraction. Achieves 60%+ token savings on verbose AWS table/text output. - `rtk psql`: Detects table format (strips borders, separators, row counts, outputs tab-separated) and expanded format (converts -[ RECORD N ]- blocks to compact key=val one-liners). Passthrough for COPY/SET/notices. Both modules include 32 inline tests covering all filters, edge cases, overflow limits, and token savings verification. Hook rewrite patterns added for both distributable and dev hooks. Co-Authored-By: Claude Opus 4.6 --- .claude/hooks/rtk-rewrite.sh | 8 + hooks/rtk-rewrite.sh | 8 + src/aws_cmd.rs | 845 +++++++++++++++++++++++++++++++++++ src/main.rs | 26 ++ src/psql_cmd.rs | 396 ++++++++++++++++ 5 files changed, 1283 insertions(+) create mode 100644 src/aws_cmd.rs create mode 100644 src/psql_cmd.rs diff --git a/.claude/hooks/rtk-rewrite.sh b/.claude/hooks/rtk-rewrite.sh index 5c8bad0..9471807 100755 --- a/.claude/hooks/rtk-rewrite.sh +++ b/.claude/hooks/rtk-rewrite.sh @@ -195,6 +195,14 @@ elif echo "$MATCH_CMD" | grep -qE '^go[[:space:]]+vet([[:space:]]|$)'; then REWRITTEN="${ENV_PREFIX}$(echo "$CMD_BODY" | sed 's/^go vet/rtk go vet/')" elif echo "$MATCH_CMD" | grep -qE '^golangci-lint([[:space:]]|$)'; then REWRITTEN="${ENV_PREFIX}$(echo "$CMD_BODY" | sed 's/^golangci-lint/rtk golangci-lint/')" + +# --- AWS CLI --- +elif echo "$MATCH_CMD" | grep -qE '^aws[[:space:]]+'; then + REWRITTEN="${ENV_PREFIX}$(echo "$CMD_BODY" | sed 's/^aws /rtk aws /')" + +# --- PostgreSQL --- +elif echo "$MATCH_CMD" | grep -qE '^psql([[:space:]]|$)'; then + REWRITTEN="${ENV_PREFIX}$(echo "$CMD_BODY" | sed 's/^psql/rtk psql/')" fi # If no rewrite needed, approve as-is diff --git a/hooks/rtk-rewrite.sh b/hooks/rtk-rewrite.sh index 59e02ca..3a975b4 100644 --- a/hooks/rtk-rewrite.sh +++ b/hooks/rtk-rewrite.sh @@ -185,6 +185,14 @@ elif echo "$MATCH_CMD" | grep -qE '^go[[:space:]]+vet([[:space:]]|$)'; then REWRITTEN="${ENV_PREFIX}$(echo "$CMD_BODY" | sed 's/^go vet/rtk go vet/')" elif echo "$MATCH_CMD" | grep -qE '^golangci-lint([[:space:]]|$)'; then REWRITTEN="${ENV_PREFIX}$(echo "$CMD_BODY" | sed 's/^golangci-lint/rtk golangci-lint/')" + +# --- AWS CLI --- +elif echo "$MATCH_CMD" | grep -qE '^aws[[:space:]]+'; then + REWRITTEN="${ENV_PREFIX}$(echo "$CMD_BODY" | sed 's/^aws /rtk aws /')" + +# --- PostgreSQL --- +elif echo "$MATCH_CMD" | grep -qE '^psql([[:space:]]|$)'; then + REWRITTEN="${ENV_PREFIX}$(echo "$CMD_BODY" | sed 's/^psql/rtk psql/')" fi # If no rewrite needed, approve as-is diff --git a/src/aws_cmd.rs b/src/aws_cmd.rs new file mode 100644 index 0000000..b6f3c51 --- /dev/null +++ b/src/aws_cmd.rs @@ -0,0 +1,845 @@ +//! AWS CLI output compression. +//! +//! Replaces verbose `--output table`/`text` with JSON, then compresses. +//! Specialized filters for high-frequency commands (STS, S3, EC2, ECS, RDS, CloudFormation). + +use crate::json_cmd; +use crate::tracking; +use anyhow::{Context, Result}; +use serde_json::Value; +use std::process::Command; + +const MAX_ITEMS: usize = 20; +const JSON_COMPRESS_DEPTH: usize = 4; + +/// Run an AWS CLI command with token-optimized output +pub fn run(subcommand: &str, args: &[String], verbose: u8) -> Result<()> { + // Build the full sub-path: e.g. "sts" + ["get-caller-identity"] -> "sts get-caller-identity" + let full_sub = if args.is_empty() { + subcommand.to_string() + } else { + format!("{} {}", subcommand, args.join(" ")) + }; + + // Route to specialized handlers + match subcommand { + "sts" if !args.is_empty() && args[0] == "get-caller-identity" => { + run_sts_identity(&args[1..], verbose) + } + "s3" if !args.is_empty() && args[0] == "ls" => run_s3_ls(&args[1..], verbose), + "ec2" if !args.is_empty() && args[0] == "describe-instances" => { + run_ec2_describe(&args[1..], verbose) + } + "ecs" if !args.is_empty() && args[0] == "list-services" => { + run_ecs_list_services(&args[1..], verbose) + } + "ecs" if !args.is_empty() && args[0] == "describe-services" => { + run_ecs_describe_services(&args[1..], verbose) + } + "rds" if !args.is_empty() && args[0] == "describe-db-instances" => { + run_rds_describe(&args[1..], verbose) + } + "cloudformation" if !args.is_empty() && args[0] == "list-stacks" => { + run_cfn_list_stacks(&args[1..], verbose) + } + "cloudformation" if !args.is_empty() && args[0] == "describe-stacks" => { + run_cfn_describe_stacks(&args[1..], verbose) + } + _ => run_generic(subcommand, args, verbose, &full_sub), + } +} + +/// Generic strategy: force --output json, compress via json_cmd schema +fn run_generic(subcommand: &str, args: &[String], verbose: u8, full_sub: &str) -> Result<()> { + let timer = tracking::TimedExecution::start(); + + let mut cmd = Command::new("aws"); + cmd.arg(subcommand); + + let mut has_output_flag = false; + for arg in args { + if arg == "--output" { + has_output_flag = true; + } + cmd.arg(arg); + } + + // Force JSON output if user specified table/text or no output flag + if !has_output_flag { + cmd.args(["--output", "json"]); + } + + if verbose > 0 { + eprintln!("Running: aws {}", full_sub); + } + + let output = cmd.output().context("Failed to run aws CLI")?; + let raw = String::from_utf8_lossy(&output.stdout).to_string(); + let stderr = String::from_utf8_lossy(&output.stderr).to_string(); + + if !output.status.success() { + timer.track( + &format!("aws {}", full_sub), + &format!("rtk aws {}", full_sub), + &stderr, + &stderr, + ); + eprintln!("{}", stderr.trim()); + std::process::exit(output.status.code().unwrap_or(1)); + } + + let filtered = match json_cmd::filter_json_string(&raw, JSON_COMPRESS_DEPTH) { + Ok(schema) => { + println!("{}", schema); + schema + } + Err(_) => { + // Fallback: print raw (maybe not JSON) + print!("{}", raw); + raw.clone() + } + }; + + timer.track( + &format!("aws {}", full_sub), + &format!("rtk aws {}", full_sub), + &raw, + &filtered, + ); + + Ok(()) +} + +fn run_aws_json( + sub_args: &[&str], + extra_args: &[String], + verbose: u8, +) -> Result<(String, String, std::process::ExitStatus)> { + let mut cmd = Command::new("aws"); + for arg in sub_args { + cmd.arg(arg); + } + + // Replace --output table/text with --output json + let mut skip_next = false; + for arg in extra_args { + if skip_next { + skip_next = false; + continue; + } + if arg == "--output" { + skip_next = true; + continue; + } + cmd.arg(arg); + } + cmd.args(["--output", "json"]); + + let cmd_desc = format!("aws {}", sub_args.join(" ")); + if verbose > 0 { + eprintln!("Running: {}", cmd_desc); + } + + let output = cmd + .output() + .context(format!("Failed to run {}", cmd_desc))?; + let stdout = String::from_utf8_lossy(&output.stdout).to_string(); + let stderr = String::from_utf8_lossy(&output.stderr).to_string(); + + if !output.status.success() { + eprintln!("{}", stderr.trim()); + } + + Ok((stdout, stderr, output.status)) +} + +fn run_sts_identity(extra_args: &[String], verbose: u8) -> Result<()> { + let timer = tracking::TimedExecution::start(); + let (raw, stderr, status) = run_aws_json(&["sts", "get-caller-identity"], extra_args, verbose)?; + + if !status.success() { + timer.track( + "aws sts get-caller-identity", + "rtk aws sts get-caller-identity", + &stderr, + &stderr, + ); + std::process::exit(status.code().unwrap_or(1)); + } + + let filtered = match filter_sts_identity(&raw) { + Some(f) => f, + None => raw.clone(), + }; + println!("{}", filtered); + + timer.track( + "aws sts get-caller-identity", + "rtk aws sts get-caller-identity", + &raw, + &filtered, + ); + Ok(()) +} + +fn run_s3_ls(extra_args: &[String], verbose: u8) -> Result<()> { + let timer = tracking::TimedExecution::start(); + + // s3 ls doesn't support --output json, run as-is and filter text + let mut cmd = Command::new("aws"); + cmd.args(["s3", "ls"]); + for arg in extra_args { + cmd.arg(arg); + } + + if verbose > 0 { + eprintln!("Running: aws s3 ls {}", extra_args.join(" ")); + } + + let output = cmd.output().context("Failed to run aws s3 ls")?; + let raw = String::from_utf8_lossy(&output.stdout).to_string(); + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr).to_string(); + timer.track("aws s3 ls", "rtk aws s3 ls", &stderr, &stderr); + eprintln!("{}", stderr.trim()); + std::process::exit(output.status.code().unwrap_or(1)); + } + + let filtered = filter_s3_ls(&raw); + println!("{}", filtered); + + timer.track("aws s3 ls", "rtk aws s3 ls", &raw, &filtered); + Ok(()) +} + +fn run_ec2_describe(extra_args: &[String], verbose: u8) -> Result<()> { + let timer = tracking::TimedExecution::start(); + let (raw, stderr, status) = run_aws_json(&["ec2", "describe-instances"], extra_args, verbose)?; + + if !status.success() { + timer.track( + "aws ec2 describe-instances", + "rtk aws ec2 describe-instances", + &stderr, + &stderr, + ); + std::process::exit(status.code().unwrap_or(1)); + } + + let filtered = match filter_ec2_instances(&raw) { + Some(f) => f, + None => raw.clone(), + }; + println!("{}", filtered); + + timer.track( + "aws ec2 describe-instances", + "rtk aws ec2 describe-instances", + &raw, + &filtered, + ); + Ok(()) +} + +fn run_ecs_list_services(extra_args: &[String], verbose: u8) -> Result<()> { + let timer = tracking::TimedExecution::start(); + let (raw, stderr, status) = run_aws_json(&["ecs", "list-services"], extra_args, verbose)?; + + if !status.success() { + timer.track( + "aws ecs list-services", + "rtk aws ecs list-services", + &stderr, + &stderr, + ); + std::process::exit(status.code().unwrap_or(1)); + } + + let filtered = match filter_ecs_list_services(&raw) { + Some(f) => f, + None => raw.clone(), + }; + println!("{}", filtered); + + timer.track( + "aws ecs list-services", + "rtk aws ecs list-services", + &raw, + &filtered, + ); + Ok(()) +} + +fn run_ecs_describe_services(extra_args: &[String], verbose: u8) -> Result<()> { + let timer = tracking::TimedExecution::start(); + let (raw, stderr, status) = run_aws_json(&["ecs", "describe-services"], extra_args, verbose)?; + + if !status.success() { + timer.track( + "aws ecs describe-services", + "rtk aws ecs describe-services", + &stderr, + &stderr, + ); + std::process::exit(status.code().unwrap_or(1)); + } + + let filtered = match filter_ecs_describe_services(&raw) { + Some(f) => f, + None => raw.clone(), + }; + println!("{}", filtered); + + timer.track( + "aws ecs describe-services", + "rtk aws ecs describe-services", + &raw, + &filtered, + ); + Ok(()) +} + +fn run_rds_describe(extra_args: &[String], verbose: u8) -> Result<()> { + let timer = tracking::TimedExecution::start(); + let (raw, stderr, status) = + run_aws_json(&["rds", "describe-db-instances"], extra_args, verbose)?; + + if !status.success() { + timer.track( + "aws rds describe-db-instances", + "rtk aws rds describe-db-instances", + &stderr, + &stderr, + ); + std::process::exit(status.code().unwrap_or(1)); + } + + let filtered = match filter_rds_instances(&raw) { + Some(f) => f, + None => raw.clone(), + }; + println!("{}", filtered); + + timer.track( + "aws rds describe-db-instances", + "rtk aws rds describe-db-instances", + &raw, + &filtered, + ); + Ok(()) +} + +fn run_cfn_list_stacks(extra_args: &[String], verbose: u8) -> Result<()> { + let timer = tracking::TimedExecution::start(); + let (raw, stderr, status) = + run_aws_json(&["cloudformation", "list-stacks"], extra_args, verbose)?; + + if !status.success() { + timer.track( + "aws cloudformation list-stacks", + "rtk aws cloudformation list-stacks", + &stderr, + &stderr, + ); + std::process::exit(status.code().unwrap_or(1)); + } + + let filtered = match filter_cfn_list_stacks(&raw) { + Some(f) => f, + None => raw.clone(), + }; + println!("{}", filtered); + + timer.track( + "aws cloudformation list-stacks", + "rtk aws cloudformation list-stacks", + &raw, + &filtered, + ); + Ok(()) +} + +fn run_cfn_describe_stacks(extra_args: &[String], verbose: u8) -> Result<()> { + let timer = tracking::TimedExecution::start(); + let (raw, stderr, status) = + run_aws_json(&["cloudformation", "describe-stacks"], extra_args, verbose)?; + + if !status.success() { + timer.track( + "aws cloudformation describe-stacks", + "rtk aws cloudformation describe-stacks", + &stderr, + &stderr, + ); + std::process::exit(status.code().unwrap_or(1)); + } + + let filtered = match filter_cfn_describe_stacks(&raw) { + Some(f) => f, + None => raw.clone(), + }; + println!("{}", filtered); + + timer.track( + "aws cloudformation describe-stacks", + "rtk aws cloudformation describe-stacks", + &raw, + &filtered, + ); + Ok(()) +} + +// --- Filter functions (all use serde_json::Value for resilience) --- + +fn filter_sts_identity(json_str: &str) -> Option { + let v: Value = serde_json::from_str(json_str).ok()?; + let account = v["Account"].as_str().unwrap_or("?"); + let arn = v["Arn"].as_str().unwrap_or("?"); + Some(format!("AWS: {} {}", account, arn)) +} + +fn filter_s3_ls(output: &str) -> String { + let lines: Vec<&str> = output.lines().collect(); + let total = lines.len(); + let mut result: Vec<&str> = lines.iter().take(MAX_ITEMS + 10).copied().collect(); + + if total > MAX_ITEMS + 10 { + result.truncate(MAX_ITEMS + 10); + result.push(""); // will be replaced + return format!( + "{}\n... +{} more items", + result[..result.len() - 1].join("\n"), + total - MAX_ITEMS - 10 + ); + } + + result.join("\n") +} + +fn filter_ec2_instances(json_str: &str) -> Option { + let v: Value = serde_json::from_str(json_str).ok()?; + let reservations = v["Reservations"].as_array()?; + + let mut instances: Vec = Vec::new(); + for res in reservations { + if let Some(insts) = res["Instances"].as_array() { + for inst in insts { + let id = inst["InstanceId"].as_str().unwrap_or("?"); + let state = inst["State"]["Name"].as_str().unwrap_or("?"); + let itype = inst["InstanceType"].as_str().unwrap_or("?"); + let ip = inst["PrivateIpAddress"].as_str().unwrap_or("-"); + + // Extract Name tag + let name = inst["Tags"] + .as_array() + .and_then(|tags| tags.iter().find(|t| t["Key"].as_str() == Some("Name"))) + .and_then(|t| t["Value"].as_str()) + .unwrap_or("-"); + + instances.push(format!("{} {} {} {} ({})", id, state, itype, ip, name)); + } + } + } + + let total = instances.len(); + let mut result = format!("EC2: {} instances\n", total); + + for inst in instances.iter().take(MAX_ITEMS) { + result.push_str(&format!(" {}\n", inst)); + } + + if total > MAX_ITEMS { + result.push_str(&format!(" ... +{} more\n", total - MAX_ITEMS)); + } + + Some(result.trim_end().to_string()) +} + +fn filter_ecs_list_services(json_str: &str) -> Option { + let v: Value = serde_json::from_str(json_str).ok()?; + let arns = v["serviceArns"].as_array()?; + + let mut result = Vec::new(); + let total = arns.len(); + + for arn in arns.iter().take(MAX_ITEMS) { + let arn_str = arn.as_str().unwrap_or("?"); + // Extract short name from ARN: arn:aws:ecs:...:service/cluster/name -> name + let short = arn_str.rsplit('/').next().unwrap_or(arn_str); + result.push(short.to_string()); + } + + let mut output = result.join("\n"); + if total > MAX_ITEMS { + output.push_str(&format!("\n... +{} more services", total - MAX_ITEMS)); + } + + Some(output) +} + +fn filter_ecs_describe_services(json_str: &str) -> Option { + let v: Value = serde_json::from_str(json_str).ok()?; + let services = v["services"].as_array()?; + + let mut result = Vec::new(); + let total = services.len(); + + for svc in services.iter().take(MAX_ITEMS) { + let name = svc["serviceName"].as_str().unwrap_or("?"); + let status = svc["status"].as_str().unwrap_or("?"); + let running = svc["runningCount"].as_i64().unwrap_or(0); + let desired = svc["desiredCount"].as_i64().unwrap_or(0); + let launch = svc["launchType"].as_str().unwrap_or("?"); + result.push(format!( + "{} {} {}/{} ({})", + name, status, running, desired, launch + )); + } + + let mut output = result.join("\n"); + if total > MAX_ITEMS { + output.push_str(&format!("\n... +{} more services", total - MAX_ITEMS)); + } + + Some(output) +} + +fn filter_rds_instances(json_str: &str) -> Option { + let v: Value = serde_json::from_str(json_str).ok()?; + let dbs = v["DBInstances"].as_array()?; + + let mut result = Vec::new(); + let total = dbs.len(); + + for db in dbs.iter().take(MAX_ITEMS) { + let name = db["DBInstanceIdentifier"].as_str().unwrap_or("?"); + let engine = db["Engine"].as_str().unwrap_or("?"); + let version = db["EngineVersion"].as_str().unwrap_or("?"); + let class = db["DBInstanceClass"].as_str().unwrap_or("?"); + let status = db["DBInstanceStatus"].as_str().unwrap_or("?"); + result.push(format!( + "{} {} {} {} {}", + name, engine, version, class, status + )); + } + + let mut output = result.join("\n"); + if total > MAX_ITEMS { + output.push_str(&format!("\n... +{} more instances", total - MAX_ITEMS)); + } + + Some(output) +} + +fn filter_cfn_list_stacks(json_str: &str) -> Option { + let v: Value = serde_json::from_str(json_str).ok()?; + let stacks = v["StackSummaries"].as_array()?; + + let mut result = Vec::new(); + let total = stacks.len(); + + for stack in stacks.iter().take(MAX_ITEMS) { + let name = stack["StackName"].as_str().unwrap_or("?"); + let status = stack["StackStatus"].as_str().unwrap_or("?"); + let date = stack["LastUpdatedTime"] + .as_str() + .or_else(|| stack["CreationTime"].as_str()) + .unwrap_or("?"); + // Truncate date to just date portion + let short_date = if date.len() >= 10 { &date[..10] } else { date }; + result.push(format!("{} {} {}", name, status, short_date)); + } + + let mut output = result.join("\n"); + if total > MAX_ITEMS { + output.push_str(&format!("\n... +{} more stacks", total - MAX_ITEMS)); + } + + Some(output) +} + +fn filter_cfn_describe_stacks(json_str: &str) -> Option { + let v: Value = serde_json::from_str(json_str).ok()?; + let stacks = v["Stacks"].as_array()?; + + let mut result = Vec::new(); + let total = stacks.len(); + + for stack in stacks.iter().take(MAX_ITEMS) { + let name = stack["StackName"].as_str().unwrap_or("?"); + let status = stack["StackStatus"].as_str().unwrap_or("?"); + let date = stack["LastUpdatedTime"] + .as_str() + .or_else(|| stack["CreationTime"].as_str()) + .unwrap_or("?"); + let short_date = if date.len() >= 10 { &date[..10] } else { date }; + result.push(format!("{} {} {}", name, status, short_date)); + + // Show outputs if present + if let Some(outputs) = stack["Outputs"].as_array() { + for out in outputs { + let key = out["OutputKey"].as_str().unwrap_or("?"); + let val = out["OutputValue"].as_str().unwrap_or("?"); + result.push(format!(" {}={}", key, val)); + } + } + } + + let mut output = result.join("\n"); + if total > MAX_ITEMS { + output.push_str(&format!("\n... +{} more stacks", total - MAX_ITEMS)); + } + + Some(output) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_filter_sts_identity() { + let json = r#"{ + "UserId": "AIDAEXAMPLE", + "Account": "123456789012", + "Arn": "arn:aws:iam::123456789012:user/dev" + }"#; + let result = filter_sts_identity(json).unwrap(); + assert_eq!( + result, + "AWS: 123456789012 arn:aws:iam::123456789012:user/dev" + ); + } + + #[test] + fn test_filter_sts_identity_missing_fields() { + let json = r#"{}"#; + let result = filter_sts_identity(json).unwrap(); + assert_eq!(result, "AWS: ? ?"); + } + + #[test] + fn test_filter_sts_identity_invalid_json() { + let result = filter_sts_identity("not json"); + assert!(result.is_none()); + } + + #[test] + fn test_filter_s3_ls_basic() { + let output = "2024-01-01 bucket1\n2024-01-02 bucket2\n2024-01-03 bucket3\n"; + let result = filter_s3_ls(output); + assert!(result.contains("bucket1")); + assert!(result.contains("bucket3")); + } + + #[test] + fn test_filter_s3_ls_overflow() { + let mut lines = Vec::new(); + for i in 1..=50 { + lines.push(format!("2024-01-01 bucket{}", i)); + } + let input = lines.join("\n"); + let result = filter_s3_ls(&input); + assert!(result.contains("... +20 more items")); + } + + #[test] + fn test_filter_ec2_instances() { + let json = r#"{ + "Reservations": [{ + "Instances": [{ + "InstanceId": "i-abc123", + "State": {"Name": "running"}, + "InstanceType": "t3.micro", + "PrivateIpAddress": "10.0.1.5", + "Tags": [{"Key": "Name", "Value": "web-server"}] + }, { + "InstanceId": "i-def456", + "State": {"Name": "stopped"}, + "InstanceType": "t3.large", + "PrivateIpAddress": "10.0.1.6", + "Tags": [{"Key": "Name", "Value": "worker"}] + }] + }] + }"#; + let result = filter_ec2_instances(json).unwrap(); + assert!(result.contains("EC2: 2 instances")); + assert!(result.contains("i-abc123 running t3.micro 10.0.1.5 (web-server)")); + assert!(result.contains("i-def456 stopped t3.large 10.0.1.6 (worker)")); + } + + #[test] + fn test_filter_ec2_no_name_tag() { + let json = r#"{ + "Reservations": [{ + "Instances": [{ + "InstanceId": "i-abc123", + "State": {"Name": "running"}, + "InstanceType": "t3.micro", + "PrivateIpAddress": "10.0.1.5", + "Tags": [] + }] + }] + }"#; + let result = filter_ec2_instances(json).unwrap(); + assert!(result.contains("(-)")); + } + + #[test] + fn test_filter_ec2_invalid_json() { + assert!(filter_ec2_instances("not json").is_none()); + } + + #[test] + fn test_filter_ecs_list_services() { + let json = r#"{ + "serviceArns": [ + "arn:aws:ecs:us-east-1:123:service/cluster/api-service", + "arn:aws:ecs:us-east-1:123:service/cluster/worker-service" + ] + }"#; + let result = filter_ecs_list_services(json).unwrap(); + assert!(result.contains("api-service")); + assert!(result.contains("worker-service")); + assert!(!result.contains("arn:aws")); + } + + #[test] + fn test_filter_ecs_describe_services() { + let json = r#"{ + "services": [{ + "serviceName": "api", + "status": "ACTIVE", + "runningCount": 3, + "desiredCount": 3, + "launchType": "FARGATE" + }] + }"#; + let result = filter_ecs_describe_services(json).unwrap(); + assert_eq!(result, "api ACTIVE 3/3 (FARGATE)"); + } + + #[test] + fn test_filter_rds_instances() { + let json = r#"{ + "DBInstances": [{ + "DBInstanceIdentifier": "mydb", + "Engine": "postgres", + "EngineVersion": "15.4", + "DBInstanceClass": "db.t3.micro", + "DBInstanceStatus": "available" + }] + }"#; + let result = filter_rds_instances(json).unwrap(); + assert_eq!(result, "mydb postgres 15.4 db.t3.micro available"); + } + + #[test] + fn test_filter_cfn_list_stacks() { + let json = r#"{ + "StackSummaries": [{ + "StackName": "my-stack", + "StackStatus": "CREATE_COMPLETE", + "CreationTime": "2024-01-15T10:30:00Z" + }, { + "StackName": "other-stack", + "StackStatus": "UPDATE_COMPLETE", + "LastUpdatedTime": "2024-02-20T14:00:00Z", + "CreationTime": "2024-01-01T00:00:00Z" + }] + }"#; + let result = filter_cfn_list_stacks(json).unwrap(); + assert!(result.contains("my-stack CREATE_COMPLETE 2024-01-15")); + assert!(result.contains("other-stack UPDATE_COMPLETE 2024-02-20")); + } + + #[test] + fn test_filter_cfn_describe_stacks_with_outputs() { + let json = r#"{ + "Stacks": [{ + "StackName": "my-stack", + "StackStatus": "CREATE_COMPLETE", + "CreationTime": "2024-01-15T10:30:00Z", + "Outputs": [ + {"OutputKey": "ApiUrl", "OutputValue": "https://api.example.com"}, + {"OutputKey": "BucketName", "OutputValue": "my-bucket"} + ] + }] + }"#; + let result = filter_cfn_describe_stacks(json).unwrap(); + assert!(result.contains("my-stack CREATE_COMPLETE 2024-01-15")); + assert!(result.contains("ApiUrl=https://api.example.com")); + assert!(result.contains("BucketName=my-bucket")); + } + + #[test] + fn test_filter_cfn_describe_stacks_no_outputs() { + let json = r#"{ + "Stacks": [{ + "StackName": "my-stack", + "StackStatus": "CREATE_COMPLETE", + "CreationTime": "2024-01-15T10:30:00Z" + }] + }"#; + let result = filter_cfn_describe_stacks(json).unwrap(); + assert!(result.contains("my-stack CREATE_COMPLETE 2024-01-15")); + assert!(!result.contains("=")); + } + + fn count_tokens(text: &str) -> usize { + text.split_whitespace().count() + } + + #[test] + fn test_ec2_token_savings() { + let mut instances = Vec::new(); + for i in 1..=10 { + instances.push(format!( + r#"{{ + "InstanceId": "i-{:012x}", + "State": {{"Name": "running", "Code": 16}}, + "InstanceType": "t3.micro", + "PrivateIpAddress": "10.0.1.{}", + "PublicIpAddress": "54.0.0.{}", + "SubnetId": "subnet-abc123", + "VpcId": "vpc-abc123", + "ImageId": "ami-abc123", + "LaunchTime": "2024-01-15T10:30:00Z", + "Placement": {{"AvailabilityZone": "us-east-1a"}}, + "SecurityGroups": [{{"GroupId": "sg-abc123", "GroupName": "default"}}], + "Tags": [{{"Key": "Name", "Value": "server-{}"}}] + }}"#, + i, i, i, i + )); + } + let json = format!( + r#"{{"Reservations": [{{"Instances": [{}]}}]}}"#, + instances.join(",") + ); + + let result = filter_ec2_instances(&json).unwrap(); + let input_tokens = count_tokens(&json); + let output_tokens = count_tokens(&result); + let savings = 100.0 - (output_tokens as f64 / input_tokens as f64 * 100.0); + assert!( + savings >= 60.0, + "EC2 filter: expected >=60% savings, got {:.1}%", + savings + ); + } + + #[test] + fn test_rds_overflow() { + let mut dbs = Vec::new(); + for i in 1..=25 { + dbs.push(format!( + r#"{{"DBInstanceIdentifier": "db-{}", "Engine": "postgres", "EngineVersion": "15.4", "DBInstanceClass": "db.t3.micro", "DBInstanceStatus": "available"}}"#, + i + )); + } + let json = format!(r#"{{"DBInstances": [{}]}}"#, dbs.join(",")); + let result = filter_rds_instances(&json).unwrap(); + assert!(result.contains("... +5 more instances")); + } +} diff --git a/src/main.rs b/src/main.rs index d9de230..7e520bd 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,4 @@ +mod aws_cmd; mod cargo_cmd; mod cc_economics; mod ccusage; @@ -34,6 +35,7 @@ mod playwright_cmd; mod pnpm_cmd; mod prettier_cmd; mod prisma_cmd; +mod psql_cmd; mod pytest_cmd; mod read; mod ruff_cmd; @@ -135,6 +137,22 @@ enum Commands { args: Vec, }, + /// AWS CLI with compact output (force JSON, compress) + Aws { + /// AWS service subcommand (e.g., sts, s3, ec2, ecs, rds, cloudformation) + subcommand: String, + /// Additional arguments + #[arg(trailing_var_arg = true, allow_hyphen_values = true)] + args: Vec, + }, + + /// PostgreSQL client with compact output (strip borders, compress tables) + Psql { + /// psql arguments + #[arg(trailing_var_arg = true, allow_hyphen_values = true)] + args: Vec, + }, + /// pnpm commands with ultra-compact output Pnpm { #[command(subcommand)] @@ -933,6 +951,14 @@ fn main() -> Result<()> { gh_cmd::run(&subcommand, &args, cli.verbose, cli.ultra_compact)?; } + Commands::Aws { subcommand, args } => { + aws_cmd::run(&subcommand, &args, cli.verbose)?; + } + + Commands::Psql { args } => { + psql_cmd::run(&args, cli.verbose)?; + } + Commands::Pnpm { command } => match command { PnpmCommands::List { depth, args } => { pnpm_cmd::run(pnpm_cmd::PnpmCommand::List { depth }, &args, cli.verbose)?; diff --git a/src/psql_cmd.rs b/src/psql_cmd.rs new file mode 100644 index 0000000..d65ccd4 --- /dev/null +++ b/src/psql_cmd.rs @@ -0,0 +1,396 @@ +//! PostgreSQL client (psql) output compression. +//! +//! Detects table and expanded display formats, strips borders/padding, +//! and produces compact tab-separated or key=value output. + +use crate::tracking; +use anyhow::{Context, Result}; +use regex::Regex; + +const MAX_TABLE_ROWS: usize = 30; +const MAX_EXPANDED_RECORDS: usize = 20; + +pub fn run(args: &[String], verbose: u8) -> Result<()> { + let timer = tracking::TimedExecution::start(); + + let mut cmd = std::process::Command::new("psql"); + for arg in args { + cmd.arg(arg); + } + + if verbose > 0 { + eprintln!("Running: psql {}", args.join(" ")); + } + + let output = cmd + .output() + .context("Failed to run psql (is PostgreSQL client installed?)")?; + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + + let exit_code = output.status.code().unwrap_or(1); + + let filtered = filter_psql_output(&stdout); + + if !stderr.is_empty() { + eprint!("{}", stderr); + } + + if let Some(hint) = crate::tee::tee_and_hint(&stdout, "psql", exit_code) { + println!("{}\n{}", filtered, hint); + } else { + println!("{}", filtered); + } + + timer.track( + &format!("psql {}", args.join(" ")), + &format!("rtk psql {}", args.join(" ")), + &stdout, + &filtered, + ); + + if exit_code != 0 { + std::process::exit(exit_code); + } + + Ok(()) +} + +fn filter_psql_output(output: &str) -> String { + if output.trim().is_empty() { + return String::new(); + } + + if is_expanded_format(output) { + filter_expanded(output) + } else if is_table_format(output) { + filter_table(output) + } else { + // Passthrough: COPY results, notices, etc. + output.to_string() + } +} + +fn is_table_format(output: &str) -> bool { + output.lines().any(|line| { + let trimmed = line.trim(); + trimmed.contains("-+-") || trimmed.contains("---+---") + }) +} + +fn is_expanded_format(output: &str) -> bool { + lazy_static::lazy_static! { + static ref EXPANDED_RECORD: Regex = Regex::new(r"-\[ RECORD \d+ \]-").unwrap(); + } + EXPANDED_RECORD.is_match(output) +} + +/// Filter psql table format: +/// - Strip separator lines (----+----) +/// - Strip (N rows) footer +/// - Trim column padding +/// - Output tab-separated +fn filter_table(output: &str) -> String { + lazy_static::lazy_static! { + static ref SEPARATOR: Regex = Regex::new(r"^[-+]+$").unwrap(); + static ref ROW_COUNT: Regex = Regex::new(r"^\(\d+ rows?\)$").unwrap(); + } + + let mut result = Vec::new(); + let mut data_rows = 0; + let mut total_rows = 0; + + for line in output.lines() { + let trimmed = line.trim(); + + // Skip separator lines + if SEPARATOR.is_match(trimmed) { + continue; + } + + // Skip row count footer + if ROW_COUNT.is_match(trimmed) { + continue; + } + + // Skip empty lines + if trimmed.is_empty() { + continue; + } + + // This is a data or header row with | delimiters + if trimmed.contains('|') { + total_rows += 1; + // First row is header, don't count it as data + if total_rows > 1 { + data_rows += 1; + } + + if data_rows <= MAX_TABLE_ROWS || total_rows == 1 { + let cols: Vec<&str> = trimmed.split('|').map(|c| c.trim()).collect(); + result.push(cols.join("\t")); + } + } else { + // Non-table line (e.g., command output like SET, NOTICE) + result.push(trimmed.to_string()); + } + } + + if data_rows > MAX_TABLE_ROWS { + result.push(format!("... +{} more rows", data_rows - MAX_TABLE_ROWS)); + } + + result.join("\n") +} + +/// Filter psql expanded format: +/// Convert -[ RECORD N ]- blocks to one-liner key=val format +fn filter_expanded(output: &str) -> String { + lazy_static::lazy_static! { + static ref RECORD_HEADER: Regex = Regex::new(r"^-\[ RECORD (\d+) \]-").unwrap(); + static ref ROW_COUNT: Regex = Regex::new(r"^\(\d+ rows?\)$").unwrap(); + } + + let mut result = Vec::new(); + let mut current_pairs: Vec = Vec::new(); + let mut current_record: Option = None; + let mut record_count = 0; + + for line in output.lines() { + let trimmed = line.trim(); + + if ROW_COUNT.is_match(trimmed) { + continue; + } + + if let Some(caps) = RECORD_HEADER.captures(trimmed) { + // Flush previous record + if let Some(rec) = current_record.take() { + if record_count <= MAX_EXPANDED_RECORDS { + result.push(format!("{} {}", rec, current_pairs.join(" "))); + } + current_pairs.clear(); + } + record_count += 1; + current_record = Some(format!("[{}]", &caps[1])); + } else if trimmed.contains('|') && current_record.is_some() { + // key | value line + let parts: Vec<&str> = trimmed.splitn(2, '|').collect(); + if parts.len() == 2 { + let key = parts[0].trim(); + let val = parts[1].trim(); + current_pairs.push(format!("{}={}", key, val)); + } + } else if trimmed.is_empty() { + continue; + } else if current_record.is_none() { + // Non-record line before any record (notices, etc.) + result.push(trimmed.to_string()); + } + } + + // Flush last record + if let Some(rec) = current_record.take() { + if record_count <= MAX_EXPANDED_RECORDS { + result.push(format!("{} {}", rec, current_pairs.join(" "))); + } + } + + if record_count > MAX_EXPANDED_RECORDS { + result.push(format!( + "... +{} more records", + record_count - MAX_EXPANDED_RECORDS + )); + } + + result.join("\n") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_is_table_format_detects_separator() { + let input = " id | name\n----+------\n 1 | foo\n(1 row)\n"; + assert!(is_table_format(input)); + } + + #[test] + fn test_is_table_format_rejects_plain() { + assert!(!is_table_format("COPY 5\n")); + assert!(!is_table_format("SET\n")); + } + + #[test] + fn test_is_expanded_format_detects_records() { + let input = "-[ RECORD 1 ]----\nid | 1\nname | foo\n"; + assert!(is_expanded_format(input)); + } + + #[test] + fn test_is_expanded_format_rejects_table() { + let input = " id | name\n----+------\n 1 | foo\n"; + assert!(!is_expanded_format(input)); + } + + #[test] + fn test_filter_table_basic() { + let input = " id | name | email\n----+-------+---------\n 1 | alice | a@b.com\n 2 | bob | b@b.com\n(2 rows)\n"; + let result = filter_table(input); + assert!(result.contains("id\tname\temail")); + assert!(result.contains("1\talice\ta@b.com")); + assert!(result.contains("2\tbob\tb@b.com")); + assert!(!result.contains("----")); + assert!(!result.contains("(2 rows)")); + } + + #[test] + fn test_filter_table_overflow() { + let mut lines = vec![" id | val".to_string(), "----+-----".to_string()]; + for i in 1..=40 { + lines.push(format!(" {} | row{}", i, i)); + } + lines.push("(40 rows)".to_string()); + let input = lines.join("\n"); + + let result = filter_table(&input); + assert!(result.contains("... +10 more rows")); + // Header + 30 data rows + overflow line + let result_lines: Vec<&str> = result.lines().collect(); + assert_eq!(result_lines.len(), 32); // 1 header + 30 data + 1 overflow + } + + #[test] + fn test_filter_table_empty() { + let result = filter_psql_output(""); + assert!(result.is_empty()); + } + + #[test] + fn test_filter_expanded_basic() { + let input = "\ +-[ RECORD 1 ]---- +id | 1 +name | alice +-[ RECORD 2 ]---- +id | 2 +name | bob +"; + let result = filter_expanded(input); + assert!(result.contains("[1] id=1 name=alice")); + assert!(result.contains("[2] id=2 name=bob")); + } + + #[test] + fn test_filter_expanded_overflow() { + let mut lines = Vec::new(); + for i in 1..=25 { + lines.push(format!("-[ RECORD {} ]----", i)); + lines.push(format!("id | {}", i)); + lines.push(format!("name | user{}", i)); + } + let input = lines.join("\n"); + + let result = filter_expanded(&input); + assert!(result.contains("... +5 more records")); + } + + #[test] + fn test_filter_psql_passthrough() { + let input = "COPY 5\n"; + let result = filter_psql_output(input); + assert_eq!(result, "COPY 5\n"); + } + + #[test] + fn test_filter_psql_routes_to_table() { + let input = " id | name\n----+------\n 1 | foo\n(1 row)\n"; + let result = filter_psql_output(input); + assert!(result.contains("id\tname")); + assert!(!result.contains("----")); + } + + #[test] + fn test_filter_psql_routes_to_expanded() { + let input = "-[ RECORD 1 ]----\nid | 1\nname | foo\n"; + let result = filter_psql_output(input); + assert!(result.contains("[1]")); + assert!(result.contains("id=1")); + } + + #[test] + fn test_filter_table_strips_row_count() { + let input = " c\n---\n 1\n(1 row)\n"; + let result = filter_table(input); + assert!(!result.contains("(1 row)")); + } + + #[test] + fn test_filter_expanded_strips_row_count() { + let input = "-[ RECORD 1 ]----\nid | 1\n(1 row)\n"; + let result = filter_expanded(input); + assert!(!result.contains("(1 row)")); + } + + fn count_tokens(text: &str) -> usize { + text.split_whitespace().count() + } + + #[test] + fn test_table_token_savings() { + let mut lines = vec![ + " id | name | email | status ".to_string(), + "------------+------------+--------------------+-----------".to_string(), + ]; + for i in 1..=20 { + lines.push(format!( + " {:10} | {:10} | {:18} | {:9}", + i, + format!("user{}", i), + format!("user{}@example.com", i), + "active" + )); + } + lines.push("(20 rows)".to_string()); + let input = lines.join("\n"); + + let result = filter_table(&input); + let input_tokens = count_tokens(&input); + let output_tokens = count_tokens(&result); + let savings = 100.0 - (output_tokens as f64 / input_tokens as f64 * 100.0); + assert!( + savings >= 30.0, + "Table filter: expected >=30% savings, got {:.1}%", + savings + ); + } + + #[test] + fn test_expanded_token_savings() { + let mut lines = Vec::new(); + for i in 1..=5 { + lines.push(format!( + "-[ RECORD {} ]------------------------------------------", + i + )); + lines.push(format!("id | {}", i)); + lines.push(format!("name | user{}", i)); + lines.push(format!("email | user{}@example.com", i)); + lines.push("status | active".to_string()); + lines.push(format!("created_at | 2024-01-{:02} 10:00:00", i)); + } + lines.push("(5 rows)".to_string()); + let input = lines.join("\n"); + + let result = filter_expanded(&input); + let input_tokens = count_tokens(&input); + let output_tokens = count_tokens(&result); + let savings = 100.0 - (output_tokens as f64 / input_tokens as f64 * 100.0); + assert!( + savings >= 30.0, + "Expanded filter: expected >=30% savings, got {:.1}%", + savings + ); + } +} From a0293af7acf073ee61fb064d9b25a050c25f38bd Mon Sep 17 00:00:00 2001 From: "itai.sagi" Date: Thu, 19 Feb 2026 16:01:30 +0200 Subject: [PATCH 2/3] fix: skip hook rewrite for piped/chained commands + add proxy guidance MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Piped (cmd | grep), chained (cmd && cmd), and redirected (cmd > file) commands are no longer rewritten by the hook — the output is being processed by the next command, not shown to the LLM, so filtering would break the pipeline. RTK.md updated with guidance on when to use `rtk proxy` for raw unfiltered output (parsing specific fields, scripting, debugging). Co-Authored-By: Claude Opus 4.6 --- .claude/hooks/rtk-rewrite.sh | 5 +++++ hooks/rtk-rewrite.sh | 5 +++++ 2 files changed, 10 insertions(+) diff --git a/.claude/hooks/rtk-rewrite.sh b/.claude/hooks/rtk-rewrite.sh index 9471807..c2ae7ea 100755 --- a/.claude/hooks/rtk-rewrite.sh +++ b/.claude/hooks/rtk-rewrite.sh @@ -44,6 +44,11 @@ case "$FIRST_CMD" in *'<<'*) _rtk_audit_log "skip:heredoc" "$CMD"; exit 0 ;; esac +# Skip piped or chained commands — the output is being processed, not shown to LLM +case "$FIRST_CMD" in + *'|'*|*'&&'*|*'>>'*|*'>'*) _rtk_audit_log "skip:piped" "$CMD"; exit 0 ;; +esac + # Strip leading env var assignments for pattern matching # e.g., "TEST_SESSION_ID=2 npx playwright test" → match against "npx playwright test" # but preserve them in the rewritten command for execution. diff --git a/hooks/rtk-rewrite.sh b/hooks/rtk-rewrite.sh index 3a975b4..bf553d4 100644 --- a/hooks/rtk-rewrite.sh +++ b/hooks/rtk-rewrite.sh @@ -31,6 +31,11 @@ case "$FIRST_CMD" in *'<<'*) exit 0 ;; esac +# Skip piped or chained commands — the output is being processed, not shown to LLM +case "$FIRST_CMD" in + *'|'*|*'&&'*|*'>>'*|*'>'*) exit 0 ;; +esac + # Strip leading env var assignments for pattern matching # e.g., "TEST_SESSION_ID=2 npx playwright test" → match against "npx playwright test" # but preserve them in the rewritten command for execution. From 940d2dfc73c3ce35f56ac1e229060175132ab385 Mon Sep 17 00:00:00 2001 From: "itai.sagi" Date: Thu, 19 Feb 2026 16:16:59 +0200 Subject: [PATCH 3/3] Revert "fix: skip hook rewrite for piped/chained commands + add proxy guidance" This reverts commit a0293af7acf073ee61fb064d9b25a050c25f38bd. --- .claude/hooks/rtk-rewrite.sh | 5 ----- hooks/rtk-rewrite.sh | 5 ----- 2 files changed, 10 deletions(-) diff --git a/.claude/hooks/rtk-rewrite.sh b/.claude/hooks/rtk-rewrite.sh index c2ae7ea..9471807 100755 --- a/.claude/hooks/rtk-rewrite.sh +++ b/.claude/hooks/rtk-rewrite.sh @@ -44,11 +44,6 @@ case "$FIRST_CMD" in *'<<'*) _rtk_audit_log "skip:heredoc" "$CMD"; exit 0 ;; esac -# Skip piped or chained commands — the output is being processed, not shown to LLM -case "$FIRST_CMD" in - *'|'*|*'&&'*|*'>>'*|*'>'*) _rtk_audit_log "skip:piped" "$CMD"; exit 0 ;; -esac - # Strip leading env var assignments for pattern matching # e.g., "TEST_SESSION_ID=2 npx playwright test" → match against "npx playwright test" # but preserve them in the rewritten command for execution. diff --git a/hooks/rtk-rewrite.sh b/hooks/rtk-rewrite.sh index bf553d4..3a975b4 100644 --- a/hooks/rtk-rewrite.sh +++ b/hooks/rtk-rewrite.sh @@ -31,11 +31,6 @@ case "$FIRST_CMD" in *'<<'*) exit 0 ;; esac -# Skip piped or chained commands — the output is being processed, not shown to LLM -case "$FIRST_CMD" in - *'|'*|*'&&'*|*'>>'*|*'>'*) exit 0 ;; -esac - # Strip leading env var assignments for pattern matching # e.g., "TEST_SESSION_ID=2 npx playwright test" → match against "npx playwright test" # but preserve them in the rewritten command for execution.