diff --git a/docs/guide/getting-started/configuration.md b/docs/guide/getting-started/configuration.md index da76023911..02c714a8f1 100644 --- a/docs/guide/getting-started/configuration.md +++ b/docs/guide/getting-started/configuration.md @@ -38,11 +38,16 @@ max_width = 120 # maximum output width ignore_dirs = [".git", "node_modules", "target", "__pycache__", ".venv", "vendor"] ignore_files = ["*.lock", "*.min.js", "*.min.css"] -[tee] -enabled = true # save raw output on failure -mode = "failures" # "failures" (default), "always", "never" -max_files = 20 # rotation: keep last N files -# directory = "/custom/tee/path" # optional override +[retriever] +mode = "sqlite" # sqlite (default) | tee (legacy files) | disabled +max_entry_bytes = 10485760 # sqlite: 10 MiB per entry +max_entries = 200 # sqlite: FIFO cap +retention_days = 30 # sqlite: 0 disables age eviction +compression = true # sqlite: gzip blobs (lossless) +# database_path = "/custom/recall.db" +tee_max_files = 20 # tee mode: rotation +tee_max_file_size = 1048576 # tee mode: per-file cap +# tee_directory = "/custom/tee/dir" [telemetry] enabled = true # anonymous daily ping — see Telemetry & Privacy for full details @@ -58,28 +63,30 @@ For full details on what is collected, opt-out options, and GDPR rights, see [Te | Variable | Description | |----------|-------------| | `RTK_DISABLED=1` | Disable RTK for a single command (`RTK_DISABLED=1 git status`) | -| `RTK_TEE_DIR` | Override the tee directory | +| `RTK_RECALL=0` | Disable the recall store for a single command | +| `RTK_RECALL_DB` | Override the recall database path | | `RTK_TELEMETRY_DISABLED=1` | Disable telemetry | | `RTK_HOOK_AUDIT=1` | Enable hook audit logging | | `SKIP_ENV_VALIDATION=1` | Skip env validation (useful with Next.js) | -## Tee system +## Recall system -When a command fails, RTK saves the full raw output to a local file and prints the path: +When a command fails — or a filter trims a long list — RTK persists the full output to an embedded database and prints a recall hint: ``` FAILED: 2/15 tests -[full output: ~/.local/share/rtk/tee/1707753600_cargo_test.log] +[full output: rtk recall 36365b69eda6] ``` -Your AI assistant can then read the file if it needs more detail, without re-running the command. +Your AI assistant runs `rtk recall ` to get back exactly what was elided (or just the hidden tail of a trimmed list), without re-running the command. Other forms: `rtk recall --full | --from N | --lines N | --grep PAT`, `rtk recall --command "cargo test"`, and `rtk recall --list`. Storage is byte-faithful (`BLOB` + lossless gzip). | Setting | Default | Description | |---------|---------|-------------| -| `tee.enabled` | `true` | Enable/disable | -| `tee.mode` | `"failures"` | `"failures"`, `"always"`, `"never"` | -| `tee.max_files` | `20` | Rotation: keep last N files | -| Min size | 500 bytes | Outputs shorter than this are not saved | +| `retriever.mode` | `"sqlite"` | `sqlite` (default), `tee` (legacy files), `disabled` | +| `retriever.max_entry_bytes` | `10485760` | Per-entry storage cap (10 MiB) | +| `retriever.max_entries` | `200` | FIFO cap on retained entries | +| `retriever.retention_days` | `30` | Age eviction in days (0 = off) | +| `retriever.compression` | `true` | gzip stored blobs (lossless) | | Max file size | 1 MB | Truncated above this | ## Excluding commands from auto-rewrite diff --git a/src/cmds/README.md b/src/cmds/README.md index ad3a6ab77a..4e891b1275 100644 --- a/src/cmds/README.md +++ b/src/cmds/README.md @@ -265,7 +265,7 @@ When a filter caps a list at N items (e.g. `take(20)`), the remaining items must **Cap values come from `src/core/truncate.rs`.** Pick the `CAP_*` matching your data class (`CAP_ERRORS`, `CAP_WARNINGS`, `CAP_LIST`, `CAP_INVENTORY`) and bind it to a local `const MAX_XXX: usize = CAP_Y;`. Derive `take(MAX_XXX)`, `> MAX_XXX`, and the offset `MAX_XXX + 1` from the local. These CAPs will later become the configuration surface for per-filter cap tuning (user-overridable via config) — keep all truncation values routed through them so that hook lands as a single switch rather than a codebase-wide hunt. A filter that genuinely needs to deviate uses **`truncate::reduced(CAP_Y, n)`** (e.g. `reduced(CAP_WARNINGS, 5)`) so it still tracks the global when reconfigured — never a bare literal, never `cap - n` (underflows once caps are runtime-configurable), and never `*`/`/` (those scale unboundedly). `reduced` falls back to the full cap if the reduction would empty the list. Each deviation needs a one-line comment stating why; if there's no real reason, just use the plain CAP. See `src/core/README.md` ("Truncation Caps") for the full rationale. -**The tee content must match what `tail` produces.** For `force_tee_tail_hint`, build the tee from the same formatted values shown in the output — not raw/intermediate data. If the filter reformats items before displaying them, pre-build a `Vec` of formatted lines and use it for both the display loop and the tee. +**The stored content must match what was shown.** For `force_tee_tail_hint`, build the recall blob from the same formatted values displayed — not raw/intermediate data — and pass the 1-based first-hidden line as `offset`. `rtk recall ` slices the stored content from that offset, so the agent gets exactly the hidden items. If the filter reformats items before displaying them, pre-build a `Vec` of formatted lines and use it for both the display loop and the recall blob. ### Stderr Handling diff --git a/src/core/README.md b/src/core/README.md index 5d6f8e5ca1..c303573e15 100644 --- a/src/core/README.md +++ b/src/core/README.md @@ -71,12 +71,16 @@ colors = true emoji = true max_width = 120 -[tee] -enabled = true -mode = "failures" # failures | always | never -max_files = 20 -max_file_size = 1048576 -directory = "/custom/tee/dir" +[retriever] +mode = "sqlite" # sqlite (default) | tee (legacy files) | disabled +max_entry_bytes = 10485760 # sqlite: 10 MiB per entry +max_entries = 200 # sqlite: FIFO cap +retention_days = 30 # sqlite: 0 disables age eviction +compression = true # sqlite: gzip blobs (lossless) +# database_path = "/custom/recall.db" +tee_max_files = 20 # tee mode: rotation +tee_max_file_size = 1048576 # tee mode: per-file cap +# tee_directory = "/custom/tee/dir" [telemetry] enabled = true @@ -115,13 +119,13 @@ Core provides infrastructure that `cmds/` and other components consume. These co Consumers must call `timer.track()` on **all** code paths — success, failure, and fallback. Calling `std::process::exit()` before `track()` loses metrics. The raw string passed to `track()` should include both stdout and stderr to produce accurate savings percentages. -### Tee (`tee_and_hint`) +### Output recovery (`tee_and_hint` + recall store) -Consumers that parse structured output (JSON, NDJSON, state machines) should call `tee::tee_and_hint()` to save raw output for LLM recovery on failure. Tee must be called before `std::process::exit()`. +Consumers that parse structured output (JSON, NDJSON, state machines) should call `tee::tee_and_hint()` to persist raw output for LLM recovery on failure. It must be called before `std::process::exit()`. -For truncation recovery on **success** (e.g., list truncated at 20 items), use `tee::force_tee_hint()` which bypasses the tee mode check and writes regardless of exit code. This ensures LLMs always have a `[full output: ...]` recovery path instead of burning tokens working around missing data. +For truncation recovery on **success** (e.g. a list capped at 20 items), use `tee::force_tee_hint()` (multi-line blocks) or `tee::force_tee_tail_hint(content, slug, offset)` (flat lists). All three persist the full output to the content-addressed recall store ([`retriever.rs`](retriever.rs)) and emit a runnable hint — `[full output: rtk recall ]` or `[+N hidden: rtk recall ]` — instead of burning tokens working around missing data. -When the truncated output is a **flat list** and the hidden items start at a predictable line, prefer `tee::force_tee_tail_hint(content, slug, offset)`. It writes the same tee file but emits a directly runnable hint — `[see remaining: tail -n +{offset} ~/path]` — so the agent jumps to exactly the first hidden item without scanning the whole file. The offset is `header_lines + MAX_CAP + 1`. Use `force_tee_hint` instead when the output has multiple sections (e.g. running + stopped containers) and no single offset cleanly covers the gap. +The agent runs `rtk recall ` to get back exactly what was elided. For `force_tee_tail_hint`, `offset` is the 1-based first hidden line (`header_lines + MAX_CAP + 1`); it is stored so the default recall returns only the hidden tail. Storage is byte-faithful (`BLOB` + lossless gzip); tune limits via the `[retriever]` config section. ### Truncation Caps (`truncate`) diff --git a/src/core/config.rs b/src/core/config.rs index ed0f00c6c9..94d8bf5a63 100644 --- a/src/core/config.rs +++ b/src/core/config.rs @@ -14,7 +14,7 @@ pub struct Config { #[serde(default)] pub filters: FilterConfig, #[serde(default)] - pub tee: crate::core::tee::TeeConfig, + pub retriever: crate::core::retriever::RetrieverConfig, #[serde(default)] pub telemetry: TelemetryConfig, #[serde(default)] diff --git a/src/core/constants.rs b/src/core/constants.rs index a5ecd3846d..96f9d06630 100644 --- a/src/core/constants.rs +++ b/src/core/constants.rs @@ -1,5 +1,6 @@ pub const RTK_DATA_DIR: &str = "rtk"; pub const HISTORY_DB: &str = "history.db"; +pub const RECALL_DB: &str = "recall.db"; pub const CONFIG_TOML: &str = "config.toml"; pub const FILTERS_TOML: &str = "filters.toml"; pub const TRUSTED_FILTERS_JSON: &str = "trusted_filters.json"; diff --git a/src/core/mod.rs b/src/core/mod.rs index aaa747e083..8c45ac4ed2 100644 --- a/src/core/mod.rs +++ b/src/core/mod.rs @@ -6,9 +6,11 @@ pub mod constants; pub mod display_helpers; pub mod filter; pub mod guard; +pub mod retriever; pub mod runner; pub mod stream; pub mod tee; +pub mod tee_file; pub mod telemetry; pub mod telemetry_cmd; pub mod toml_filter; diff --git a/src/core/retriever.rs b/src/core/retriever.rs new file mode 100644 index 0000000000..cebc72aca3 --- /dev/null +++ b/src/core/retriever.rs @@ -0,0 +1,636 @@ +//! Content-addressed recall store backing `rtk recall`. + +use super::constants::{RECALL_DB, RTK_DATA_DIR}; +use crate::core::config::Config; +use anyhow::{Context, Result}; +use flate2::read::GzDecoder; +use flate2::write::GzEncoder; +use flate2::Compression; +use rusqlite::{params, Connection, OptionalExtension}; +use serde::{Deserialize, Serialize}; +use std::io::{Read, Write}; +use std::path::PathBuf; +use std::time::{SystemTime, UNIX_EPOCH}; + +const DEFAULT_MAX_ENTRY_BYTES: usize = 10 * 1024 * 1024; +const DEFAULT_MAX_ENTRIES: usize = 200; +const DEFAULT_RETENTION_DAYS: u32 = 30; +pub const MIN_FAILURE_BYTES: usize = 500; +const HASH_HEX_LEN: usize = 12; +const DEFAULT_TEE_MAX_FILES: usize = 20; +const DEFAULT_TEE_MAX_FILE_SIZE: usize = 1_048_576; + +#[derive(Debug, Clone, Copy, PartialEq, Default, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum RecoveryMode { + #[default] + Sqlite, + Tee, + Disabled, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default)] +pub struct RetrieverConfig { + pub mode: RecoveryMode, + pub max_entry_bytes: usize, + pub max_entries: usize, + pub retention_days: u32, + pub compression: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub database_path: Option, + pub tee_max_files: usize, + pub tee_max_file_size: usize, + #[serde(skip_serializing_if = "Option::is_none")] + pub tee_directory: Option, +} + +impl Default for RetrieverConfig { + fn default() -> Self { + Self { + mode: RecoveryMode::Sqlite, + max_entry_bytes: DEFAULT_MAX_ENTRY_BYTES, + max_entries: DEFAULT_MAX_ENTRIES, + retention_days: DEFAULT_RETENTION_DAYS, + compression: true, + database_path: None, + tee_max_files: DEFAULT_TEE_MAX_FILES, + tee_max_file_size: DEFAULT_TEE_MAX_FILE_SIZE, + tee_directory: None, + } + } +} + +pub struct StoredRef { + pub hash: String, + pub hidden_lines: usize, +} + +pub enum Stored { + Saved(StoredRef), + Unavailable, + Empty, +} + +fn now_secs() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs() as i64) + .unwrap_or(0) +} + +fn content_hash(command: &str, content: &[u8]) -> String { + use sha2::{Digest, Sha256}; + let mut hasher = Sha256::new(); + hasher.update(command.as_bytes()); + hasher.update([0u8]); + hasher.update(content); + let hex = format!("{:x}", hasher.finalize()); + hex[..HASH_HEX_LEN].to_string() +} + +fn count_lines(bytes: &[u8]) -> usize { + if bytes.is_empty() { + return 0; + } + let newlines = bytes.iter().filter(|&&b| b == b'\n').count(); + if *bytes.last().unwrap() == b'\n' { + newlines + } else { + newlines + 1 + } +} + +fn slice_from_line(bytes: &[u8], from: usize) -> &[u8] { + if from <= 1 { + return bytes; + } + let mut seen = 0usize; + for (i, &b) in bytes.iter().enumerate() { + if b == b'\n' { + seen += 1; + if seen == from - 1 { + return &bytes[i + 1..]; + } + } + } + &[] +} + +fn slice_first_lines(bytes: &[u8], n: usize) -> &[u8] { + if n == 0 { + return &[]; + } + let mut seen = 0usize; + for (i, &b) in bytes.iter().enumerate() { + if b == b'\n' { + seen += 1; + if seen == n { + return &bytes[..=i]; + } + } + } + bytes +} + +fn grep_bytes(input: &[u8], pattern: &str) -> Vec { + use regex::bytes::Regex; + let re = Regex::new(pattern) + .or_else(|_| Regex::new(®ex::escape(pattern))) + .ok(); + let Some(re) = re else { + return input.to_vec(); + }; + let mut out = Vec::new(); + for line in input.split(|&b| b == b'\n') { + if re.is_match(line) { + out.extend_from_slice(line); + out.push(b'\n'); + } + } + out +} + +fn gzip(data: &[u8]) -> Result> { + let mut enc = GzEncoder::new(Vec::new(), Compression::default()); + enc.write_all(data).context("gzip write")?; + enc.finish().context("gzip finish") +} + +fn gunzip(data: &[u8]) -> Result> { + let mut out = Vec::new(); + GzDecoder::new(data) + .read_to_end(&mut out) + .context("gunzip")?; + Ok(out) +} + +fn db_path(cfg: &RetrieverConfig) -> Result { + if let Ok(p) = std::env::var("RTK_RECALL_DB") { + return Ok(PathBuf::from(p)); + } + if let Some(ref p) = cfg.database_path { + return Ok(p.clone()); + } + let data_dir = dirs::data_local_dir().unwrap_or_else(|| PathBuf::from(".")); + Ok(data_dir.join(RTK_DATA_DIR).join(RECALL_DB)) +} + +fn open(cfg: &RetrieverConfig) -> Result { + let path = db_path(cfg)?; + if let Some(parent) = path.parent() { + let _ = std::fs::create_dir_all(parent); + } + let conn = Connection::open(&path).with_context(|| format!("open {}", path.display()))?; + // best-effort: NFS / read-only filesystems may reject WAL + let _ = conn.execute_batch("PRAGMA journal_mode=WAL; PRAGMA busy_timeout=5000;"); + init_schema(&conn)?; + Ok(conn) +} + +fn init_schema(conn: &Connection) -> Result<()> { + conn.execute_batch( + "CREATE TABLE IF NOT EXISTS recall ( + hash TEXT PRIMARY KEY, + command TEXT NOT NULL, + cwd TEXT, + exit_code INTEGER, + created_at INTEGER NOT NULL, + total_lines INTEGER NOT NULL, + shown_upto INTEGER NOT NULL, + byte_size INTEGER NOT NULL, + truncated INTEGER NOT NULL, + codec TEXT NOT NULL, + blob BLOB NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_recall_command ON recall(command, created_at DESC); + CREATE INDEX IF NOT EXISTS idx_recall_created ON recall(created_at);", + ) + .context("init recall schema") +} + +fn evict(conn: &Connection, cfg: &RetrieverConfig) { + if cfg.retention_days > 0 { + let cutoff = now_secs() - (cfg.retention_days as i64) * 86_400; + let _ = conn.execute("DELETE FROM recall WHERE created_at < ?1", params![cutoff]); + } + if cfg.max_entries > 0 { + let count: i64 = conn + .query_row("SELECT COUNT(*) FROM recall", [], |r| r.get(0)) + .unwrap_or(0); + let excess = count - cfg.max_entries as i64; + if excess > 0 { + let _ = conn.execute( + "DELETE FROM recall WHERE hash IN ( + SELECT hash FROM recall ORDER BY created_at ASC, hash ASC LIMIT ?1 + )", + params![excess], + ); + } + } +} + +pub fn store(content: &[u8], command: &str, exit_code: i32, shown_upto: usize) -> Stored { + if content.is_empty() { + return Stored::Empty; + } + let Ok(config) = Config::load() else { + return Stored::Unavailable; + }; + match store_inner( + &config.retriever, + content, + command, + exit_code, + shown_upto.max(1), + ) { + Ok(r) => Stored::Saved(r), + Err(_) => Stored::Unavailable, + } +} + +fn store_inner( + cfg: &RetrieverConfig, + content: &[u8], + command: &str, + exit_code: i32, + shown_upto: usize, +) -> Result { + let total_lines = count_lines(content); + let (payload, truncated) = if content.len() > cfg.max_entry_bytes { + (&content[..cfg.max_entry_bytes], true) + } else { + (content, false) + }; + let hash = content_hash(command, content); + let (blob, codec): (Vec, &str) = if cfg.compression { + match gzip(payload) { + Ok(z) => (z, "gzip"), + Err(_) => (payload.to_vec(), "raw"), + } + } else { + (payload.to_vec(), "raw") + }; + let cwd = std::env::current_dir() + .ok() + .map(|p| p.to_string_lossy().into_owned()); + + let conn = open(cfg)?; + conn.execute( + "INSERT OR REPLACE INTO recall + (hash, command, cwd, exit_code, created_at, total_lines, shown_upto, byte_size, truncated, codec, blob) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)", + params![ + hash, + command, + cwd, + exit_code, + now_secs(), + total_lines as i64, + shown_upto as i64, + content.len() as i64, + truncated as i64, + codec, + blob + ], + ) + .context("insert recall row")?; + evict(&conn, cfg); + + Ok(StoredRef { + hash, + hidden_lines: total_lines.saturating_sub(shown_upto.saturating_sub(1)), + }) +} + +struct Row { + shown_upto: usize, + truncated: bool, + codec: String, + blob: Vec, +} + +fn map_row(r: &rusqlite::Row) -> rusqlite::Result { + Ok(Row { + shown_upto: r.get::<_, i64>(0)? as usize, + truncated: r.get::<_, i64>(1)? != 0, + codec: r.get(2)?, + blob: r.get(3)?, + }) +} + +const SELECT_COLS: &str = "shown_upto, truncated, codec, blob"; + +fn load_by_hash(conn: &Connection, hash: &str) -> Result> { + let sql = format!( + "SELECT {SELECT_COLS} FROM recall WHERE hash = ?1 OR hash LIKE ?1 || '%' \ + ORDER BY (hash = ?1) DESC, length(hash) ASC LIMIT 1" + ); + Ok(conn.query_row(&sql, params![hash], map_row).optional()?) +} + +fn load_latest_by_command(conn: &Connection, command: &str) -> Result> { + let sql = format!( + "SELECT {SELECT_COLS} FROM recall WHERE command = ?1 OR command LIKE '%' || ?1 || '%' \ + ORDER BY created_at DESC LIMIT 1" + ); + Ok(conn.query_row(&sql, params![command], map_row).optional()?) +} + +fn decode(row: &Row) -> Result> { + match row.codec.as_str() { + "gzip" => gunzip(&row.blob), + _ => Ok(row.blob.clone()), + } +} + +pub struct RecallArgs<'a> { + pub hash: Option<&'a str>, + pub command: Option<&'a str>, + pub full: bool, + pub from: Option, + pub lines: Option, + pub grep: Option<&'a str>, + pub list: bool, +} + +pub fn run_recall(args: RecallArgs) -> Result { + let cfg = Config::load().unwrap_or_default().retriever; + let conn = match open(&cfg) { + Ok(c) => c, + Err(e) => { + eprintln!("rtk recall: store unavailable: {e}"); + return Ok(1); + } + }; + + if args.list { + return list_entries(&conn); + } + + let row = match (args.hash, args.command) { + (Some(h), _) => load_by_hash(&conn, h)?, + (None, Some(c)) => load_latest_by_command(&conn, c)?, + (None, None) => { + eprintln!("rtk recall: provide a , --command , or --list"); + return Ok(2); + } + }; + + let Some(row) = row else { + eprintln!("rtk recall: no matching entry (try `rtk recall --list`)"); + return Ok(1); + }; + + let full = decode(&row)?; + let sliced: Vec = if args.full { + full.clone() + } else if let Some(n) = args.from { + slice_from_line(&full, n).to_vec() + } else if let Some(n) = args.lines { + slice_first_lines(&full, n).to_vec() + } else { + slice_from_line(&full, row.shown_upto).to_vec() + }; + let out = match args.grep { + Some(pat) => grep_bytes(&sliced, pat), + None => sliced, + }; + + let stdout = std::io::stdout(); + let _ = stdout.lock().write_all(&out); + + if row.truncated && args.full { + eprintln!( + "rtk recall: note: output exceeded the {}-byte cap and was stored truncated", + cfg.max_entry_bytes + ); + } + Ok(0) +} + +fn list_entries(conn: &Connection) -> Result { + let mut stmt = conn.prepare( + "SELECT hash, command, total_lines, shown_upto, exit_code, truncated \ + FROM recall ORDER BY created_at DESC LIMIT 50", + )?; + let rows = stmt.query_map([], |r| { + Ok(( + r.get::<_, String>(0)?, + r.get::<_, String>(1)?, + r.get::<_, i64>(2)?, + r.get::<_, i64>(3)?, + r.get::<_, Option>(4)?, + r.get::<_, i64>(5)?, + )) + })?; + + println!( + "{:<14} {:<26} {:>7} {:>7} {:>5} TRUNC", + "HASH", "COMMAND", "LINES", "HIDDEN", "EXIT" + ); + let mut n = 0; + for row in rows { + let (hash, command, total, shown, exit, truncated) = row?; + let hidden = total.saturating_sub(shown.saturating_sub(1)).max(0); + let cmd = if command.chars().count() > 26 { + let head: String = command.chars().take(25).collect(); + format!("{head}…") + } else { + command + }; + println!( + "{:<14} {:<26} {:>7} {:>7} {:>5} {}", + hash, + cmd, + total, + hidden, + exit.map(|e| e.to_string()).unwrap_or_else(|| "-".into()), + if truncated != 0 { "yes" } else { "" } + ); + n += 1; + } + if n == 0 { + println!("(no recall entries)"); + } + Ok(0) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn temp_cfg(dir: &std::path::Path) -> RetrieverConfig { + RetrieverConfig { + database_path: Some(dir.join("recall_test.db")), + ..RetrieverConfig::default() + } + } + + #[test] + fn test_count_lines() { + assert_eq!(count_lines(b""), 0); + assert_eq!(count_lines(b"abc"), 1); + assert_eq!(count_lines(b"a\nb\nc"), 3); + assert_eq!(count_lines(b"a\nb\nc\n"), 3); + assert_eq!(count_lines(b"\n"), 1); + } + + #[test] + fn test_slice_from_line() { + let b = b"l1\nl2\nl3\n"; + assert_eq!(slice_from_line(b, 1), b); + assert_eq!(slice_from_line(b, 2), b"l2\nl3\n"); + assert_eq!(slice_from_line(b, 3), b"l3\n"); + assert_eq!(slice_from_line(b, 4), b""); + assert_eq!(slice_from_line(b, 99), b""); + } + + #[test] + fn test_slice_first_lines() { + let b = b"l1\nl2\nl3\n"; + assert_eq!(slice_first_lines(b, 0), b""); + assert_eq!(slice_first_lines(b, 1), b"l1\n"); + assert_eq!(slice_first_lines(b, 2), b"l1\nl2\n"); + assert_eq!(slice_first_lines(b, 99), b); + } + + #[test] + fn test_content_hash_deterministic() { + let a = content_hash("cmd", b"output"); + assert_eq!(a, content_hash("cmd", b"output")); + assert_eq!(a.len(), HASH_HEX_LEN); + assert_ne!(a, content_hash("cmd2", b"output")); + assert_ne!(a, content_hash("cmd", b"output2")); + } + + #[test] + fn test_grep_bytes() { + let input = b"alpha\nbeta\ngamma\n"; + assert_eq!(grep_bytes(input, "et"), b"beta\n"); + assert_eq!(grep_bytes(input, "^g"), b"gamma\n"); + } + + #[test] + fn test_gzip_roundtrip_arbitrary_bytes() { + let cases: Vec> = vec![ + b"hello\n".to_vec(), + vec![0xff, 0xfe, 0x00, 0x01, 0x80], + b"crlf\r\nline\r\n".to_vec(), + b"lone\rcr".to_vec(), + "emoji😀漢字".as_bytes().to_vec(), + b"no trailing newline".to_vec(), + (0u8..=255).collect(), + ]; + for c in cases { + let z = gzip(&c).expect("gzip"); + assert_eq!(gunzip(&z).expect("gunzip"), c, "gzip must be byte-exact"); + } + } + + #[test] + fn test_store_fetch_byte_faithful() { + let dir = tempfile::tempdir().unwrap(); + let cfg = temp_cfg(dir.path()); + let mut nasty = Vec::new(); + nasty.extend_from_slice(b"line1\r\n"); + nasty.extend_from_slice(&[0xff, 0x00, 0xfe]); + nasty.extend_from_slice("漢字\n".as_bytes()); + nasty.extend_from_slice(b"no-eol-tail"); + + let stored = store_inner(&cfg, &nasty, "nasty-cmd", 0, 1).expect("store"); + let conn = open(&cfg).unwrap(); + let row = load_by_hash(&conn, &stored.hash).unwrap().expect("row"); + assert_eq!( + decode(&row).unwrap(), + nasty, + "stored bytes must round-trip exactly" + ); + } + + #[test] + fn test_store_fetch_raw_codec() { + let dir = tempfile::tempdir().unwrap(); + let cfg = RetrieverConfig { + compression: false, + ..temp_cfg(dir.path()) + }; + let data = vec![0u8, 1, 2, 255, b'\n', b'x']; + let stored = store_inner(&cfg, &data, "c", 0, 1).unwrap(); + let conn = open(&cfg).unwrap(); + let row = load_by_hash(&conn, &stored.hash).unwrap().unwrap(); + assert_eq!(row.codec, "raw"); + assert_eq!(decode(&row).unwrap(), data); + } + + #[test] + fn test_delta_recall_returns_only_missed() { + let dir = tempfile::tempdir().unwrap(); + let cfg = temp_cfg(dir.path()); + let content = b"i1\ni2\ni3\ni4\ni5\n"; + let stored = store_inner(&cfg, content, "list", 0, 3).unwrap(); + assert_eq!(stored.hidden_lines, 3); + let conn = open(&cfg).unwrap(); + let row = load_by_hash(&conn, &stored.hash).unwrap().unwrap(); + let full = decode(&row).unwrap(); + assert_eq!(slice_from_line(&full, row.shown_upto), b"i3\ni4\ni5\n"); + } + + #[test] + fn test_truncation_cap_flagged() { + let dir = tempfile::tempdir().unwrap(); + let cfg = RetrieverConfig { + max_entry_bytes: 10, + ..temp_cfg(dir.path()) + }; + let big = vec![b'a'; 100]; + let stored = store_inner(&cfg, &big, "big", 0, 1).unwrap(); + let conn = open(&cfg).unwrap(); + let row = load_by_hash(&conn, &stored.hash).unwrap().unwrap(); + assert!(row.truncated); + assert_eq!(decode(&row).unwrap().len(), 10); + } + + #[test] + fn test_fifo_count_eviction() { + let dir = tempfile::tempdir().unwrap(); + let cfg = RetrieverConfig { + max_entries: 3, + retention_days: 0, + ..temp_cfg(dir.path()) + }; + for i in 0..5 { + let content = format!("output-{i}"); + store_inner(&cfg, content.as_bytes(), &format!("cmd{i}"), 0, 1).unwrap(); + } + let conn = open(&cfg).unwrap(); + let count: i64 = conn + .query_row("SELECT COUNT(*) FROM recall", [], |r| r.get(0)) + .unwrap(); + assert_eq!(count, 3, "FIFO cap should retain only max_entries"); + } + + #[test] + fn test_dedup_same_content_same_hash() { + let dir = tempfile::tempdir().unwrap(); + let cfg = temp_cfg(dir.path()); + let a = store_inner(&cfg, b"same output\n", "cmd", 0, 1).unwrap(); + let b = store_inner(&cfg, b"same output\n", "cmd", 0, 1).unwrap(); + assert_eq!(a.hash, b.hash); + let conn = open(&cfg).unwrap(); + let count: i64 = conn + .query_row("SELECT COUNT(*) FROM recall", [], |r| r.get(0)) + .unwrap(); + assert_eq!(count, 1, "identical output must dedupe to one row"); + } + + #[test] + fn test_load_by_hash_prefix() { + let dir = tempfile::tempdir().unwrap(); + let cfg = temp_cfg(dir.path()); + let stored = store_inner(&cfg, b"hello world\n", "cmd", 0, 1).unwrap(); + let conn = open(&cfg).unwrap(); + let prefix = &stored.hash[..6]; + assert!(load_by_hash(&conn, prefix).unwrap().is_some()); + } +} diff --git a/src/core/tee.rs b/src/core/tee.rs index 0b01bcae3b..a9da5001b0 100644 --- a/src/core/tee.rs +++ b/src/core/tee.rs @@ -1,269 +1,74 @@ -//! Raw output recovery -- saves unfiltered output to disk on command failure. +//! Recovery-hint dispatch — routes to the sqlite store or legacy tee per `[retriever] mode`. -use super::constants::RTK_DATA_DIR; use crate::core::config::Config; -use std::path::PathBuf; +use crate::core::retriever::{self, RecoveryMode, Stored, MIN_FAILURE_BYTES}; -/// Minimum output size to tee (smaller outputs don't need recovery) -const MIN_TEE_SIZE: usize = 500; +const RECALL_UNAVAILABLE_HINT: &str = + "[recall unavailable — rtk proxy to bypass filtering]"; -/// Default max files to keep in tee directory -const DEFAULT_MAX_FILES: usize = 20; - -/// Default max file size (1MB) -const DEFAULT_MAX_FILE_SIZE: usize = 1_048_576; - -/// Sanitize a command slug for use in filenames. -/// Replaces non-alphanumeric chars (except underscore/hyphen) with underscore, -/// truncates at 40 chars. -fn sanitize_slug(slug: &str) -> String { - let sanitized: String = slug - .chars() - .map(|c| { - if c.is_ascii_alphanumeric() || c == '_' || c == '-' { - c - } else { - '_' - } - }) - .collect(); - if sanitized.len() > 40 { - sanitized[..40].to_string() - } else { - sanitized +fn active_mode() -> RecoveryMode { + if matches!(std::env::var("RTK_RECALL").ok().as_deref(), Some("0")) + || matches!(std::env::var("RTK_TEE").ok().as_deref(), Some("0")) + { + return RecoveryMode::Disabled; } + Config::load().map(|c| c.retriever.mode).unwrap_or_default() } -/// Get the tee directory, respecting config and env overrides. -fn get_tee_dir(config: &Config) -> Option { - // Env var override - if let Ok(dir) = std::env::var("RTK_TEE_DIR") { - return Some(PathBuf::from(dir)); - } - - // Config override - if let Some(ref dir) = config.tee.directory { - return Some(dir.clone()); - } - - // Default: ~/.local/share/rtk/tee/ - dirs::data_local_dir().map(|d| d.join(RTK_DATA_DIR).join("tee")) -} - -/// Rotate old tee files: keep only the last `max_files`, delete oldest. -fn cleanup_old_files(dir: &std::path::Path, max_files: usize) { - let mut entries: Vec<_> = std::fs::read_dir(dir) - .ok() - .into_iter() - .flatten() - .filter_map(|e| e.ok()) - .filter(|e| e.path().extension().is_some_and(|ext| ext == "log")) - .collect(); - - if entries.len() <= max_files { - return; - } - - // Sort by filename (which starts with epoch timestamp = chronological) - entries.sort_by_key(|e| e.file_name()); - - let to_remove = entries.len() - max_files; - for entry in entries.iter().take(to_remove) { - let _ = std::fs::remove_file(entry.path()); - } -} - -/// Check if tee should be skipped based on config, mode, exit code, and size. -/// Returns None if should skip, Some(tee_dir) if should proceed. -fn should_tee( - config: &TeeConfig, - raw_len: usize, - exit_code: i32, - tee_dir: Option, -) -> Option { - if !config.enabled { +pub fn tee_and_hint(raw: &str, command_slug: &str, exit_code: i32) -> Option { + if exit_code == 0 || raw.len() < MIN_FAILURE_BYTES { return None; } - - match config.mode { - TeeMode::Never => return None, - TeeMode::Failures => { - if exit_code == 0 { - return None; + match active_mode() { + RecoveryMode::Disabled => Some(RECALL_UNAVAILABLE_HINT.to_string()), + RecoveryMode::Tee => super::tee_file::tee_and_hint(raw, command_slug), + RecoveryMode::Sqlite => { + match retriever::store(raw.as_bytes(), command_slug, exit_code, 1) { + Stored::Saved(s) => Some(format!("[full output: rtk recall {}]", s.hash)), + Stored::Unavailable => Some(RECALL_UNAVAILABLE_HINT.to_string()), + Stored::Empty => None, } } - TeeMode::Always => {} - } - - if raw_len < MIN_TEE_SIZE { - return None; - } - - tee_dir -} - -/// Write raw output to a tee file in the given directory. -/// Returns file path on success. -fn write_tee_file( - raw: &str, - command_slug: &str, - tee_dir: &std::path::Path, - max_file_size: usize, - max_files: usize, -) -> Option { - std::fs::create_dir_all(tee_dir).ok()?; - - let slug = sanitize_slug(command_slug); - let epoch = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .ok()? - .as_secs(); - let filename = format!("{}_{}.log", epoch, slug); - let filepath = tee_dir.join(filename); - - // Truncate at max_file_size (find a safe UTF-8 char boundary) - let content = if raw.len() > max_file_size { - let boundary = raw - .char_indices() - .take_while(|(i, _)| *i < max_file_size) - .last() - .map(|(i, c)| i + c.len_utf8()) - .unwrap_or(0); - format!( - "{}\n\n--- truncated at {} bytes ---", - &raw[..boundary], - max_file_size - ) - } else { - raw.to_string() - }; - - std::fs::write(&filepath, content).ok()?; - - // Rotate old files - cleanup_old_files(tee_dir, max_files); - - Some(filepath) -} - -/// Write raw output to tee file if conditions are met. -/// Returns file path on success, None if skipped/failed. -pub fn tee_raw(raw: &str, command_slug: &str, exit_code: i32) -> Option { - // Check RTK_TEE=0 env override (disable) - if std::env::var("RTK_TEE").ok().as_deref() == Some("0") { - return None; - } - - let config = Config::load().ok()?; - let tee_dir = get_tee_dir(&config)?; - - let tee_dir = should_tee(&config.tee, raw.len(), exit_code, Some(tee_dir))?; - - write_tee_file( - raw, - command_slug, - &tee_dir, - config.tee.max_file_size, - config.tee.max_files, - ) -} - -fn display_path(path: &std::path::Path) -> String { - if let Some(home) = dirs::home_dir() { - if let Ok(relative) = path.strip_prefix(&home) { - return format!("~/{}", relative.display()); - } } - path.display().to_string() } -fn format_hint(path: &std::path::Path) -> String { - format!("[full output: {}]", display_path(path)) -} - -/// Convenience: tee + format hint in one call. -/// Returns hint string if file was written, None if skipped. -pub fn tee_and_hint(raw: &str, command_slug: &str, exit_code: i32) -> Option { - let path = tee_raw(raw, command_slug, exit_code)?; - Some(format_hint(&path)) -} - -fn force_tee_path(content: &str, command_slug: &str) -> Option { - if std::env::var("RTK_TEE").ok().as_deref() == Some("0") { - return None; - } - +pub fn force_tee_hint(content: &str, command_slug: &str) -> Option { if content.is_empty() { return None; } - - let config = Config::load().ok()?; - - if !config.tee.enabled { - return None; + match active_mode() { + RecoveryMode::Disabled => Some(RECALL_UNAVAILABLE_HINT.to_string()), + RecoveryMode::Tee => super::tee_file::force_tee_hint(content, command_slug), + RecoveryMode::Sqlite => match retriever::store(content.as_bytes(), command_slug, 0, 1) { + Stored::Saved(s) => Some(format!("[full output: rtk recall {}]", s.hash)), + Stored::Unavailable => Some(RECALL_UNAVAILABLE_HINT.to_string()), + Stored::Empty => None, + }, } - - let tee_dir = get_tee_dir(&config)?; - let tee_dir = std::fs::create_dir_all(&tee_dir).ok().and(Some(tee_dir))?; - - write_tee_file( - content, - command_slug, - &tee_dir, - config.tee.max_file_size, - config.tee.max_files, - ) } -/// Returns `[full output: ~/path]`, or None if tee is disabled/skipped. -pub fn force_tee_hint(raw: &str, command_slug: &str) -> Option { - let path = force_tee_path(raw, command_slug)?; - Some(format_hint(&path)) -} - -/// Returns `[see remaining: tail -n +{line_offset} ~/path]`, or None if tee is disabled/skipped. pub fn force_tee_tail_hint( content: &str, command_slug: &str, line_offset: usize, ) -> Option { - let path = force_tee_path(content, command_slug)?; - Some(format!( - "[see remaining: tail -n +{} {}]", - line_offset, - display_path(&path) - )) -} - -/// TeeMode controls when tee writes files. -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Default)] -#[serde(rename_all = "lowercase")] -pub enum TeeMode { - #[default] - Failures, - Always, - Never, -} - -/// Configuration for the tee feature. -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] -pub struct TeeConfig { - pub enabled: bool, - pub mode: TeeMode, - pub max_files: usize, - pub max_file_size: usize, - #[serde(skip_serializing_if = "Option::is_none")] - pub directory: Option, -} - -impl Default for TeeConfig { - fn default() -> Self { - Self { - enabled: true, - mode: TeeMode::default(), - max_files: DEFAULT_MAX_FILES, - max_file_size: DEFAULT_MAX_FILE_SIZE, - directory: None, + if content.is_empty() { + return None; + } + match active_mode() { + RecoveryMode::Disabled => Some(RECALL_UNAVAILABLE_HINT.to_string()), + RecoveryMode::Tee => { + super::tee_file::force_tee_tail_hint(content, command_slug, line_offset) + } + RecoveryMode::Sqlite => { + match retriever::store(content.as_bytes(), command_slug, 0, line_offset) { + Stored::Saved(s) => Some(format!( + "[+{} hidden: rtk recall {}]", + s.hidden_lines, s.hash + )), + Stored::Unavailable => Some(RECALL_UNAVAILABLE_HINT.to_string()), + Stored::Empty => None, + } } } } @@ -271,257 +76,25 @@ impl Default for TeeConfig { #[cfg(test)] mod tests { use super::*; - use std::fs; - - #[test] - fn test_sanitize_slug() { - assert_eq!(sanitize_slug("cargo_test"), "cargo_test"); - assert_eq!(sanitize_slug("cargo test"), "cargo_test"); - assert_eq!(sanitize_slug("cargo-test"), "cargo-test"); - assert_eq!(sanitize_slug("go/test/./pkg"), "go_test___pkg"); - // Truncate at 40 - let long = "a".repeat(50); - assert_eq!(sanitize_slug(&long).len(), 40); - } - - #[test] - fn test_should_tee_disabled() { - let config = TeeConfig { - enabled: false, - ..TeeConfig::default() - }; - let dir = PathBuf::from("/tmp/tee"); - assert!(should_tee(&config, 1000, 1, Some(dir)).is_none()); - } - - #[test] - fn test_should_tee_never_mode() { - let config = TeeConfig { - mode: TeeMode::Never, - ..TeeConfig::default() - }; - let dir = PathBuf::from("/tmp/tee"); - assert!(should_tee(&config, 1000, 1, Some(dir)).is_none()); - } - - #[test] - fn test_should_tee_skip_small_output() { - let config = TeeConfig::default(); - let dir = PathBuf::from("/tmp/tee"); - // Below MIN_TEE_SIZE (500) - assert!(should_tee(&config, 100, 1, Some(dir)).is_none()); - } - - #[test] - fn test_should_tee_skip_success_in_failures_mode() { - let config = TeeConfig::default(); // mode = Failures - let dir = PathBuf::from("/tmp/tee"); - assert!(should_tee(&config, 1000, 0, Some(dir)).is_none()); - } - - #[test] - fn test_should_tee_proceed_on_failure() { - let config = TeeConfig::default(); // mode = Failures - let dir = PathBuf::from("/tmp/tee"); - assert!(should_tee(&config, 1000, 1, Some(dir)).is_some()); - } - - #[test] - fn test_should_tee_always_mode_success() { - let config = TeeConfig { - mode: TeeMode::Always, - ..TeeConfig::default() - }; - let dir = PathBuf::from("/tmp/tee"); - assert!(should_tee(&config, 1000, 0, Some(dir)).is_some()); - } - - #[test] - fn test_write_tee_file_creates_file() { - let tmpdir = tempfile::tempdir().unwrap(); - let content = "error: test failed\n".repeat(50); - let result = write_tee_file( - &content, - "cargo_test", - tmpdir.path(), - DEFAULT_MAX_FILE_SIZE, - 20, - ); - assert!(result.is_some()); - - let path = result.unwrap(); - assert!(path.exists()); - let written = fs::read_to_string(&path).unwrap(); - assert!(written.contains("error: test failed")); - } - - #[test] - fn test_write_tee_file_truncation() { - let tmpdir = tempfile::tempdir().unwrap(); - let big_output = "x".repeat(2000); - // Set max_file_size to 1000 bytes - let result = write_tee_file(&big_output, "test", tmpdir.path(), 1000, 20); - assert!(result.is_some()); - - let path = result.unwrap(); - let content = fs::read_to_string(&path).unwrap(); - assert!(content.contains("--- truncated at 1000 bytes ---")); - assert!(content.len() < 2000); - } - - #[test] - fn test_write_tee_file_truncation_utf8_boundary() { - let tmpdir = tempfile::tempdir().unwrap(); - // Create a string where the truncation point falls inside a multi-byte char. - // Japanese chars are 3 bytes each in UTF-8. - // 332 chars * 3 bytes = 996 bytes, then one more = 999 bytes. - // With max_file_size=998, the cut falls mid-character. - let japanese = "\u{6F22}".repeat(333); // 999 bytes of 3-byte chars - assert_eq!(japanese.len(), 999); - - // Truncate at 998 — falls in the middle of the 333rd character - let result = write_tee_file(&japanese, "test_utf8", tmpdir.path(), 998, 20); - assert!(result.is_some()); - - let path = result.unwrap(); - let content = fs::read_to_string(&path).unwrap(); - assert!(content.contains("--- truncated at 998 bytes ---")); - // Should contain 332 full characters (996 bytes), not panic - assert!(content.starts_with(&"\u{6F22}".repeat(332))); - } - - #[test] - fn test_write_tee_file_truncation_emoji() { - let tmpdir = tempfile::tempdir().unwrap(); - // Emoji are 4 bytes each in UTF-8 - let emojis = "\u{1F600}".repeat(100); // 400 bytes - assert_eq!(emojis.len(), 400); - - // Truncate at 201 — falls mid-emoji (4-byte boundary is at 200, 204) - let result = write_tee_file(&emojis, "test_emoji", tmpdir.path(), 201, 20); - assert!(result.is_some()); - - let path = result.unwrap(); - let content = fs::read_to_string(&path).unwrap(); - assert!(content.contains("--- truncated at 201 bytes ---")); - // The emoji portion should be exactly 200 bytes (50 emojis), - // rounded down from 201 to the nearest char boundary - let target = "\u{1F600}".repeat(50); - assert!(content.starts_with(&target)); - } - - #[test] - fn test_cleanup_old_files() { - let tmpdir = tempfile::tempdir().unwrap(); - let dir = tmpdir.path(); - - // Create 25 .log files - for i in 0..25 { - let filename = format!("{:010}_{}.log", 1000000 + i, "test"); - fs::write(dir.join(&filename), "content").unwrap(); - } - - cleanup_old_files(dir, 20); - - let remaining: Vec<_> = fs::read_dir(dir).unwrap().filter_map(|e| e.ok()).collect(); - assert_eq!(remaining.len(), 20); - - // Oldest 5 should be removed - for i in 0..5 { - let filename = format!("{:010}_{}.log", 1000000 + i, "test"); - assert!(!dir.join(&filename).exists()); - } - // Newest 20 should remain - for i in 5..25 { - let filename = format!("{:010}_{}.log", 1000000 + i, "test"); - assert!(dir.join(&filename).exists()); - } - } - - #[test] - fn test_format_hint() { - let path = PathBuf::from("/tmp/rtk/tee/123_cargo_test.log"); - let hint = format_hint(&path); - assert!(hint.starts_with("[full output: ")); - assert!(hint.ends_with(']')); - assert!(hint.contains("123_cargo_test.log")); - } - - #[test] - fn test_tee_config_default() { - let config = TeeConfig::default(); - assert!(config.enabled); - assert_eq!(config.mode, TeeMode::Failures); - assert_eq!(config.max_files, 20); - assert_eq!(config.max_file_size, 1_048_576); - assert!(config.directory.is_none()); - } - - #[test] - fn test_tee_config_deserialize() { - let toml_str = r#" -enabled = true -mode = "always" -max_files = 10 -max_file_size = 524288 -directory = "/tmp/rtk-tee" -"#; - let config: TeeConfig = toml::from_str(toml_str).unwrap(); - assert!(config.enabled); - assert_eq!(config.mode, TeeMode::Always); - assert_eq!(config.max_files, 10); - assert_eq!(config.max_file_size, 524288); - assert_eq!(config.directory, Some(PathBuf::from("/tmp/rtk-tee"))); - - // Round-trip - let serialized = toml::to_string_pretty(&config).unwrap(); - let deserialized: TeeConfig = toml::from_str(&serialized).unwrap(); - assert_eq!(deserialized.mode, TeeMode::Always); - assert_eq!(deserialized.max_files, 10); - } - - #[test] - fn test_tee_mode_serde() { - // Test all modes via JSON - let mode: TeeMode = serde_json::from_str(r#""always""#).unwrap(); - assert_eq!(mode, TeeMode::Always); - - let mode: TeeMode = serde_json::from_str(r#""failures""#).unwrap(); - assert_eq!(mode, TeeMode::Failures); - - let mode: TeeMode = serde_json::from_str(r#""never""#).unwrap(); - assert_eq!(mode, TeeMode::Never); - } #[test] - fn test_force_tee_hint_skip_empty() { - let hint = force_tee_hint("", "test_cmd"); - assert!(hint.is_none(), "Should skip empty content"); + fn test_tee_and_hint_skips_success() { + let big = "x".repeat(1000); + assert!(tee_and_hint(&big, "cmd", 0).is_none()); } #[test] - fn test_force_tee_hint_respects_env_disable() { - // When RTK_TEE=0, force_tee_hint should return None - std::env::set_var("RTK_TEE", "0"); - let large_output = "x".repeat(1000); - let hint = force_tee_hint(&large_output, "test_cmd"); - std::env::remove_var("RTK_TEE"); - assert!(hint.is_none(), "Should respect RTK_TEE=0"); + fn test_tee_and_hint_skips_tiny_failure() { + assert!(tee_and_hint("tiny", "cmd", 1).is_none()); } #[test] - fn test_force_tee_tail_hint_skip_empty() { - let hint = force_tee_tail_hint("", "test_cmd", 22); - assert!(hint.is_none(), "Should skip empty content"); + fn test_force_tee_hint_skips_empty() { + assert!(force_tee_hint("", "cmd").is_none()); } #[test] - fn test_force_tee_tail_hint_format() { - let path = std::path::PathBuf::from("/tmp/rtk/tee/123_docker_images.log"); - let display = display_path(&path); - let hint = format!("[see remaining: tail -n +{} {}]", 22, display); - assert!(hint.starts_with("[see remaining: tail -n +22 ")); - assert!(hint.ends_with(']')); - assert!(hint.contains("123_docker_images.log")); + fn test_force_tee_tail_hint_skips_empty() { + assert!(force_tee_tail_hint("", "cmd", 5).is_none()); } } diff --git a/src/core/tee_file.rs b/src/core/tee_file.rs new file mode 100644 index 0000000000..f3b2dcf829 --- /dev/null +++ b/src/core/tee_file.rs @@ -0,0 +1,125 @@ +//! Legacy file-based recovery ("tee" mode) — may be deprecated. Prefer the +//! sqlite recall store (`[retriever] mode = "sqlite"`); see retriever.rs. + +use crate::core::config::Config; +use crate::core::constants::RTK_DATA_DIR; +use crate::core::retriever::RetrieverConfig; +use std::path::{Path, PathBuf}; +use std::time::{SystemTime, UNIX_EPOCH}; + +fn sanitize_slug(slug: &str) -> String { + let sanitized: String = slug + .chars() + .map(|c| { + if c.is_ascii_alphanumeric() || c == '_' || c == '-' { + c + } else { + '_' + } + }) + .collect(); + if sanitized.len() > 40 { + sanitized[..40].to_string() + } else { + sanitized + } +} + +fn get_tee_dir(cfg: &RetrieverConfig) -> Option { + if let Ok(dir) = std::env::var("RTK_TEE_DIR") { + return Some(PathBuf::from(dir)); + } + if let Some(ref dir) = cfg.tee_directory { + return Some(dir.clone()); + } + dirs::data_local_dir().map(|d| d.join(RTK_DATA_DIR).join("tee")) +} + +fn cleanup_old_files(dir: &Path, max_files: usize) { + let mut entries: Vec<_> = std::fs::read_dir(dir) + .ok() + .into_iter() + .flatten() + .filter_map(|e| e.ok()) + .filter(|e| e.path().extension().is_some_and(|ext| ext == "log")) + .collect(); + if entries.len() <= max_files { + return; + } + entries.sort_by_key(|e| e.file_name()); + let to_remove = entries.len() - max_files; + for entry in entries.iter().take(to_remove) { + let _ = std::fs::remove_file(entry.path()); + } +} + +fn write_tee_file( + raw: &str, + slug: &str, + dir: &Path, + max_file_size: usize, + max_files: usize, +) -> Option { + std::fs::create_dir_all(dir).ok()?; + let slug = sanitize_slug(slug); + let epoch = SystemTime::now().duration_since(UNIX_EPOCH).ok()?.as_secs(); + let filepath = dir.join(format!("{}_{}.log", epoch, slug)); + let content = if raw.len() > max_file_size { + let boundary = raw + .char_indices() + .take_while(|(i, _)| *i < max_file_size) + .last() + .map(|(i, c)| i + c.len_utf8()) + .unwrap_or(0); + format!( + "{}\n\n--- truncated at {} bytes ---", + &raw[..boundary], + max_file_size + ) + } else { + raw.to_string() + }; + std::fs::write(&filepath, content).ok()?; + cleanup_old_files(dir, max_files); + Some(filepath) +} + +fn display_path(path: &Path) -> String { + if let Some(home) = dirs::home_dir() { + if let Ok(relative) = path.strip_prefix(&home) { + return format!("~/{}", relative.display()); + } + } + path.display().to_string() +} + +fn write(content: &str, slug: &str) -> Option { + let cfg = Config::load().ok()?.retriever; + let dir = get_tee_dir(&cfg)?; + write_tee_file( + content, + slug, + &dir, + cfg.tee_max_file_size, + cfg.tee_max_files, + ) +} + +pub fn tee_and_hint(raw: &str, slug: &str) -> Option { + let path = write(raw, slug)?; + Some(format!("[full output: {}]", display_path(&path))) +} + +pub fn force_tee_hint(content: &str, slug: &str) -> Option { + let path = write(content, slug)?; + Some(format!("[full output: {}]", display_path(&path))) +} + +pub fn force_tee_tail_hint(content: &str, slug: &str, line_offset: usize) -> Option { + let path = write(content, slug)?; + Some(format!( + "[see remaining: tail -n +{} {}]", + line_offset, + display_path(&path) + )) +} diff --git a/src/main.rs b/src/main.rs index d919afe8e3..525c61f3ec 100644 --- a/src/main.rs +++ b/src/main.rs @@ -623,6 +623,26 @@ enum Commands { args: Vec, }, + /// Recall output a filter elided, by content hash + Recall { + hash: Option, + /// Return the complete output, not just the missed part + #[arg(long)] + full: bool, + #[arg(long)] + from: Option, + #[arg(long, conflicts_with = "from")] + lines: Option, + /// Filter recalled lines by regex + #[arg(long)] + grep: Option, + /// Latest entry for a command instead of a hash + #[arg(long, conflicts_with = "hash")] + command: Option, + #[arg(long)] + list: bool, + }, + /// Read stdin, apply filter, print filtered output (Unix pipe mode) Pipe { /// Filter name (cargo-test, pytest, grep, find, git-log, etc.) @@ -1171,6 +1191,7 @@ const RTK_META_COMMANDS: &[&str] = &[ "init", "config", "proxy", + "recall", "run", "hook", "hook-audit", @@ -2325,6 +2346,24 @@ fn run_cli() -> Result { } } + Commands::Recall { + hash, + full, + from, + lines, + grep, + command, + list, + } => core::retriever::run_recall(core::retriever::RecallArgs { + hash: hash.as_deref(), + command: command.as_deref(), + full, + from, + lines, + grep: grep.as_deref(), + list, + })?, + Commands::Proxy { args } => { use std::io::{Read, Write}; use std::process::Stdio;